diff --git a/contracts/PoolSelector.sol b/contracts/PoolSelector.sol index d2054637..391ec178 100644 --- a/contracts/PoolSelector.sol +++ b/contracts/PoolSelector.sol @@ -22,6 +22,9 @@ contract PoolSelector is IPoolSelector, AccessControlUpgradeable { mapping(uint8 => uint256) public poolWeights; + bytes32 public constant POOL_WEIGHTS_OPERATOR = keccak256("POOL_WEIGHTS_OPERATOR"); + error CallerNotPoolWeightsOperator(); + /// @custom:oz-upgrades-unsafe-allow constructor constructor() { _disableInitializers(); @@ -108,12 +111,14 @@ contract PoolSelector is IPoolSelector, AccessControlUpgradeable { /** * @notice update the target weights of existing pools - * @dev only `Manager` can call, + * @dev only `POOL_WEIGHTS_OPERATOR's` can call, * @param _poolTargets new target weights of pools * `_poolTargets` array provide pool target in the same order of poolIDs that are stored in poolIdArray of poolUtils */ function updatePoolWeights(uint256[] calldata _poolTargets) external { - UtilLib.onlyManagerRole(msg.sender, staderConfig); + if(!onlyPoolWeightsOperatorRole(msg.sender)){ + revert CallerNotPoolWeightsOperator(); + } uint8[] memory poolIdArray = IPoolUtils(staderConfig.getPoolUtils()).getPoolIdArray(); uint256 poolCount = poolIdArray.length; uint256 poolTargetLength = _poolTargets.length; @@ -145,4 +150,8 @@ contract PoolSelector is IPoolSelector, AccessControlUpgradeable { staderConfig = IStaderConfig(_staderConfig); emit UpdatedStaderConfig(_staderConfig); } + + function onlyPoolWeightsOperatorRole(address account) private view returns (bool) { + return hasRole(POOL_WEIGHTS_OPERATOR, account); + } } diff --git a/test/foundry_tests/PoolSelector.t.sol b/test/foundry_tests/PoolSelector.t.sol index baf7ef39..537968b2 100644 --- a/test/foundry_tests/PoolSelector.t.sol +++ b/test/foundry_tests/PoolSelector.t.sol @@ -23,6 +23,8 @@ contract PoolSelectorTest is Test { PoolSelector poolSelector; PoolUtilsMockForDepositFlow poolUtils; + error CallerNotPoolWeightsOperator(); + function setUp() public { vm.clearMockedCalls(); staderAdmin = vm.addr(100); @@ -55,6 +57,7 @@ contract PoolSelectorTest is Test { poolUtils = new PoolUtilsMockForDepositFlow(address(0), address(staderConfig)); vm.startPrank(staderAdmin); + poolSelector.grantRole(poolSelector.POOL_WEIGHTS_OPERATOR(), staderManager); staderConfig.updatePoolUtils(address(poolUtils)); staderConfig.updateStakePoolManager(staderStakePoolManager); staderConfig.grantRole(staderConfig.MANAGER(), staderManager); @@ -95,7 +98,7 @@ contract PoolSelectorTest is Test { invalidSizePoolWeight[1] = 4000; invalidSizePoolWeight[2] = 4000; - vm.expectRevert(UtilLib.CallerNotManager.selector); + vm.expectRevert(CallerNotPoolWeightsOperator.selector); poolSelector.updatePoolWeights(poolWeight); vm.startPrank(staderManager);