check allowInstantWithdrawals in pp canWithdraw

stakedotlink · Oct 21, 2024 · b72805a · b72805a
1 parent 7a1c252
commit b72805a
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/contracts/core/priorityPool/PriorityPool.sol b/contracts/core/priorityPool/PriorityPool.sol
@@ -220,10 +220,19 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
         uint256 canUnqueue = paused()
             ? 0
             : MathUpgradeable.min(getQueuedTokens(_account, _distributionAmount), totalQueued);
-        uint256 stLINKCanWithdraw = MathUpgradeable.min(
-            stakingPool.balanceOf(_account),
-            stakingPool.canWithdraw() + totalQueued - canUnqueue
-        );
+        uint256 stLINKCanWithdraw = poolStatus == PoolStatus.CLOSED
+            ? 0
+            : MathUpgradeable.min(
+                stakingPool.balanceOf(_account),
+                (
+                    ((allowInstantWithdrawals && withdrawalPool.getTotalQueuedWithdrawals() == 0) ||
+                        _account == address(withdrawalPool))
+                        ? stakingPool.canWithdraw()
+                        : 0
+                ) +
+                    totalQueued -
+                    canUnqueue
+            );
         return canUnqueue + stLINKCanWithdraw;
     }
 
@@ -693,7 +702,11 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
             toWithdraw -= toWithdrawFromQueue;
         }
 
-        if (toWithdraw != 0 && allowInstantWithdrawals) {
+        if (
+            toWithdraw != 0 &&
+            allowInstantWithdrawals &&
+            withdrawalPool.getTotalQueuedWithdrawals() == 0
+        ) {
             uint256 toWithdrawFromPool = MathUpgradeable.min(stakingPool.canWithdraw(), toWithdraw);
             if (toWithdrawFromPool != 0) {
                 stakingPool.withdraw(address(this), address(this), toWithdrawFromPool, _data);

diff --git a/test/core/priorityPool/priority-pool.test.ts b/test/core/priorityPool/priority-pool.test.ts
@@ -749,9 +749,11 @@ describe('PriorityPool', () => {
 
     await strategy.setMinDeposits(0)
     await pp.deposit(toEther(2000), true, ['0x'])
-    assert.equal(fromEther(await pp.canWithdraw(accounts[0], 0)), 2000)
+    assert.equal(fromEther(await pp.canWithdraw(accounts[0], 0)), 1000)
     await strategy.setMaxDeposits(toEther(1100))
     await pp.depositQueuedTokens(toEther(100), toEther(1000), ['0x'])
+    assert.equal(fromEther(await pp.canWithdraw(accounts[0], 0)), 900)
+    await pp.setAllowInstantWithdrawals(true)
     assert.equal(fromEther(await pp.canWithdraw(accounts[0], 0)), 1900)
     await pp.pauseForUpdate()
     assert.equal(fromEther(await pp.canWithdraw(accounts[0], 0)), 1000)