From 6454c05beec77a165211f081b50905c516fc7777 Mon Sep 17 00:00:00 2001 From: quaq <56312047+0x0aa0@users.noreply.github.com> Date: Mon, 15 Apr 2024 10:48:40 -0500 Subject: [PATCH] feat: allowlist mapping (#226) --- src/EjectionManager.sol | 36 +++++++++++++++-------------- src/interfaces/IEjectionManager.sol | 8 +++---- test/unit/EjectionManagerUnit.t.sol | 17 ++++++++------ 3 files changed, 33 insertions(+), 28 deletions(-) diff --git a/src/EjectionManager.sol b/src/EjectionManager.sol index d6ed1b46..a9a4bc77 100644 --- 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; } -} +} \ No newline at end of file diff --git a/src/interfaces/IEjectionManager.sol b/src/interfaces/IEjectionManager.sol index 94aebd0f..83299484 100644 --- 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); -} +} \ No newline at end of file diff --git a/test/unit/EjectionManagerUnit.t.sol b/test/unit/EjectionManagerUnit.t.sol index 59e5fd01..82a22daa 100644 --- 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,6 +42,9 @@ 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))), @@ -49,7 +52,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer { 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 {