handle withdrawals better

Layr-Labs · Sep 21, 2024 · 101c3b3 · 101c3b3
1 parent ad11f4d
commit 101c3b3
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 14 deletions.
diff --git a/src/contracts/core/DelegationManager.sol b/src/contracts/core/DelegationManager.sol
@@ -422,7 +422,6 @@ contract DelegationManager is
     ) external onlyEigenPodManager {
         PrincipalShares existingPrincipalShares = PrincipalShares.wrap(existingPrincipalShares);
 
-        // todo: set beaconChainScalingFactors
         DelegatedShares delegatedSharesBefore = existingPrincipalShares.toDelegatedShares(
             stakerScalingFactors[staker][beaconChainETHStrategy]
         );
@@ -575,8 +574,11 @@ contract DelegationManager is
 
         for (uint256 i = 0; i < withdrawal.strategies.length; i++) {
             IShareManager shareManager = _getShareManager(withdrawal.strategies[i]);
-            // TODO: handle beaconchain slashing
-            uint256 sharesToWithdraw = Shares.unwrap(withdrawal.delegatedShares[i].toShares(totalMagnitudes[i]));
+            uint256 sharesToWithdraw = Shares.unwrap(
+                withdrawal.delegatedShares[i]
+                .scaleForCompleteWithdrawal(stakerScalingFactors[withdrawal.staker][withdrawal.strategies[i]])
+                .toShares(totalMagnitudes[i])
+            );
             if (receiveAsTokens) {
                 // Withdraws `shares` in `strategy` to `withdrawer`. If the shares are virtual beaconChainETH shares,
                 // then a call is ultimately forwarded to the `staker`s EigenPod; otherwise a call is ultimately forwarded
@@ -674,8 +676,7 @@ contract DelegationManager is
             IShareManager shareManager = _getShareManager(strategies[i]);
 
             // delegatedShares for staker to place into queueWithdrawal
-            // TODO: handle beaconchain slashing
-            delegatedSharesToWithdraw[i] = Shares.wrap(sharesToWithdraw[i]).toDelegatedShares(totalMagnitudes[i]);
+            DelegatedShares delegatedSharesToRemove = Shares.wrap(sharesToWithdraw[i]).toDelegatedShares(totalMagnitudes[i]);
             uint256 principalSharesToWithdraw = PrincipalShares.unwrap(
                 delegatedSharesToWithdraw[i].toPrincipalShares(stakerScalingFactors[staker][strategies[i]])
             );
@@ -693,10 +694,12 @@ contract DelegationManager is
                     operator: operator, 
                     staker: staker, 
                     strategy: strategies[i], 
-                    delegatedShares: DelegatedShares.unwrap(delegatedSharesToWithdraw[i])
+                    delegatedShares: DelegatedShares.unwrap(delegatedSharesToRemove)
                 });
             }
 
+            delegatedSharesToWithdraw[i] = delegatedSharesToRemove.scaleForQueueWithdrawal(stakerScalingFactors[staker][strategies[i]]);
+
             // Remove active shares from EigenPodManager/StrategyManager
             // EigenPodManager: this call will revert if it would reduce the Staker's virtual beacon chain ETH shares below zero
             // StrategyManager: this call will revert if `principalSharesToDecrement` exceeds the Staker's current deposit shares in `strategies[i]`
@@ -763,7 +766,6 @@ contract DelegationManager is
         uint256 existingPrincipalShares,
         uint256 addedShares
     ) internal returns (uint256) {
-        // TODO: overflow analysis
         uint256 newDepositScalingFactor;
 
         if (existingPrincipalShares == 0) {
@@ -782,8 +784,8 @@ contract DelegationManager is
             //
             // newShares
             //     = newPrincipalShares.toDelegatedShares(stakerScalingFactors[staker][strategy).toShares(totalMagnitude)
-            //     = newPrincipalShares * newDepositScalingFactor / WAD * totalMagnitude / WAD 
-            //     = (existingPrincipalShares + addedShares) * newDepositScalingFactor / WAD * totalMagnitude / WAD 
+            //     = newPrincipalShares * newDepositScalingFactor / WAD * totalMagnitude / WAD * beaonChainScalingFactor / WAD
+            //     = (existingPrincipalShares + addedShares) * newDepositScalingFactor / WAD * totalMagnitude / WAD * beaconChainScalingFactor / WAD
             //
             // we can solve for
             //

diff --git a/src/contracts/libraries/SlashingLib.sol b/src/contracts/libraries/SlashingLib.sol
@@ -39,6 +39,9 @@ uint32 constant DEALLOCATION_DELAY = 17.5 days;
  *          - These live in the storage of the StrategyManager/EigenPodManager
  *              - `stakerStrategyShares` in the SM is the staker's principalShares that have not been queued for withdrawal in a strategy
  *              - `podOwnerShares` in the EPM is the staker's principalShares that have not been queued for withdrawal in the beaconChainETHStrategy
+ * 
+ * Note that `withdrawal.delegatedShares` is scaled for the beaconChainETHStrategy to divide by the beaconChainScalingFactor upon queueing
+ * and multiply by the beaconChainScalingFactor upon withdrawal
  */
 
 type Shares is uint256;
@@ -77,6 +80,24 @@ library SlashingLib {
         );
     }
 
+    function toShares(DelegatedShares delegatedShares, uint256 magnitude) internal view returns (Shares) {
+        return Shares.wrap(DelegatedShares.unwrap(delegatedShares).mulWad(magnitude));
+    }
+
+    function scaleForQueueWithdrawal(DelegatedShares delegatedShares, StakerScalingFactors storage ssf) internal view returns (DelegatedShares) {
+        return DelegatedShares.wrap(
+            DelegatedShares.unwrap(delegatedShares)
+                .divWad(ssf.getBeaconChainScalingFactor())
+        );
+    }
+
+    function scaleForCompleteWithdrawal(DelegatedShares delegatedShares, StakerScalingFactors storage ssf) internal view returns (DelegatedShares) {
+        return DelegatedShares.wrap(
+            DelegatedShares.unwrap(delegatedShares)
+                .mulWad(ssf.getBeaconChainScalingFactor())
+        );
+    }
+
     /// @dev beaconChainScalingFactor = 0 -> WAD for all non beaconChainETH strategies
     function toDelegatedShares(
         PrincipalShares principalShares,
@@ -88,11 +109,6 @@ library SlashingLib {
                 .mulWad(ssf.getBeaconChainScalingFactor())
         );
     }
-
-
-    function toShares(DelegatedShares delegatedShares, uint256 magnitude) internal view returns (Shares) {
-        return Shares.wrap(DelegatedShares.unwrap(delegatedShares).mulWad(magnitude));
-    }
 
     function toDelegatedShares(Shares shares, uint256 magnitude) internal view returns (DelegatedShares) {
         return DelegatedShares.wrap(Shares.unwrap(shares).divWad(magnitude));