feat: remove unnecessary steps while deposit and claim

stader-labs · Nov 26, 2024 · 57e8fc4 · 57e8fc4
1 parent c2b620a
commit 57e8fc4
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 69 deletions.
diff --git a/contracts/OperatorRewardsCollector.sol b/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;
@@ -71,7 +72,12 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
     }
 
     function claimLiquidation(address operator) public override {
-        _transferBackUtilizedSD(operator);
+        IPoolUtils poolUtils = IPoolUtils(staderConfig.getPoolUtils());
+        uint8 poolId = poolUtils.getOperatorPoolId(operator);
+        address permissionlessNodeRegistry = staderConfig.getPermissionlessNodeRegistry();
+
+        if (INodeRegistry(permissionlessNodeRegistry).POOL_ID() == poolId) _transferBackUtilizedSD(operator);
+
         _completeLiquidationIfExists(operator);
     }
 

diff --git a/contracts/PermissionedNodeRegistry.sol b/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

diff --git a/test/foundry_tests/NodeELRewardVault.t.sol b/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);
+    }
 }
diff --git a/test/foundry_tests/PermissionedNodeRegistry.t.sol b/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,

diff --git a/test/foundry_tests/ValidatorWithdrawalVault.t.sol b/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);
+    }
 }