feat(op sets): update stakes when forceUnregister (#300)

stevennevins · web-flow · commit 86f0928dd603 · 2024-08-21T14:27:00.000-04:00
* feat: update stakes handle direct deregistration on AVSDirectory

* test: update stake for quorum if operator directly unregistered from the AVSDirectory

* chore: simplify setup

* chore: simplify setup

* chore: make service manager immutable on stakeRegistry
diff --git a/src/StakeRegistry.sol b/src/StakeRegistry.sol
@@ -2,6 +2,8 @@
 pragma solidity ^0.8.12;
 
 import {IDelegationManager} from "eigenlayer-contracts/src/contracts/interfaces/IDelegationManager.sol";
+import {IAVSDirectory} from "eigenlayer-contracts/src/contracts/interfaces/IAVSDirectory.sol";
+import {IServiceManager} from "./interfaces/IServiceManager.sol";
 
 import {StakeRegistryStorage, IStrategy} from "./StakeRegistryStorage.sol";
 
@@ -40,8 +42,10 @@ contract StakeRegistry is StakeRegistryStorage {
 
     constructor(
         IRegistryCoordinator _registryCoordinator,
-        IDelegationManager _delegationManager
-    ) StakeRegistryStorage(_registryCoordinator, _delegationManager) {}
+        IDelegationManager _delegationManager,
+        IAVSDirectory _avsDirectory,
+        IServiceManager _serviceManager
+    ) StakeRegistryStorage(_registryCoordinator, _delegationManager, _avsDirectory, _serviceManager) {}
 
     /*******************************************************************************
                       EXTERNAL FUNCTIONS - REGISTRY COORDINATOR
@@ -148,6 +152,10 @@ contract StakeRegistry is StakeRegistryStorage {
     ) external onlyRegistryCoordinator returns (uint192) {
         uint192 quorumsToRemove;
 
+        bool isOperatorSetAVS = avsDirectory.isOperatorSetAVS(
+                address(serviceManager)
+            );
+
         /**
          * For each quorum, update the operator's stake and record the delta
          * in the quorum's total stake.
@@ -163,9 +171,21 @@ contract StakeRegistry is StakeRegistryStorage {
             // Fetch the operator's current stake, applying weighting parameters and checking
             // against the minimum stake requirements for the quorum.
             (uint96 stakeWeight, bool hasMinimumStake) = _weightOfOperatorForQuorum(quorumNumber, operator);
-
             // If the operator no longer meets the minimum stake, set their stake to zero and mark them for removal
-            if (!hasMinimumStake) {
+            /// also handle setting the operator's stake to 0 and remove them from the quorum 
+            /// if they directly unregistered from the AVSDirectory bubbles up info via registry coordinator to deregister them
+            bool operatorRegistered;
+            // Convert quorumNumber to operatorSetId
+            uint32 operatorSetId = uint32(quorumNumber);
+
+            // Get the AVSDirectory address from the RegistryCoordinator
+            // Query the AVSDirectory to check if the operator is directly unregistered
+            operatorRegistered = avsDirectory.isMember(
+                operator,
+                IAVSDirectory.OperatorSet(address(serviceManager), operatorSetId)
+            );
+
+            if (!hasMinimumStake || (isOperatorSetAVS && !operatorRegistered)) {
                 stakeWeight = 0;
                 quorumsToRemove = uint192(quorumsToRemove.setBit(quorumNumber));
             }
diff --git a/src/StakeRegistryStorage.sol b/src/StakeRegistryStorage.sol
@@ -2,6 +2,8 @@
 pragma solidity ^0.8.12;
 
 import {IDelegationManager} from "eigenlayer-contracts/src/contracts/interfaces/IDelegationManager.sol";
+import {IAVSDirectory} from "eigenlayer-contracts/src/contracts/interfaces/IAVSDirectory.sol";
+import {IServiceManager} from "./interfaces/IServiceManager.sol";
 import {IStrategyManager, IStrategy} from "eigenlayer-contracts/src/contracts/interfaces/IStrategyManager.sol";
 
 import {IRegistryCoordinator} from "./interfaces/IRegistryCoordinator.sol";
@@ -24,6 +26,12 @@ abstract contract StakeRegistryStorage is IStakeRegistry {
     /// @notice The address of the Delegation contract for EigenLayer.
     IDelegationManager public immutable delegation;
 
+    /// @notice The address of the Delegation contract for EigenLayer.
+    IAVSDirectory public immutable avsDirectory;
+
+    /// @notice the address of the ServiceManager associtated with the stake registries
+    IServiceManager public immutable serviceManager;
+
     /// @notice the coordinator contract that this registry is associated with
     address public immutable registryCoordinator;
 
@@ -47,10 +55,14 @@ abstract contract StakeRegistryStorage is IStakeRegistry {
 
     constructor(
         IRegistryCoordinator _registryCoordinator, 
-        IDelegationManager _delegationManager
+        IDelegationManager _delegationManager,
+        IAVSDirectory _avsDirectory,
+        IServiceManager _serviceManager
     ) {
         registryCoordinator = address(_registryCoordinator);
         delegation = _delegationManager;
+        avsDirectory = _avsDirectory;
+        serviceManager = _serviceManager;
     }
 
     // storage gap for upgradeability
diff --git a/src/interfaces/IRegistryCoordinator.sol b/src/interfaces/IRegistryCoordinator.sol
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity ^0.8.12;
 
+import {IServiceManager} from "./IServiceManager.sol";
 import {IBLSApkRegistry} from "./IBLSApkRegistry.sol";
 import {IStakeRegistry} from "./IStakeRegistry.sol";
 import {IIndexRegistry} from "./IIndexRegistry.sol";
@@ -150,4 +151,6 @@ interface IRegistryCoordinator {
 
     /// @notice The owner of the registry coordinator
     function owner() external view returns (address);
+
+    function serviceManager() external view returns (IServiceManager);
 }
diff --git a/test/harnesses/StakeRegistryHarness.sol b/test/harnesses/StakeRegistryHarness.sol
@@ -7,8 +7,10 @@ import "../../src/StakeRegistry.sol";
 contract StakeRegistryHarness is StakeRegistry {
     constructor(
         IRegistryCoordinator _registryCoordinator,
-        IDelegationManager _delegationManager
-    ) StakeRegistry(_registryCoordinator, _delegationManager) {
+        IDelegationManager _delegationManager,
+        IAVSDirectory _avsDirectory,
+        IServiceManager _serviceManager
+    ) StakeRegistry(_registryCoordinator, _delegationManager, _avsDirectory, _serviceManager) {
     }
 
     function recordOperatorStakeUpdate(bytes32 operatorId, uint8 quorumNumber, uint96 newStake) external returns(int256) {
diff --git a/test/integration/IntegrationDeployer.t.sol b/test/integration/IntegrationDeployer.t.sol
@@ -330,7 +330,7 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
         cheats.stopPrank();
 
         StakeRegistry stakeRegistryImplementation = new StakeRegistry(
-            IRegistryCoordinator(registryCoordinator), IDelegationManager(delegationManager)
+            IRegistryCoordinator(registryCoordinator), IDelegationManager(delegationManager), IAVSDirectory(avsDirectory), IServiceManager(serviceManager)
         );
         BLSApkRegistry blsApkRegistryImplementation =
             new BLSApkRegistry(IRegistryCoordinator(registryCoordinator));
diff --git a/test/mocks/RegistryCoordinatorMock.sol b/test/mocks/RegistryCoordinatorMock.sol
@@ -68,4 +68,6 @@ contract RegistryCoordinatorMock is IRegistryCoordinator {
     function quorumUpdateBlockNumber(uint8 quorumNumber) external view returns (uint256) {}
 
     function owner() external view returns (address) {}
-}
+
+    function serviceManager() external view returns (IServiceManager){}
+}
diff --git a/test/unit/RegistryCoordinatorMigration.t.sol b/test/unit/RegistryCoordinatorMigration.t.sol
@@ -197,4 +197,89 @@ contract RegistryCoordinatorMigrationUnit is MockAVSDeployer, IServiceManagerBas
         assertTrue(avsDirectory.isOperatorSet(address(serviceManager), quorumNumber), "Operator set was not created for the quorum");
 
     }
+
+    function test_updateOperatorsForQuorumsAfterDirectUnregister() public {
+        vm.prank(proxyAdmin.owner());
+        proxyAdmin.upgrade(
+            TransparentUpgradeableProxy(payable(address(avsDirectory))),
+            address(avsDirectoryMock)
+        );
+        uint256 pseudoRandomNumber = uint256(keccak256("pseudoRandomNumber"));
+        _registerRandomOperators(pseudoRandomNumber);
+
+        vm.prank(proxyAdmin.owner());
+        proxyAdmin.upgrade(
+            TransparentUpgradeableProxy(payable(address(avsDirectory))),
+            address(avsDirectoryHarness)
+        );
+
+        uint256 quorumCount = registryCoordinator.quorumCount();
+        for (uint256 i = 0; i < quorumCount; i++) {
+            uint256 operatorCount = indexRegistry.totalOperatorsForQuorum(uint8(i));
+            bytes32[] memory operatorIds =
+                indexRegistry.getOperatorListAtBlockNumber(uint8(i), uint32(block.number));
+            assertEq(operatorCount, operatorIds.length, "Operator Id length mismatch"); // sanity check
+            for (uint256 j = 0; j < operatorCount; j++) {
+                address operatorAddress =
+                 registryCoordinator.blsApkRegistry().getOperatorFromPubkeyHash(operatorIds[j]);
+                AVSDirectoryHarness(address(avsDirectory)).setAvsOperatorStatus(
+                    address(serviceManager),
+                    operatorAddress,
+                    IAVSDirectory.OperatorAVSRegistrationStatus.REGISTERED
+                );
+            }
+        }
+
+        (
+            uint32[] memory operatorSetsToCreate,
+            uint32[][] memory operatorSetIdsToMigrate,
+            address[] memory operators
+        ) = serviceManager.getOperatorsToMigrate();
+        cheats.startPrank(serviceManagerOwner);
+        serviceManager.migrateAndCreateOperatorSetIds(operatorSetsToCreate);
+        serviceManager.migrateToOperatorSets(operatorSetIdsToMigrate, operators);
+        cheats.stopPrank();
+
+        bytes32[] memory registeredOperators = indexRegistry.getOperatorListAtBlockNumber(defaultQuorumNumber, uint32(block.number));
+        uint256 preNumOperators = registeredOperators.length;
+        address[] memory registeredOperatorAddresses = new address[](registeredOperators.length);
+        for (uint256 i = 0; i < registeredOperators.length; i++) {
+            registeredOperatorAddresses[i] = registryCoordinator.blsApkRegistry().pubkeyHashToOperator(registeredOperators[i]);
+        }
+
+        uint32[] memory operatorSetsToUnregister = new uint32[](1);
+        operatorSetsToUnregister[0] = defaultQuorumNumber;
+
+        vm.prank(operators[0]);
+        avsDirectory.forceDeregisterFromOperatorSets(
+            operators[0], 
+            address(serviceManager), 
+            operatorSetsToUnregister,
+            ISignatureUtils.SignatureWithSaltAndExpiry({
+                signature: new bytes(0),
+                salt: bytes32(0),
+                expiry: 0
+            })
+        );
+        // sanity check if the operator was unregistered from the intended operator set
+        bool operatorIsUnRegistered = !avsDirectory.isMember(operators[0], IAVSDirectory.OperatorSet({
+            avs: address(serviceManager),
+            operatorSetId: defaultQuorumNumber
+        }));
+        bool isOperatorSetAVS = avsDirectory.isOperatorSetAVS(address(serviceManager));
+        assertTrue(isOperatorSetAVS, "ServiceManager is not an operator set AVS");
+        assertTrue(operatorIsUnRegistered, "Operator wasnt unregistered from op set");
+
+        address[][] memory registeredOperatorAddresses2D = new address[][](1);
+        registeredOperatorAddresses2D[0] = registeredOperatorAddresses;
+        bytes memory quorumNumbers = new bytes(1);
+        quorumNumbers[0] = bytes1(defaultQuorumNumber);
+        registryCoordinator.updateOperatorsForQuorum(registeredOperatorAddresses2D, quorumNumbers);
+
+        registeredOperators = indexRegistry.getOperatorListAtBlockNumber(defaultQuorumNumber, uint32(block.number));
+        uint256 postRegisteredOperators = registeredOperators.length;
+
+        assertEq(preNumOperators-1, postRegisteredOperators, "");
+
+    }
 }
diff --git a/test/unit/StakeRegistryUnit.t.sol b/test/unit/StakeRegistryUnit.t.sol
@@ -52,7 +52,7 @@ contract StakeRegistryUnitTests is MockAVSDeployer, IStakeRegistryEvents {
         );
 
         stakeRegistryImplementation = new StakeRegistryHarness(
-            IRegistryCoordinator(address(registryCoordinator)), delegationMock
+            IRegistryCoordinator(address(registryCoordinator)), delegationMock, avsDirectoryMock, serviceManager
         );
 
         stakeRegistry = StakeRegistryHarness(
diff --git a/test/utils/MockAVSDeployer.sol b/test/utils/MockAVSDeployer.sol
@@ -222,7 +222,7 @@ contract MockAVSDeployer is Test {
         cheats.startPrank(proxyAdminOwner);
 
         stakeRegistryImplementation =
-            new StakeRegistryHarness(IRegistryCoordinator(registryCoordinator), delegationMock);
+            new StakeRegistryHarness(IRegistryCoordinator(registryCoordinator), delegationMock, avsDirectory, serviceManager);
 
         proxyAdmin.upgrade(
             TransparentUpgradeableProxy(payable(address(stakeRegistry))),
