add strats on rejoin

Layr-Labs · Sep 18, 2024 · 39b9383 · 39b9383
1 parent 1076dc5
commit 39b9383
Showing 1 changed file with 7 additions and 15 deletions.
diff --git a/src/contracts/core/StakeRootCompendium.sol b/src/contracts/core/StakeRootCompendium.sol
@@ -52,16 +52,7 @@ contract StakeRootCompendium is StakeRootCompendiumStorage {
                 cumulativeChargePerStrategyLastPaid: uint96(cumulativeChargePerStrategy)
             });
 
-            // empty their strategies and multipliers if they were force removed before
-            address[] memory keys = new address[](operatorSetToStrategyAndMultipliers[operatorSet.avs][operatorSet.operatorSetId].length());
-            for (uint256 i = 0; i < keys.length; i++) {
-                (address key,) = operatorSetToStrategyAndMultipliers[operatorSet.avs][operatorSet.operatorSetId].at(i);
-                keys[i] = key;
-            }
-            for (uint256 i = 0; i < keys.length; i++) {
-                operatorSetToStrategyAndMultipliers[operatorSet.avs][operatorSet.operatorSetId].remove(keys[i]);
-            }
-
+            _updateTotalStrategies(0, operatorSetToStrategyAndMultipliers[operatorSet.avs][operatorSet.operatorSetId].length());
             operatorSetToIndex[operatorSet.avs][operatorSet.operatorSetId].push(uint32(block.timestamp), uint224(operatorSets.length)); 
             operatorSets.push(operatorSet);
         }
@@ -109,18 +100,19 @@ contract StakeRootCompendium is StakeRootCompendiumStorage {
 
     /// @inheritdoc IStakeRootCompendium
     function removeStrategiesAndMultipliers(uint32 operatorSetId, IStrategy[] calldata strategies) external {
-        // update the deposit balance for the operator set whenever number of strategies is changed
         OperatorSet memory operatorSet = OperatorSet({avs: msg.sender, operatorSetId: operatorSetId});
+        // update the deposit balance for the operator set whenever number of strategies is changed
         _updateDepositInfo(operatorSet);
 
         // note below either all strategies are removed or none are removed and transaction reverts
-        uint256 numStrategiesBefore = operatorSetToStrategyAndMultipliers[msg.sender][operatorSetId].length();
+        uint256 numStrategiesBefore = operatorSetToStrategyAndMultipliers[operatorSet.avs][operatorSetId].length();
         for (uint256 i = 0; i < strategies.length; i++) {
             require(
-                operatorSetToStrategyAndMultipliers[msg.sender][operatorSetId].remove(address(strategies[i])),
+                operatorSetToStrategyAndMultipliers[operatorSet.avs][operatorSetId].remove(address(strategies[i])),
                 "StakeRootCompendium.removeStrategiesAndMultipliers: strategy not found"
             );
         }
+
         _updateTotalStrategies(numStrategiesBefore, numStrategiesBefore - strategies.length);
     }
 
@@ -182,8 +174,8 @@ contract StakeRootCompendium is StakeRootCompendiumStorage {
                 }
             }
         }
-        // todo use gas metering to make this call incentive neutral and simply refund any excess balance to AVS'?
-        // send the penalty to the caller
+        // note this reward is subject to a race condition where anyone can claim the penalty if they submit their transaction first
+        // it is the caller's responsibility to use private mempool to protect themselves from reverted transactions
         (bool success, ) = payable(msg.sender).call{value: penalty}("");
         require(success, "StakeRootCompendium.updateBalances: eth transfer failed");
     }