chore: remove redundancy

stader-labs · Dec 6, 2024 · 2118e39 · 2118e39
1 parent 963a6a5
commit 2118e39
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 36 deletions.
diff --git a/contracts/OperatorRewardsCollector.sol b/contracts/OperatorRewardsCollector.sol
@@ -53,11 +53,8 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
      * @dev This function first checks for any unpaid liquidations for the operator and repays them if necessary. Then, it transfers any remaining balance to the operator's reward address.
      */
     function claim() external {
-        IPoolUtils poolUtils = IPoolUtils(staderConfig.getPoolUtils());
-        uint8 poolId = poolUtils.getOperatorPoolId(msg.sender);
-        address permissionlessNodeRegistry = staderConfig.getPermissionlessNodeRegistry();
         uint256 amount;
-        if (INodeRegistry(permissionlessNodeRegistry).POOL_ID() == poolId) {
+        if (_isPermissionlessCaller(msg.sender)) {
             claimLiquidation(msg.sender);
             amount = balances[msg.sender] > withdrawableInEth(msg.sender)
                 ? withdrawableInEth(msg.sender)
@@ -75,17 +72,12 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
      * @param _amount amount of ETH to claim
      */
     function claimWithAmount(uint256 _amount) external {
-        IPoolUtils poolUtils = IPoolUtils(staderConfig.getPoolUtils());
-        uint8 poolId = poolUtils.getOperatorPoolId(msg.sender);
-        address permissionlessNodeRegistry = staderConfig.getPermissionlessNodeRegistry();
-
-        if (INodeRegistry(permissionlessNodeRegistry).POOL_ID() == poolId) {
+        if (_isPermissionlessCaller(msg.sender)) {
             claimLiquidation(msg.sender);
             uint256 maxWithdrawableInEth = withdrawableInEth(msg.sender);
-            if (_amount > maxWithdrawableInEth || _amount > balances[msg.sender]) revert InsufficientBalance();
-        } else {
-            if (_amount > balances[msg.sender]) revert InsufficientBalance();
+            if (_amount > maxWithdrawableInEth) revert InsufficientBalance();
         }
+        if (_amount > balances[msg.sender]) revert InsufficientBalance();
         _claim(msg.sender, _amount);
     }
 
@@ -132,6 +124,14 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
         return balances[operator];
     }
 
+    function _isPermissionlessCaller(address caller) internal returns (bool) {
+        IPoolUtils poolUtils = IPoolUtils(staderConfig.getPoolUtils());
+        uint8 poolId = poolUtils.getOperatorPoolId(caller);
+        address permissionlessNodeRegistry = staderConfig.getPermissionlessNodeRegistry();
+
+        return INodeRegistry(permissionlessNodeRegistry).POOL_ID() == poolId;
+    }
+
     /**
      * @notice Completes any pending liquidation for an operator if exists.
      * @dev Internal function to handle liquidation completion.

diff --git a/contracts/PermissionedNodeRegistry.sol b/contracts/PermissionedNodeRegistry.sol
@@ -432,17 +432,6 @@ contract PermissionedNodeRegistry is
         emit UpdatedOperatorName(msg.sender, _operatorName);
     }
 
-    /**
-     * @notice update the maximum non terminal key limit per operator
-     * @dev only `MANAGER` role can call
-     * @param _maxNonTerminalKeyPerOperator updated maximum non terminal key per operator limit
-     */
-    function updateMaxNonTerminalKeyPerOperator(uint64 _maxNonTerminalKeyPerOperator) external override {
-        UtilLib.onlyManagerRole(msg.sender, staderConfig);
-        maxNonTerminalKeyPerOperator = _maxNonTerminalKeyPerOperator;
-        emit UpdatedMaxNonTerminalKeyPerOperator(maxNonTerminalKeyPerOperator);
-    }
-
     /**
      * @notice update number of validator keys that can be added in a single tx by the operator
      * @dev only `OPERATOR` role can call

diff --git a/contracts/interfaces/IPermissionedNodeRegistry.sol b/contracts/interfaces/IPermissionedNodeRegistry.sol
@@ -62,8 +62,6 @@ interface IPermissionedNodeRegistry {
 
     function markValidatorStatusAsPreDeposit(bytes calldata _pubkey) external;
 
-    function updateMaxNonTerminalKeyPerOperator(uint64 _maxNonTerminalKeyPerOperator) external;
-
     function updateInputKeyCountLimit(uint16 _inputKeyCountLimit) external;
 
     function proposeRewardAddress(address _operatorAddress, address _newRewardAddress) external;

diff --git a/test/foundry_tests/PermissionedNodeRegistry.t.sol b/test/foundry_tests/PermissionedNodeRegistry.t.sol
@@ -375,17 +375,6 @@ contract PermissionedNodeRegistryTest is Test {
         assertEq(nodeRegistry.inputKeyCountLimit(), _keyCountLimit);
     }
 
-    function test_updateMaxNonTerminalKeyPerOperator(uint64 _maxNonTerminalKeyPerOperator) public {
-        vm.prank(staderManager);
-        nodeRegistry.updateMaxNonTerminalKeyPerOperator(_maxNonTerminalKeyPerOperator);
-        assertEq(nodeRegistry.maxNonTerminalKeyPerOperator(), _maxNonTerminalKeyPerOperator);
-    }
-
-    function testFail_updateMaxNonTerminalKeyPerOperator(uint64 _maxNonTerminalKeyPerOperator) public {
-        nodeRegistry.updateMaxNonTerminalKeyPerOperator(_maxNonTerminalKeyPerOperator);
-        assertEq(nodeRegistry.maxNonTerminalKeyPerOperator(), _maxNonTerminalKeyPerOperator);
-    }
-
     function test_updateVerifiedKeysBatchSize(uint256 _verifiedKeysBatchSize) public {
         vm.prank(operator);
         nodeRegistry.updateVerifiedKeysBatchSize(_verifiedKeysBatchSize);