feat: remove unnecessary steps while deposit and claim

Author: blockgroot

contracts/OperatorRewardsCollector.sol

@@ -16,6 +16,7 @@ import { ISDUtilityPool, UserData, OperatorLiquidation } from "./interfaces/ISDU
 import { ISDCollateral } from "./interfaces/SDCollateral/ISDCollateral.sol";
 import { IWETH } from "./interfaces/IWETH.sol";
 import { IStaderOracle } from "../contracts/interfaces/IStaderOracle.sol";
+import { IPoolUtils } from "../contracts/interfaces/IPoolUtils.sol";
 
 contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpgradeable {
     IStaderConfig public staderConfig;
@@ -52,10 +53,18 @@ 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 {
-        claimLiquidation(msg.sender);
-        uint256 amount = balances[msg.sender] > withdrawableInEth(msg.sender)
-            ? withdrawableInEth(msg.sender)
-            : balances[msg.sender];
+        IPoolUtils poolUtils = IPoolUtils(staderConfig.getPoolUtils());
+        uint8 poolId = poolUtils.getOperatorPoolId(msg.sender);
+        address permissionlessNodeRegistry = staderConfig.getPermissionlessNodeRegistry();
+        uint256 amount;
+        if (INodeRegistry(permissionlessNodeRegistry).POOL_ID() == poolId) {
+            claimLiquidation(msg.sender);
+            amount = balances[msg.sender] > withdrawableInEth(msg.sender)
+                ? withdrawableInEth(msg.sender)
+                : balances[msg.sender];
+        } else {
+            amount = balances[msg.sender];
+        }
         _claim(msg.sender, amount);
     }
 
@@ -66,7 +75,17 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
      * @param _amount amount of ETH to claim
      */
     function claimWithAmount(uint256 _amount) external {
-        claimLiquidation(msg.sender);
+        IPoolUtils poolUtils = IPoolUtils(staderConfig.getPoolUtils());
+        uint8 poolId = poolUtils.getOperatorPoolId(msg.sender);
+        address permissionlessNodeRegistry = staderConfig.getPermissionlessNodeRegistry();
+
+        if (INodeRegistry(permissionlessNodeRegistry).POOL_ID() == poolId) {
+            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();
+        }
         _claim(msg.sender, _amount);
     }
 
@@ -181,9 +200,6 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
      * @param amount The amount to be claimed.
      */
     function _claim(address operator, uint256 amount) internal {
-        uint256 maxWithdrawableInEth = withdrawableInEth(operator);
-        if (amount > maxWithdrawableInEth || amount > balances[operator]) revert InsufficientBalance();
-
         balances[operator] -= amount;
 
         // If there's an amount to send, transfer it to the operator's rewards address

contracts/PermissionedNodeRegistry.sol

@@ -720,22 +720,6 @@ contract PermissionedNodeRegistry is
             revert InvalidKeyCount();
         }
         totalKeys = getOperatorTotalKeys(_operatorId);
-        uint256 totalNonTerminalKeys = getOperatorTotalNonTerminalKeys(msg.sender, 0, totalKeys);
-        if ((totalNonTerminalKeys + keyCount) > maxNonTerminalKeyPerOperator) {
-            revert MaxKeyLimitReached();
-        }
-
-        //checks if operator has enough SD collateral for adding `keyCount` keys
-        //SD threshold for permissioned NOs is 0 for phase1
-        if (
-            !ISDCollateral(staderConfig.getSDCollateral()).hasEnoughSDCollateral(
-                msg.sender,
-                POOL_ID,
-                totalNonTerminalKeys + keyCount
-            )
-        ) {
-            revert NotEnoughSDCollateral();
-        }
     }
 
     // operator in active state

lib/openzeppelin-foundry-upgrades

@@ -0,0 +1 @@
+Subproject commit 16e0ae21e0e39049f619f2396fa28c57fad07368

test/fork/reward-claim-optimisation.ts

@@ -0,0 +1,56 @@
+import { ethers, network } from "hardhat";
+import "dotenv/config";
+import { impersonateAccount, setBalance } from "@nomicfoundation/hardhat-network-helpers";
+
+const PROXY_OWNER = "0x1112D5C55670Cb5144BF36114C20a122908068B9"
+const PROXY_ADMIN = "0x67B12264Ca3e0037Fc7E22F2457b42643a04C86e";
+const OPERATOR_REWARDS_COLLECTOR_ADDRESS = "0x84ffDC9De310144D889540A49052F6d1AdB2C335";
+const OPERATOR = "0xb851788Fa34B0d9215F54531061D4e2e06A74AEE"
+ 
+async function setForkBlock(blockNumber: number) {
+  await network.provider.request({
+    method: "hardhat_reset",
+    params: [
+      {
+        forking: {
+          jsonRpcUrl: process.env.PROVIDER_URL_MAINNET,
+          blockNumber: blockNumber,
+        },
+      },
+    ],
+  });
+}
+
+async function configureNewContract (contractName: String, contractAddress: String) {
+  await setBalance(PROXY_OWNER, ethers.parseEther("1"))
+  await impersonateAccount(PROXY_OWNER)
+
+  const impersonatedProxyOwner = await ethers.getSigner(PROXY_OWNER);
+  
+  const contractFactory = await ethers.getContractFactory(contractName);
+  const contractImpl = await contractFactory.deploy();
+  console.log(`${contractName} Implementation deployed to:`, await contractImpl.getAddress());
+
+  const proxyAdminContract = await ethers.getContractAt("ProxyAdmin", PROXY_ADMIN);
+  await proxyAdminContract.connect(impersonatedProxyOwner).upgrade(contractAddress, await contractImpl.getAddress());
+
+  const contract = await ethers.getContractAt(contractName, contractAddress)
+  
+  return contract;
+}
+
+describe("Gas Coverage", function () {
+  it("should consume less gas after upgrade", async () => {
+    await setForkBlock(21270988);
+    await setBalance(OPERATOR, ethers.parseEther("100"))
+    await impersonateAccount(OPERATOR);
+    const impersonatedOperator = await ethers.getSigner(OPERATOR);
+
+    const newOperatorRewardsCollector = await configureNewContract("OperatorRewardsCollector", OPERATOR_REWARDS_COLLECTOR_ADDRESS);
+    
+    // Firing a txn with updated contracts
+    const claimTxn = await newOperatorRewardsCollector.connect(impersonatedOperator).claim();
+    const claimTxnReceipt = await claimTxn.wait();
+    console.log("Claim Txn Gas Estimate:", claimTxnReceipt.gasUsed);
+  });
+});

test/foundry_tests/NodeELRewardVault.t.sol

@@ -19,6 +19,7 @@ import { PoolUtilsMock } from "../mocks/PoolUtilsMock.sol";
 import { StakePoolManagerMock } from "../mocks/StakePoolManagerMock.sol";
 import { StaderOracleMock } from "../mocks/StaderOracleMock.sol";
 import { SDUtilityPoolMock } from "../mocks/SDUtilityPoolMock.sol";
+import { PermissionlessNodeRegistryMock } from "../mocks/PermissionlessNodeRegistryMock.sol";
 
 contract NodeELRewardVaultTest is Test {
     address private constant OPERATOR_ADDRESSS = address(500);
@@ -59,6 +60,7 @@ contract NodeELRewardVaultTest is Test {
         address operator = OPERATOR_ADDRESSS;
         mockStaderOracle(staderOracleMock);
         mockSdUtilityPool(sdUtilityPoolMock, operator);
+        mockPermissionlessNodeRegistry(vm.addr(105));
 
         OperatorRewardsCollector operatorRCImpl = new OperatorRewardsCollector();
         TransparentUpgradeableProxy operatorRCProxy = new TransparentUpgradeableProxy(
@@ -95,6 +97,7 @@ contract NodeELRewardVaultTest is Test {
         staderConfig.updateSDCollateral(address(sdCollateral));
         staderConfig.updateSDUtilityPool(sdUtilityPoolMock);
         staderConfig.updateStaderOracle(staderOracleMock);
+        staderConfig.updatePermissionlessNodeRegistry(vm.addr(105));
         staderConfig.grantRole(staderConfig.MANAGER(), staderManager);
         vaultFactory.grantRole(vaultFactory.NODE_REGISTRY_CONTRACT(), address(poolUtils.nodeRegistry()));
         vm.stopPrank();
@@ -265,4 +268,11 @@ contract NodeELRewardVaultTest is Test {
         bytes memory mockCode = address(implementation).code;
         vm.etch(staderOracleMock, mockCode);
     }
+
+    function mockPermissionlessNodeRegistry(address _permissionlessNodeRegistry) private {
+        emit log_named_address("permissionlessNodeRegistry", _permissionlessNodeRegistry);
+        PermissionlessNodeRegistryMock nodeRegistryMock = new PermissionlessNodeRegistryMock();
+        bytes memory mockCode = address(nodeRegistryMock).code;
+        vm.etch(_permissionlessNodeRegistry, mockCode);
+    }
 }

test/foundry_tests/PermissionedNodeRegistry.t.sol

@@ -227,24 +227,6 @@ contract PermissionedNodeRegistryTest is Test {
         assertEq(nodeRegistry.isExistingPubkey(pubkeys[0]), true);
     }
 
-    function testAddValidatorKeysNotEnoughSDCollateral() public {
-        (
-            bytes[] memory pubkeys,
-            bytes[] memory preDepositSignature,
-            bytes[] memory depositSignature
-        ) = getValidatorKeys();
-        vm.startPrank(permissionedNO);
-        nodeRegistry.onboardNodeOperator("testOP", payable(address(this)));
-        vm.mockCall(
-            address(sdCollateral),
-            abi.encodeWithSelector(ISDCollateral.hasEnoughSDCollateral.selector),
-            abi.encode(false)
-        );
-        vm.expectRevert(INodeRegistry.NotEnoughSDCollateral.selector);
-        nodeRegistry.addValidatorKeys(pubkeys, preDepositSignature, depositSignature);
-        vm.stopPrank();
-    }
-
     function test_addValidatorKeysWithMisMatchingInputs() public {
         bytes[] memory pubkeys = new bytes[](1);
         bytes[] memory preDepositSignature = new bytes[](1);
@@ -271,40 +253,6 @@ contract PermissionedNodeRegistryTest is Test {
         vm.stopPrank();
     }
 
-    function test_addValidatorKeysOPCrossingMaxNonTerminalKeys() public {
-        (
-            bytes[] memory pubkeys,
-            bytes[] memory preDepositSignature,
-            bytes[] memory depositSignature
-        ) = getValidatorKeys();
-        vm.prank(staderManager);
-        nodeRegistry.updateMaxNonTerminalKeyPerOperator(2);
-        vm.startPrank(permissionedNO);
-        nodeRegistry.onboardNodeOperator("testOP", payable(address(this)));
-        vm.expectRevert(INodeRegistry.MaxKeyLimitReached.selector);
-        nodeRegistry.addValidatorKeys(pubkeys, preDepositSignature, depositSignature);
-        vm.stopPrank();
-    }
-
-    function test_addValidatorKeysWithInsufficientSDCollateral() public {
-        (
-            bytes[] memory pubkeys,
-            bytes[] memory preDepositSignature,
-            bytes[] memory depositSignature
-        ) = getValidatorKeys();
-
-        vm.startPrank(permissionedNO);
-        nodeRegistry.onboardNodeOperator("testOP", payable(address(this)));
-        vm.mockCall(
-            address(sdCollateral),
-            abi.encodeWithSelector(ISDCollateral.hasEnoughSDCollateral.selector),
-            abi.encode(false)
-        );
-        vm.expectRevert(INodeRegistry.NotEnoughSDCollateral.selector);
-        nodeRegistry.addValidatorKeys(pubkeys, preDepositSignature, depositSignature);
-        vm.stopPrank();
-    }
-
     function test_markReadyToDepositValidator() public {
         (
             bytes[] memory pubkeys,

test/foundry_tests/ValidatorWithdrawalVault.t.sol

@@ -21,6 +21,7 @@ import "../mocks/StakePoolManagerMock.sol";
 
 import { SDUtilityPoolMock } from "../mocks/SDUtilityPoolMock.sol";
 import { StaderOracleMock } from "../mocks/StaderOracleMock.sol";
+import { PermissionlessNodeRegistryMock } from "../mocks/PermissionlessNodeRegistryMock.sol";
 
 contract ValidatorWithdrawalVaultTest is Test {
     address private constant OPERATOR_ADDRESS = address(500);
@@ -52,6 +53,7 @@ contract ValidatorWithdrawalVaultTest is Test {
 
         mockStaderOracle(staderOracleMock);
         mockSDUtilityPool(sdUtilityPoolMock, operator);
+        mockPermissionlessNodeRegistry(vm.addr(105));
 
         ProxyAdmin proxyAdmin = new ProxyAdmin();
 
@@ -93,6 +95,7 @@ contract ValidatorWithdrawalVaultTest is Test {
         staderConfig.updateOperatorRewardsCollector(address(operatorRC));
         staderConfig.updateValidatorWithdrawalVaultImplementation(address(withdrawVaultImpl));
         staderConfig.updateStaderOracle(staderOracleMock);
+        staderConfig.updatePermissionlessNodeRegistry(vm.addr(105));
         staderConfig.grantRole(staderConfig.MANAGER(), staderManager);
         vaultFactory.grantRole(vaultFactory.NODE_REGISTRY_CONTRACT(), address(poolUtils.nodeRegistry()));
         vm.stopPrank();
@@ -384,4 +387,11 @@ contract ValidatorWithdrawalVaultTest is Test {
             abi.encode(UserData(0 ether, 4 ether, 1, 0))
         );
     }
+
+    function mockPermissionlessNodeRegistry(address _permissionlessNodeRegistry) private {
+        emit log_named_address("permissionlessNodeRegistry", _permissionlessNodeRegistry);
+        PermissionlessNodeRegistryMock nodeRegistryMock = new PermissionlessNodeRegistryMock();
+        bytes memory mockCode = address(nodeRegistryMock).code;
+        vm.etch(_permissionlessNodeRegistry, mockCode);
+    }
 }