made forceUndelegate queue a withdrawal for each strategy (#345)

Gajesh2007 · ChaoticWalrus · wadealexc · web-flow · commit 23232072538f · 2023-12-17T18:34:29.000-08:00
* changes * added back comments * chore: fix tests to work with modified behavior (#378) * chore: fix tests to work with modified behavior integration tests in particular are now slightly more flexible * fix: remove memory overwrite * docs: update dmgr docs --------- Co-authored-by: wadealexc <pragma-services@proton.me> --------- Co-authored-by: ChaoticWalrus <93558947+ChaoticWalrus@users.noreply.github.com> Co-authored-by: wadealexc <pragma-services@proton.me>
diff --git a/docs/core/DelegationManager.md b/docs/core/DelegationManager.md
@@ -229,24 +229,24 @@ function undelegate(
 ) 
     external 
     onlyWhenNotPaused(PAUSED_ENTER_WITHDRAWAL_QUEUE)
-    returns (bytes32 withdrawalRoot)
+    returns (bytes32[] memory withdrawalRoots)
 ```
 
-`undelegate` can be called by a Staker to undelegate themselves, or by a Staker's delegated Operator (or that Operator's `delegationApprover`). Undelegation (i) queues a withdrawal on behalf of the Staker for all their delegated shares, and (ii) decreases the Operator's delegated shares according to the amounts and strategies being withdrawn.
+`undelegate` can be called by a Staker to undelegate themselves, or by a Staker's delegated Operator (or that Operator's `delegationApprover`). Undelegation (i) queues withdrawals on behalf of the Staker for all their delegated shares, and (ii) decreases the Operator's delegated shares according to the amounts and strategies being withdrawn.
 
-If the Staker has active shares in either the `EigenPodManager` or `StrategyManager`, they are removed while the withdrawal is in the queue.
+If the Staker has active shares in either the `EigenPodManager` or `StrategyManager`, they are removed while the withdrawal is in the queue - and an individual withdrawal is queued for each strategy removed.
 
-The withdrawal can be completed by the Staker after `withdrawalDelayBlocks`, and does not require the Staker to "fully exit" from the system -- the Staker may choose to receive their shares back in full once the withdrawal is completed (see [`completeQueuedWithdrawal`](#completequeuedwithdrawal) for details).
+The withdrawals can be completed by the Staker after `withdrawalDelayBlocks`. This does not require the Staker to "fully exit" from the system -- the Staker may choose to receive their shares back in full once withdrawals are completed (see [`completeQueuedWithdrawal`](#completequeuedwithdrawal) for details).
 
 Note that becoming an Operator is irreversible! Although Operators can withdraw, they cannot use this method to undelegate from themselves.
 
 *Effects*: 
 * Any shares held by the Staker in the `EigenPodManager` and `StrategyManager` are removed from the Operator's delegated shares.
 * The Staker is undelegated from the Operator
 * If the Staker has no delegatable shares, there is no withdrawal queued or further effects
-* A `Withdrawal` is queued for the Staker, tracking the strategies and shares being withdrawn
-    * The Staker's withdrawal nonce is increased
-    * The hash of the `Withdrawal` is marked as "pending"
+* For each strategy being withdrawn, a `Withdrawal` is queued for the Staker:
+    * The Staker's withdrawal nonce is increased by 1 for each `Withdrawal`
+    * The hash of each `Withdrawal` is marked as "pending"
 * See [`EigenPodManager.removeShares`](./EigenPodManager.md#eigenpodmanagerremoveshares)
 * See [`StrategyManager.removeShares`](./StrategyManager.md#removeshares)
 
diff --git a/src/contracts/core/DelegationManager.sol b/src/contracts/core/DelegationManager.sol
@@ -217,7 +217,7 @@ contract DelegationManager is Initializable, OwnableUpgradeable, Pausable, Deleg
      * a staker from their operator. Undelegation immediately removes ALL active shares/strategies from
      * both the staker and operator, and places the shares and strategies in the withdrawal queue
      */
-    function undelegate(address staker) external onlyWhenNotPaused(PAUSED_ENTER_WITHDRAWAL_QUEUE) returns (bytes32) {
+    function undelegate(address staker) external onlyWhenNotPaused(PAUSED_ENTER_WITHDRAWAL_QUEUE) returns (bytes32[] memory withdrawalRoots) {
         require(isDelegated(staker), "DelegationManager.undelegate: staker must be delegated to undelegate");
         require(!isOperator(staker), "DelegationManager.undelegate: operators cannot be undelegated");
         require(staker != address(0), "DelegationManager.undelegate: cannot undelegate zero address");
@@ -231,8 +231,7 @@ contract DelegationManager is Initializable, OwnableUpgradeable, Pausable, Deleg
 
         // Gather strategies and shares to remove from staker/operator during undelegation
         // Undelegation removes ALL currently-active strategies and shares
-        (IStrategy[] memory strategies, uint256[] memory shares)
-            = getDelegatableShares(staker);
+        (IStrategy[] memory strategies, uint256[] memory shares) = getDelegatableShares(staker);
 
         // emit an event if this action was not initiated by the staker themselves
         if (msg.sender != staker) {
@@ -243,19 +242,28 @@ contract DelegationManager is Initializable, OwnableUpgradeable, Pausable, Deleg
         emit StakerUndelegated(staker, operator);
         delegatedTo[staker] = address(0);
 
-        // if no delegatable shares, return zero root, and don't queue a withdrawal
+        // if no delegatable shares, return an empty array, and don't queue a withdrawal
         if (strategies.length == 0) {
-            return bytes32(0);
+            withdrawalRoots = new bytes32[](0);
         } else {
-            // Remove all strategies/shares from staker and operator and place into queue
-            return _removeSharesAndQueueWithdrawal({
-                staker: staker,
-                operator: operator,
-                withdrawer: staker,
-                strategies: strategies,
-                shares: shares
-            });
+            withdrawalRoots = new bytes32[](strategies.length);
+            for (uint256 i = 0; i < strategies.length; i++) {
+                IStrategy[] memory singleStrategy = new IStrategy[](1);
+                uint256[] memory singleShare = new uint256[](1);
+                singleStrategy[0] = strategies[i];
+                singleShare[0] = shares[i];
+
+                withdrawalRoots[i] = _removeSharesAndQueueWithdrawal({
+                    staker: staker,
+                    operator: operator,
+                    withdrawer: staker,
+                    strategies: singleStrategy,
+                    shares: singleShare
+                });
+            }
         }
+
+        return withdrawalRoots;
     }
 
      /**
diff --git a/src/contracts/interfaces/IDelegationManager.sol b/src/contracts/interfaces/IDelegationManager.sol
@@ -249,7 +249,7 @@ interface IDelegationManager is ISignatureUtils {
      * @dev Reverts if the caller is not the staker, nor the operator who the staker is delegated to, nor the operator's specified "delegationApprover"
      * @dev Reverts if the `staker` is already undelegated.
      */
-    function undelegate(address staker) external returns (bytes32 withdrawalRoot);
+    function undelegate(address staker) external returns (bytes32[] memory withdrawalRoot);
 
     /**
      * Allows a staker to withdraw some shares. Withdrawn shares/strategies are immediately removed
diff --git a/src/test/integration/IntegrationChecks.t.sol b/src/test/integration/IntegrationChecks.t.sol
@@ -52,11 +52,16 @@ contract IntegrationCheckUtils is IntegrationBase {
         // ... check that each withdrawal was successfully enqueued, that the returned roots
         //     match the hashes of each withdrawal, and that the staker and operator have
         //     reduced shares.
-        assert_AllWithdrawalsPending(withdrawalRoots, "staker withdrawals should now be pending");
-        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots, "calculated withdrawals should match returned roots");
-        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals, "staker should have increased nonce by withdrawals.length");
-        assert_Snap_Removed_OperatorShares(operator, strategies, shares, "failed to remove operator shares");
-        assert_Snap_Removed_StakerShares(staker, strategies, shares, "failed to remove staker shares");
+        assert_AllWithdrawalsPending(withdrawalRoots,
+            "check_QueuedWithdrawal_State: staker withdrawals should now be pending");
+        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots,
+            "check_QueuedWithdrawal_State: calculated withdrawals should match returned roots");
+        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals,
+            "check_QueuedWithdrawal_State: staker should have increased nonce by withdrawals.length");
+        assert_Snap_Removed_OperatorShares(operator, strategies, shares,
+            "check_QueuedWithdrawal_State: failed to remove operator shares");
+        assert_Snap_Removed_StakerShares(staker, strategies, shares,
+            "check_QueuedWithdrawal_State: failed to remove staker shares");
     }
 
     function check_Undelegate_State(
@@ -72,13 +77,18 @@ contract IntegrationCheckUtils is IntegrationBase {
         // ... check that the staker is undelegated, all strategies from which the staker is deposited are unqeuued,
         //     that the returned root matches the hashes for each strategy and share amounts, and that the staker
         //     and operator have reduced shares
-        assertEq(withdrawalRoots.length, 1, "should only be one withdrawal root");
-        assertFalse(delegationManager.isDelegated(address(staker)), "staker should not be delegated");
-        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots, "calculated withdrawl should match returned root");
-        assert_AllWithdrawalsPending(withdrawalRoots, "stakers withdrawal should now be pending");
-        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals, "staker should have increased nonce by 1");
-        assert_Snap_Removed_OperatorShares(operator, strategies, shares, "failed to remove operator shares");
-        assert_Snap_Removed_StakerShares(staker, strategies, shares, "failed to remove staker shares");
+        assertFalse(delegationManager.isDelegated(address(staker)),
+            "check_Undelegate_State: staker should not be delegated");
+        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots,
+            "check_Undelegate_State: calculated withdrawl should match returned root");
+        assert_AllWithdrawalsPending(withdrawalRoots,
+            "check_Undelegate_State: stakers withdrawal should now be pending");
+        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals,
+            "check_Undelegate_State: staker should have increased nonce by withdrawals.length");
+        assert_Snap_Removed_OperatorShares(operator, strategies, shares,
+            "check_Undelegate_State: failed to remove operator shares");
+        assert_Snap_Removed_StakerShares(staker, strategies, shares,
+            "check_Undelegate_State: failed to remove staker shares");
     }
 
     function check_Withdrawal_AsTokens_State(
diff --git a/src/test/integration/User.t.sol b/src/test/integration/User.t.sol
@@ -171,20 +171,26 @@ contract User is Test {
     function undelegate() public createSnapshot virtual returns(IDelegationManager.Withdrawal[] memory){
         emit log(_name(".undelegate"));
 
-        IDelegationManager.Withdrawal[] memory withdrawal = new IDelegationManager.Withdrawal[](1);
-        withdrawal[0] = _getExpectedWithdrawalStructForStaker(address(this));
+        IDelegationManager.Withdrawal[] memory expectedWithdrawals = _getExpectedWithdrawalStructsForStaker(address(this));
         delegationManager.undelegate(address(this));
-        return withdrawal;
+
+        for (uint i = 0; i < expectedWithdrawals.length; i++) {
+            emit log("expecting withdrawal:");
+            emit log_named_uint("nonce: ", expectedWithdrawals[i].nonce);
+            emit log_named_address("strat: ", address(expectedWithdrawals[i].strategies[0]));
+            emit log_named_uint("shares: ", expectedWithdrawals[i].shares[0]);
+        }
+        
+        return expectedWithdrawals;
     }
 
     /// @dev Force undelegate staker
     function forceUndelegate(User staker) public createSnapshot virtual returns(IDelegationManager.Withdrawal[] memory){
         emit log_named_string(_name(".forceUndelegate: "), staker.NAME());
 
-        IDelegationManager.Withdrawal[] memory withdrawal = new IDelegationManager.Withdrawal[](1);
-        withdrawal[0] = _getExpectedWithdrawalStructForStaker(address(staker));
+        IDelegationManager.Withdrawal[] memory expectedWithdrawals = _getExpectedWithdrawalStructsForStaker(address(staker));
         delegationManager.undelegate(address(staker));
-        return withdrawal;
+        return expectedWithdrawals;
     }
 
     /// @dev Queues a single withdrawal for every share and strategy pair
@@ -317,20 +323,33 @@ contract User is Test {
         return abi.encodePacked(bytes1(uint8(1)), bytes11(0), address(pod));
     }
 
+    /// @notice Gets the expected withdrawals to be created when the staker is undelegated via a call to `DelegationManager.undelegate()`
     /// @notice Assumes staker and withdrawer are the same and that all strategies and shares are withdrawn
-    function _getExpectedWithdrawalStructForStaker(address staker) internal view returns (IDelegationManager.Withdrawal memory) {
-        (IStrategy[] memory strategies, uint[] memory shares)
+    function _getExpectedWithdrawalStructsForStaker(address staker) internal returns (IDelegationManager.Withdrawal[] memory) {
+        (IStrategy[] memory strategies, uint256[] memory shares)
             = delegationManager.getDelegatableShares(staker);
 
-        return IDelegationManager.Withdrawal({
-            staker: staker,
-            delegatedTo: delegationManager.delegatedTo(staker),
-            withdrawer: staker,
-            nonce: delegationManager.cumulativeWithdrawalsQueued(staker),
-            startBlock: uint32(block.number),
-            strategies: strategies,
-            shares: shares
-        });
+        IDelegationManager.Withdrawal[] memory expectedWithdrawals = new IDelegationManager.Withdrawal[](strategies.length);
+        address delegatedTo = delegationManager.delegatedTo(staker);
+        uint256 nonce = delegationManager.cumulativeWithdrawalsQueued(staker);
+        
+        for (uint256 i = 0; i < strategies.length; ++i) {
+            IStrategy[] memory singleStrategy = new IStrategy[](1);
+            uint256[] memory singleShares = new uint256[](1);
+            singleStrategy[0] = strategies[i];
+            singleShares[0] = shares[i];
+            expectedWithdrawals[i] = IDelegationManager.Withdrawal({
+                staker: staker,
+                delegatedTo: delegatedTo,
+                withdrawer: staker,
+                nonce: (nonce + i),
+                startBlock: uint32(block.number),
+                strategies: singleStrategy,
+                shares: singleShares
+            });
+        }
+
+        return expectedWithdrawals;
     }
 
     function _name(string memory s) internal view returns (string memory) {
diff --git a/src/test/integration/tests/Deposit_Delegate_Queue_Complete.t.sol b/src/test/integration/tests/Deposit_Delegate_Queue_Complete.t.sol
@@ -56,8 +56,8 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
-        for (uint i = 0; i < withdrawals.length; i++) {
-            uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
+        for (uint256 i = 0; i < withdrawals.length; i++) {
+            uint256[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
             IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
             check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], strategies, shares, tokens, expectedTokens);
         }
@@ -115,7 +115,7 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
-        for (uint i = 0; i < withdrawals.length; i++) {
+        for (uint256 i = 0; i < withdrawals.length; i++) {
             staker.completeWithdrawalAsShares(withdrawals[i]);
             check_Withdrawal_AsShares_State(staker, operator, withdrawals[i], strategies, shares);
         }
@@ -182,8 +182,8 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
         // 4. Complete withdrawals
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        for (uint i = 0; i < withdrawals.length; i++) {
-            uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
+        for (uint256 i = 0; i < withdrawals.length; i++) {
+            uint256[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
             IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
             check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], withdrawStrats, withdrawShares, tokens, expectedTokens);
         }
diff --git a/src/test/integration/tests/Deposit_Delegate_Redelegate_Complete.t.sol b/src/test/integration/tests/Deposit_Delegate_Redelegate_Complete.t.sol
@@ -55,8 +55,10 @@ contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationCheckUti
         // 4. Complete withdrawal as shares
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        staker.completeWithdrawalAsShares(withdrawals[0]);
-        check_Withdrawal_AsShares_Undelegated_State(staker, operator1, withdrawals[0], strategies, shares);
+        for (uint256 i = 0; i < withdrawals.length; ++i) {
+            staker.completeWithdrawalAsShares(withdrawals[i]);
+            check_Withdrawal_AsShares_Undelegated_State(staker, operator1, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares);
+        }
 
         // 5. Delegate to a new operator
         staker.delegateTo(operator2);
@@ -76,7 +78,7 @@ contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationCheckUti
         for (uint i = 0; i < withdrawals.length; i++) {
             uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
             IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
-            check_Withdrawal_AsTokens_State(staker, operator2, withdrawals[i], strategies, shares, tokens, expectedTokens);
+            check_Withdrawal_AsTokens_State(staker, operator2, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares, tokens, expectedTokens);
         }
     }
 
diff --git a/src/test/integration/tests/Deposit_Delegate_Undelegate_Complete.t.sol b/src/test/integration/tests/Deposit_Delegate_Undelegate_Complete.t.sol
diff --git a/src/test/mocks/DelegationManagerMock.sol b/src/test/mocks/DelegationManagerMock.sol
diff --git a/src/test/unit/DelegationUnit.t.sol b/src/test/unit/DelegationUnit.t.sol
