From 2118e39136259a99f0835b40fc52ed1365a3c1fc Mon Sep 17 00:00:00 2001
From: blockgroot <170620375+blockgroot@users.noreply.github.com>
Date: Fri, 6 Dec 2024 07:16:59 +0000
Subject: [PATCH] chore: remove redundancy

---
 contracts/OperatorRewardsCollector.sol        | 24 +++++++++----------
 contracts/PermissionedNodeRegistry.sol        | 11 ---------
 .../interfaces/IPermissionedNodeRegistry.sol  |  2 --
 .../PermissionedNodeRegistry.t.sol            | 11 ---------
 4 files changed, 12 insertions(+), 36 deletions(-)

diff --git a/contracts/OperatorRewardsCollector.sol b/contracts/OperatorRewardsCollector.sol
index 9323f324..cb08a2b6 100644
--- 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
index 545e4e38..a24ae3d6 100644
--- 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
index f0f76ee9..d7ae04d0 100644
--- 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
index e4f79fcc..559ab445 100644
--- 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);