feat: allowlist mapping (#226)

Layr-Labs · Apr 15, 2024 · 6454c05 · 6454c05
1 parent f224cf9
commit 6454c05
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 28 deletions.
diff --git a/src/EjectionManager.sol b/src/EjectionManager.sol
@@ -20,8 +20,8 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
     /// @notice the StakeRegistry contract that keeps track of quorum stake
     IStakeRegistry public immutable stakeRegistry;
 
-    /// @notice Address permissioned to eject operators under a ratelimit
-    address public ejector;
+    /// @notice Addresses permissioned to eject operators under a ratelimit
+    mapping(address => bool) public isEjector;
 
     /// @notice Keeps track of the total stake ejected for a quorum 
     mapping(uint8 => StakeEjection[]) public stakeEjectedForQuorum;
@@ -40,17 +40,18 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
 
     /**
      * @param _owner will hold the owner role
-     * @param _ejector will hold the ejector role
+     * @param _ejectors will hold the ejector role
      * @param _quorumEjectionParams are the ratelimit parameters for the quorum at each index
      */
     function initialize(
         address _owner, 
-        address _ejector,
+        address[] memory _ejectors,
         QuorumEjectionParams[] memory _quorumEjectionParams
     ) external initializer {
         _transferOwnership(_owner);
-        _setEjector(_ejector);
-
+        for(uint8 i = 0; i < _ejectors.length; i++) {
+            _setEjector(_ejectors[i], true);
+        }
         for(uint8 i = 0; i < _quorumEjectionParams.length; i++) {
             _setQuorumEjectionParams(i, _quorumEjectionParams[i]);
         }
@@ -63,7 +64,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
      * @dev The owner can eject operators without recording of stake ejection
      */
     function ejectOperators(bytes32[][] memory _operatorIds) external {
-        require(msg.sender == ejector || msg.sender == owner(), "Ejector: Only owner or ejector can eject");
+        require(isEjector[msg.sender] || msg.sender == owner(), "Ejector: Only owner or ejector can eject");
 
         for(uint i = 0; i < _operatorIds.length; ++i) {
             uint8 quorumNumber = uint8(i);
@@ -77,7 +78,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
 
                 //if caller is ejector enforce ratelimit
                 if(
-                    msg.sender == ejector && 
+                    isEjector[msg.sender] && 
                     quorumEjectionParams[quorumNumber].rateLimitWindow > 0 &&
                     stakeForEjection + operatorStake > amountEjectable
                 ){
@@ -88,7 +89,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
                     broke = true;
                     break;
                 }
-                
+
                 //try-catch used to prevent race condition of operator deregistering before ejection
                 try registryCoordinator.ejectOperator(
                     registryCoordinator.getOperatorFromId(_operatorIds[i][j]),
@@ -102,7 +103,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
             }
 
             //record the stake ejected if ejector and ratelimit enforced
-            if(!broke && msg.sender == ejector){ 
+            if(!broke && isEjector[msg.sender]){ 
                 stakeEjectedForQuorum[quorumNumber].push(StakeEjection({
                     timestamp: block.timestamp,
                     stakeEjected: stakeForEjection
@@ -124,9 +125,10 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
     /**
      * @notice Sets the address permissioned to eject operators under a ratelimit
      * @param _ejector The address to permission
+     * @param _status The status to set for the given address
      */
-    function setEjector(address _ejector) external onlyOwner() {
-        _setEjector(_ejector);
+    function setEjector(address _ejector, bool _status) external onlyOwner() {
+        _setEjector(_ejector, _status);
     }
 
     ///@dev internal function to set the quorum ejection params
@@ -136,9 +138,9 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
     }
 
     ///@dev internal function to set the ejector
-    function _setEjector(address _ejector) internal {
-        emit EjectorUpdated(ejector, _ejector);
-        ejector = _ejector;
+    function _setEjector(address _ejector, bool _status) internal {
+        isEjector[_ejector] = _status;
+        emit EjectorUpdated(_ejector, _status);
     }
 
     /**
@@ -154,7 +156,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
             return totalEjectable;
         }
         i = stakeEjectedForQuorum[_quorumNumber].length - 1;
-    
+
         while(stakeEjectedForQuorum[_quorumNumber][i].timestamp > cutoffTime) {
             totalEjected += stakeEjectedForQuorum[_quorumNumber][i].stakeEjected;
             if(i == 0){
@@ -169,4 +171,4 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
         }
         return totalEjectable - totalEjected;   
     }
-}
+}
diff --git a/src/interfaces/IEjectionManager.sol b/src/interfaces/IEjectionManager.sol
@@ -20,14 +20,14 @@ interface IEjectionManager {
     }
 
     ///@notice Emitted when the ejector address is set
-    event EjectorUpdated(address previousAddress, address newAddress);
+    event EjectorUpdated(address ejector, bool status);
     ///@notice Emitted when the ratelimit parameters for a quorum are set
     event QuorumEjectionParamsSet(uint8 quorumNumber, uint32 rateLimitWindow, uint16 ejectableStakePercent);
     ///@notice Emitted when an operator is ejected
     event OperatorEjected(bytes32 operatorId, uint8 quorumNumber);
     ///@notice Emitted when an operator ejection fails
     event FailedOperatorEjection(bytes32 operatorId, uint8 quorumNumber, bytes err);
-    
+
    /**
      * @notice Ejects operators from the AVSs registryCoordinator under a ratelimit
      * @param _operatorIds The ids of the operators to eject for each quorum
@@ -45,11 +45,11 @@ interface IEjectionManager {
      * @notice Sets the address permissioned to eject operators under a ratelimit
      * @param _ejector The address to permission
      */
-    function setEjector(address _ejector) external;
+    function setEjector(address _ejector, bool _status) external;
 
     /**
      * @notice Returns the amount of stake that can be ejected for a quorum at the current block.timestamp
      * @param _quorumNumber The quorum number to view ejectable stake for
      */
     function amountEjectableForQuorum(uint8 _quorumNumber) external view returns (uint256);
-}
+}
diff --git a/test/unit/EjectionManagerUnit.t.sol b/test/unit/EjectionManagerUnit.t.sol
@@ -8,7 +8,7 @@ import "../utils/MockAVSDeployer.sol";
 
 contract EjectionManagerUnitTests is MockAVSDeployer {
 
-    event EjectorUpdated(address previousAddress, address newAddress);
+    event EjectorUpdated(address ejector, bool status);
     event QuorumEjectionParamsSet(uint8 quorumNumber, uint32 rateLimitWindow, uint16 ejectableStakePercent);
     event OperatorEjected(bytes32 operatorId, uint8 quorumNumber);
     event FailedOperatorEjection(bytes32 operatorId, uint8 quorumNumber, bytes err);
@@ -42,14 +42,17 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
 
         ejectionManagerImplementation = new EjectionManager(registryCoordinator, stakeRegistry);
 
+        address[] memory ejectors = new address[](1);
+        ejectors[0] = ejector;
+
         cheats.prank(proxyAdminOwner);
         proxyAdmin.upgradeAndCall(
             TransparentUpgradeableProxy(payable(address(ejectionManager))),
             address(ejectionManagerImplementation),
             abi.encodeWithSelector(
                 EjectionManager.initialize.selector,
                 registryCoordinatorOwner,
-                ejector,
+                ejectors,
                 quorumEjectionParams
             )
         );
@@ -82,7 +85,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
                 emit OperatorEjected(operatorIds[i][j], i);
             }
         }
-        
+
         cheats.prank(ejector);
         ejectionManager.ejectOperators(operatorIds);
 
@@ -352,12 +355,12 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
 
     function testSetEjector() public {
         cheats.expectEmit(true, true, true, true, address(ejectionManager));
-        emit EjectorUpdated(ejector, address(0));
+        emit EjectorUpdated(address(0), true);
 
         cheats.prank(registryCoordinatorOwner);
-        ejectionManager.setEjector(address(0));
+        ejectionManager.setEjector(address(0), true);
 
-        assertEq(ejectionManager.ejector(), address(0));
+        assertEq(ejectionManager.isEjector(address(0)), true);
     }
 
     function test_Revert_NotPermissioned() public {
@@ -370,7 +373,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         ejectionManager.setQuorumEjectionParams(0, _quorumEjectionParams);
 
         cheats.expectRevert("Ownable: caller is not the owner");
-        ejectionManager.setEjector(address(0));
+        ejectionManager.setEjector(address(0), true);
     }
 
     function _registerOperaters(uint8 numOperators, uint96 stake) internal {