avoid revert in withdrawal edge case

stakedotlink · Nov 25, 2024 · 302148d · 302148d
1 parent ec3a1fc
commit 302148d
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 9 deletions.
diff --git a/contracts/core/interfaces/IWithdrawalPool.sol b/contracts/core/interfaces/IWithdrawalPool.sol
@@ -4,6 +4,8 @@ pragma solidity 0.8.15;
 interface IWithdrawalPool {
     function getTotalQueuedWithdrawals() external view returns (uint256);
 
+    function minWithdrawalAmount() external view returns (uint256);
+
     function deposit(uint256 _amount) external;
 
     function queueWithdrawal(address _account, uint256 _amount) external;

diff --git a/contracts/core/priorityPool/PriorityPool.sol b/contracts/core/priorityPool/PriorityPool.sol
@@ -54,7 +54,7 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
     uint256 public totalQueued;
     // total number of tokens deposited into the staking pool or swapped for LSTs since the last distribution
     uint256 public depositsSinceLastUpdate;
-    // total number of shares received for depositSinceLastUpdate
+    // total number of shares received for depositsSinceLastUpdate
     uint256 private sharesSinceLastUpdate;
 
     // list of all accounts that have ever queued tokens
@@ -102,6 +102,7 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
     error InvalidAmount();
     error StatusAlreadySet();
     error InsufficientLiquidity();
+    error WithdrawFailed();
 
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
@@ -238,8 +239,8 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
         if (msg.sender == address(token)) {
             _deposit(_sender, _value, shouldQueue, data);
         } else if (msg.sender == address(stakingPool)) {
-            uint256 amountQueued = _withdraw(_sender, _value, shouldQueue);
-            token.safeTransfer(_sender, _value - amountQueued);
+            uint256 amountWithdrawn = _withdraw(_sender, _value, shouldQueue, true);
+            token.safeTransfer(_sender, amountWithdrawn);
         } else {
             revert UnauthorizedToken();
         }
@@ -314,7 +315,12 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
                 address(this),
                 toWithdraw
             );
-            toWithdraw = _withdraw(account, toWithdraw, _shouldQueueWithdrawal);
+            toWithdraw -= _withdraw(
+                account,
+                toWithdraw,
+                _shouldQueueWithdrawal,
+                toWithdraw == _amountToWithdraw
+            );
         }
 
         token.safeTransfer(account, _amountToWithdraw - toWithdraw);
@@ -654,16 +660,20 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
      * @param _account account to withdraw for
      * @param _amount amount to withdraw
      * @param _shouldQueueWithdrawal whether a withdrawal should be queued if the the full amount cannot be satisfied
-     * @return the amount of tokens that were queued for withdrawal
+     * @param _shouldRevertOnZero whether call should revert if nothing is withdrawn or queued for withdrawal
+     * @return the amount of tokens that were withdrawn
      **/
     function _withdraw(
         address _account,
         uint256 _amount,
-        bool _shouldQueueWithdrawal
+        bool _shouldQueueWithdrawal,
+        bool _shouldRevertOnZero
     ) internal returns (uint256) {
         if (poolStatus == PoolStatus.CLOSED) revert WithdrawalsDisabled();
 
         uint256 toWithdraw = _amount;
+        uint256 withdrawn;
+        uint256 queued;
 
         if (totalQueued != 0) {
             uint256 toWithdrawFromQueue = toWithdraw <= totalQueued ? toWithdraw : totalQueued;
@@ -672,15 +682,24 @@ contract PriorityPool is UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeabl
             depositsSinceLastUpdate += toWithdrawFromQueue;
             sharesSinceLastUpdate += stakingPool.getSharesByStake(toWithdrawFromQueue);
             toWithdraw -= toWithdrawFromQueue;
+            withdrawn = toWithdrawFromQueue;
         }
 
         if (toWithdraw != 0) {
             if (!_shouldQueueWithdrawal) revert InsufficientLiquidity();
-            withdrawalPool.queueWithdrawal(_account, toWithdraw);
+
+            if (toWithdraw >= withdrawalPool.minWithdrawalAmount()) {
+                withdrawalPool.queueWithdrawal(_account, toWithdraw);
+                queued = toWithdraw;
+            } else {
+                IERC20Upgradeable(address(stakingPool)).safeTransfer(_account, toWithdraw);
+            }
         }
 
-        emit Withdraw(_account, _amount - toWithdraw);
-        return toWithdraw;
+        if (_shouldRevertOnZero && withdrawn + queued == 0) revert WithdrawFailed();
+        if (_amount != toWithdraw) emit Withdraw(_account, _amount - toWithdraw);
+
+        return withdrawn;
     }
 
     /**