From 8796840fd8c09639917a94ed55d16a1ed89597df Mon Sep 17 00:00:00 2001 From: Yash Patil Date: Fri, 6 Oct 2023 14:31:15 -0400 Subject: [PATCH 1/2] Migrate multiple queuedWithdrawals, reorganize root calculation --- src/contracts/core/DelegationManager.sol | 56 +++++++++++-------- src/contracts/core/StrategyManager.sol | 15 +++-- src/contracts/interfaces/IStrategyManager.sol | 2 +- src/test/mocks/StrategyManagerMock.sol | 2 +- 4 files changed, 43 insertions(+), 32 deletions(-) diff --git a/src/contracts/core/DelegationManager.sol b/src/contracts/core/DelegationManager.sol index af7eba7a7..126285b65 100644 --- a/src/contracts/core/DelegationManager.sol +++ b/src/contracts/core/DelegationManager.sol @@ -365,33 +365,41 @@ contract DelegationManager is Initializable, OwnableUpgradeable, Pausable, Deleg /// @notice Migrates an existing queued withdrawal from the StrategyManager contract to this contract. /// @dev This function is expected to be removed in the next upgrade, after all queued withdrawals have been migrated. - function migrateQueuedWithdrawal(IStrategyManager.DeprecatedStruct_QueuedWithdrawal memory strategyManagerWithdrawalToMigrate) external { - // check for existence and delete the old storage - bytes32 oldWithdrawalRoot = strategyManager.calculateWithdrawalRoot(strategyManagerWithdrawalToMigrate); - strategyManager.migrateQueuedWithdrawal(oldWithdrawalRoot); - - address staker = strategyManagerWithdrawalToMigrate.staker; - // Create queue entry and increment withdrawal nonce - uint256 nonce = cumulativeWithdrawalsQueued[staker]; - cumulativeWithdrawalsQueued[staker]++; - - Withdrawal memory migratedWithdrawal = Withdrawal({ - staker: staker, - delegatedTo: strategyManagerWithdrawalToMigrate.delegatedAddress, - withdrawer: strategyManagerWithdrawalToMigrate.withdrawerAndNonce.withdrawer, - nonce: nonce, - startBlock: strategyManagerWithdrawalToMigrate.withdrawalStartBlock, - strategies: strategyManagerWithdrawalToMigrate.strategies, - shares: strategyManagerWithdrawalToMigrate.shares - }); + function migrateQueuedWithdrawals(IStrategyManager.DeprecatedStruct_QueuedWithdrawal[] memory strategyManagerWithdrawalsToMigrate) external { + for(uint256 i = 0; i < strategyManagerWithdrawalsToMigrate.length;) { + IStrategyManager.DeprecatedStruct_QueuedWithdrawal memory strategyManagerWithdrawalToMigrate = strategyManagerWithdrawalsToMigrate[i]; + // Delete withdrawal root from strateyManager + (bool isDeleted, bytes32 oldWithdrawalRoot) = strategyManager.migrateQueuedWithdrawal(strategyManagerWithdrawalToMigrate); + // If old storage is deleted from strategyManager + if (isDeleted) { + address staker = strategyManagerWithdrawalToMigrate.staker; + // Create queue entry and increment withdrawal nonce + uint256 nonce = cumulativeWithdrawalsQueued[staker]; + cumulativeWithdrawalsQueued[staker]++; + + Withdrawal memory migratedWithdrawal = Withdrawal({ + staker: staker, + delegatedTo: strategyManagerWithdrawalToMigrate.delegatedAddress, + withdrawer: strategyManagerWithdrawalToMigrate.withdrawerAndNonce.withdrawer, + nonce: nonce, + startBlock: strategyManagerWithdrawalToMigrate.withdrawalStartBlock, + strategies: strategyManagerWithdrawalToMigrate.strategies, + shares: strategyManagerWithdrawalToMigrate.shares + }); - // create the new storage - bytes32 newRoot = calculateWithdrawalRoot(migratedWithdrawal); - pendingWithdrawals[newRoot] = true; + // create the new storage + bytes32 newRoot = calculateWithdrawalRoot(migratedWithdrawal); + pendingWithdrawals[newRoot] = true; - emit WithdrawalQueued(newRoot, migratedWithdrawal); + emit WithdrawalQueued(newRoot, migratedWithdrawal); - emit WithdrawalMigrated(oldWithdrawalRoot, newRoot); + emit WithdrawalMigrated(oldWithdrawalRoot, newRoot); + } + unchecked { + ++i; + } + } + } /** diff --git a/src/contracts/core/StrategyManager.sol b/src/contracts/core/StrategyManager.sol index 3bdd706aa..3653b06b9 100644 --- a/src/contracts/core/StrategyManager.sol +++ b/src/contracts/core/StrategyManager.sol @@ -206,12 +206,15 @@ contract StrategyManager is /// @notice Function called by the DelegationManager as part of the process of transferring existing queued withdrawals from this contract to that contract. /// @dev This function is expected to be removed in the next upgrade, after all queued withdrawals have been migrated. - function migrateQueuedWithdrawal(bytes32 existingWithdrawalRoot) external onlyDelegationManager { - // check for existence - require(withdrawalRootPending[existingWithdrawalRoot], "StrategyManager.migrateQueuedWithdrawal: withdrawal does not exist"); - - // delete the withdrawal - withdrawalRootPending[existingWithdrawalRoot] = false; + function migrateQueuedWithdrawal(DeprecatedStruct_QueuedWithdrawal memory queuedWithdrawal) external onlyDelegationManager returns(bool, bytes32) { + bytes32 existingWithdrawalRoot = calculateWithdrawalRoot(queuedWithdrawal); + bool isDeleted; + // Delete the withdrawal root if it exists + if (withdrawalRootPending[existingWithdrawalRoot]) { + withdrawalRootPending[existingWithdrawalRoot] = false; + isDeleted = true; + } + return (isDeleted, existingWithdrawalRoot); } /** diff --git a/src/contracts/interfaces/IStrategyManager.sol b/src/contracts/interfaces/IStrategyManager.sol index 52debbe44..396d86780 100644 --- a/src/contracts/interfaces/IStrategyManager.sol +++ b/src/contracts/interfaces/IStrategyManager.sol @@ -139,7 +139,7 @@ interface IStrategyManager { address delegatedAddress; } - function migrateQueuedWithdrawal(bytes32 existingWithdrawalRoot) external; + function migrateQueuedWithdrawal(DeprecatedStruct_QueuedWithdrawal memory queuedWithdrawal) external returns (bool, bytes32); function calculateWithdrawalRoot(DeprecatedStruct_QueuedWithdrawal memory queuedWithdrawal) external pure returns (bytes32); } diff --git a/src/test/mocks/StrategyManagerMock.sol b/src/test/mocks/StrategyManagerMock.sol index 7f086ba1c..b2d8f21ed 100644 --- a/src/test/mocks/StrategyManagerMock.sol +++ b/src/test/mocks/StrategyManagerMock.sol @@ -112,7 +112,7 @@ contract StrategyManagerMock is function removeStrategiesFromDepositWhitelist(IStrategy[] calldata /*strategiesToRemoveFromWhitelist*/) external pure {} - function migrateQueuedWithdrawal(bytes32 existingWithdrawalRoot) external {} + function migrateQueuedWithdrawal(DeprecatedStruct_QueuedWithdrawal memory queuedWithdrawal) external returns (bool, bytes32) {} function calculateWithdrawalRoot(DeprecatedStruct_QueuedWithdrawal memory queuedWithdrawal) external pure returns (bytes32) {} } \ No newline at end of file From 0641603b05c88a2588dd46fe3ae04cfd8834293a Mon Sep 17 00:00:00 2001 From: Yash Patil Date: Fri, 6 Oct 2023 14:45:32 -0400 Subject: [PATCH 2/2] fix comment --- src/contracts/core/DelegationManager.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/contracts/core/DelegationManager.sol b/src/contracts/core/DelegationManager.sol index 126285b65..ce0d33b3a 100644 --- a/src/contracts/core/DelegationManager.sol +++ b/src/contracts/core/DelegationManager.sol @@ -363,7 +363,7 @@ contract DelegationManager is Initializable, OwnableUpgradeable, Pausable, Deleg emit WithdrawalCompleted(withdrawalRoot); } - /// @notice Migrates an existing queued withdrawal from the StrategyManager contract to this contract. + /// @notice Migrates an array of queued withdrawals from the StrategyManager contract to this contract. /// @dev This function is expected to be removed in the next upgrade, after all queued withdrawals have been migrated. function migrateQueuedWithdrawals(IStrategyManager.DeprecatedStruct_QueuedWithdrawal[] memory strategyManagerWithdrawalsToMigrate) external { for(uint256 i = 0; i < strategyManagerWithdrawalsToMigrate.length;) {