From 896ae0c804d3df0b53db092ae05ec7232e9c1fda Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Wed, 11 Sep 2024 10:44:27 +0000 Subject: [PATCH 1/7] feat: remove unnecessary steps while deposit and claim --- .gitmodules | 3 + contracts/OperatorRewardsCollector.sol | 32 ++++++++--- contracts/PermissionedNodeRegistry.sol | 16 ------ lib/openzeppelin-foundry-upgrades | 1 + test/fork/reward-claim-optimisation.ts | 56 +++++++++++++++++++ test/foundry_tests/NodeELRewardVault.t.sol | 10 ++++ .../PermissionedNodeRegistry.t.sol | 52 ----------------- .../ValidatorWithdrawalVault.t.sol | 10 ++++ 8 files changed, 104 insertions(+), 76 deletions(-) create mode 160000 lib/openzeppelin-foundry-upgrades create mode 100644 test/fork/reward-claim-optimisation.ts diff --git a/.gitmodules b/.gitmodules index 888d42dc..a0e0ceb5 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,6 @@ [submodule "lib/forge-std"] path = lib/forge-std url = https://github.com/foundry-rs/forge-std +[submodule "lib/openzeppelin-foundry-upgrades"] + path = lib/openzeppelin-foundry-upgrades + url = https://github.com/OpenZeppelin/openzeppelin-foundry-upgrades \ No newline at end of file diff --git a/contracts/OperatorRewardsCollector.sol b/contracts/OperatorRewardsCollector.sol index 7a318eb2..9323f324 100644 --- 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; @@ -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 diff --git a/contracts/PermissionedNodeRegistry.sol b/contracts/PermissionedNodeRegistry.sol index 5b2f9b56..545e4e38 100644 --- 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/lib/openzeppelin-foundry-upgrades b/lib/openzeppelin-foundry-upgrades new file mode 160000 index 00000000..16e0ae21 --- /dev/null +++ b/lib/openzeppelin-foundry-upgrades @@ -0,0 +1 @@ +Subproject commit 16e0ae21e0e39049f619f2396fa28c57fad07368 diff --git a/test/fork/reward-claim-optimisation.ts b/test/fork/reward-claim-optimisation.ts new file mode 100644 index 00000000..d1cfb73e --- /dev/null +++ b/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); + }); +}); diff --git a/test/foundry_tests/NodeELRewardVault.t.sol b/test/foundry_tests/NodeELRewardVault.t.sol index 5eb65f30..5dd64da1 100644 --- 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 index a0da00f0..e4f79fcc 100644 --- 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 index 8b15e979..dabb95e6 100644 --- 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); + } } From 963a6a5b1bd515d21001df16c51de77faaae62bf Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Fri, 6 Dec 2024 05:09:00 +0000 Subject: [PATCH 2/7] test: increase test coverage --- .../OperatorRewardsCollector.t.sol | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/test/foundry_tests/OperatorRewardsCollector.t.sol b/test/foundry_tests/OperatorRewardsCollector.t.sol index 47c20a35..0d341a30 100644 --- a/test/foundry_tests/OperatorRewardsCollector.t.sol +++ b/test/foundry_tests/OperatorRewardsCollector.t.sol @@ -30,6 +30,8 @@ contract OperatorRewardsCollectorTest is Test { event SDWithdrawn(address indexed operator, uint256 sdAmount); event SDRepaid(address operator, uint256 repayAmount); + error InsufficientBalance(); + address staderAdmin; address staderManager; address staderTreasury; @@ -173,6 +175,102 @@ contract OperatorRewardsCollectorTest is Test { vm.stopPrank(); } + function test_Claim_PermissionedPoolOperator(uint256 amount) public { + vm.assume(amount < 100000 ether); + + operatorRewardsCollector.depositFor{ value: amount }(staderManager); + assertEq(operatorRewardsCollector.balances(staderManager), amount); + vm.mockCall( + address(staderOracle), + abi.encodeWithSelector(IStaderOracle.getSDPriceInETH.selector), + abi.encode(1e14) + ); + + vm.mockCall( + address(poolUtils), + abi.encodeWithSelector(IPoolUtils.getOperatorPoolId.selector, staderManager), + abi.encode(uint8(2)) // Assuming `POOL_ID()` for PermissionedNodeRegistry is `2` + ); + + vm.mockCall( + address(permissionlessNodeRegistryMock), + abi.encodeWithSelector(INodeRegistry.POOL_ID.selector), + abi.encode(uint8(1)) + ); + + vm.startPrank(staderManager); + operatorRewardsCollector.claim(); + + assertEq(operatorRewardsCollector.balances(staderManager), 0); + vm.stopPrank(); + } + + function test_ClaimWithAmount_PermissionlessNodeRegistry(uint256 amount, uint256 claimAmount) public { + // Assume reasonable values for deposits and claims + vm.assume(amount > 0 && amount < 100000 ether); + + operatorRewardsCollector.depositFor{ value: amount }(staderManager); + assertEq(operatorRewardsCollector.balances(staderManager), amount); + + // Mock the operator's pool ID to match the PermissionlessPool + vm.mockCall( + address(poolUtils), + abi.encodeWithSelector(IPoolUtils.getOperatorPoolId.selector, staderManager), + abi.encode(uint8(1)) + ); + + vm.mockCall( + address(permissionlessNodeRegistryMock), + abi.encodeWithSelector(INodeRegistry.POOL_ID.selector), + abi.encode(uint8(1)) + ); + uint256 withdrawableAmount = amount / 2; + vm.startPrank(staderManager); + + if (!(claimAmount > withdrawableAmount || claimAmount > amount)) { + // Case: Claiming valid amount + operatorRewardsCollector.claimWithAmount(claimAmount); + assertEq(operatorRewardsCollector.balances(staderManager), amount - claimAmount); + } + + vm.stopPrank(); + } + + function test_ClaimWithAmount_PermissionedNodeRegistry(uint256 amount, uint256 claimAmount) public { + // Assume reasonable values for deposits and claims + vm.assume(amount > 0 && amount < 100000 ether); + + operatorRewardsCollector.depositFor{ value: amount }(staderManager); + assertEq(operatorRewardsCollector.balances(staderManager), amount); + + // Mock the operator's pool ID to match the PermissionedPool + vm.mockCall( + address(poolUtils), + abi.encodeWithSelector(IPoolUtils.getOperatorPoolId.selector, staderManager), + abi.encode(uint8(2)) + ); + + vm.mockCall( + address(permissionlessNodeRegistryMock), + abi.encodeWithSelector(INodeRegistry.POOL_ID.selector), + abi.encode(uint8(1)) + ); + + vm.startPrank(staderManager); + + if (claimAmount > amount) { + // Case: Attempting to claim more than balance, expect revert + vm.expectRevert(InsufficientBalance.selector); + operatorRewardsCollector.claimWithAmount(claimAmount); + } else { + // Case: Claiming valid amount + operatorRewardsCollector.claimWithAmount(claimAmount); + assertEq(operatorRewardsCollector.balances(staderManager), amount - claimAmount); + } + + vm.stopPrank(); + } + function test_claimLiquidationZeroAmount(uint256 amount) public { vm.assume(amount < 100000 ether); 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 3/7] 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); From 3618d8ac5a42c4326ac1993e4bddf23ab20f96ba Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Fri, 6 Dec 2024 12:11:37 +0000 Subject: [PATCH 4/7] test: add test for branches to cover e2e --- .../OperatorRewardsCollector.t.sol | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/test/foundry_tests/OperatorRewardsCollector.t.sol b/test/foundry_tests/OperatorRewardsCollector.t.sol index 0d341a30..d14b5555 100644 --- a/test/foundry_tests/OperatorRewardsCollector.t.sol +++ b/test/foundry_tests/OperatorRewardsCollector.t.sol @@ -205,33 +205,43 @@ contract OperatorRewardsCollectorTest is Test { vm.stopPrank(); } - function test_ClaimWithAmount_PermissionlessNodeRegistry(uint256 amount, uint256 claimAmount) public { - // Assume reasonable values for deposits and claims - vm.assume(amount > 0 && amount < 100000 ether); + function test_ClaimWithAmount_PermissionlessNodeRegistry() public { + uint256 amount = 10 ether; // Deposit amount + // Deposit funds into the contract operatorRewardsCollector.depositFor{ value: amount }(staderManager); assertEq(operatorRewardsCollector.balances(staderManager), amount); + uint256 withdrawableAmount = operatorRewardsCollector.withdrawableInEth(staderManager); // Withdrawable amount + // Mock the operator's pool ID to match the PermissionlessPool vm.mockCall( address(poolUtils), abi.encodeWithSelector(IPoolUtils.getOperatorPoolId.selector, staderManager), - abi.encode(uint8(1)) + abi.encode(uint8(1)) // PermissionlessPool ); + // Mock Permissionless Node Registry POOL_ID vm.mockCall( address(permissionlessNodeRegistryMock), abi.encodeWithSelector(INodeRegistry.POOL_ID.selector), - abi.encode(uint8(1)) + abi.encode(uint8(1)) // PermissionlessNodeRegistry ); - uint256 withdrawableAmount = amount / 2; + vm.startPrank(staderManager); - if (!(claimAmount > withdrawableAmount || claimAmount > amount)) { - // Case: Claiming valid amount - operatorRewardsCollector.claimWithAmount(claimAmount); - assertEq(operatorRewardsCollector.balances(staderManager), amount - claimAmount); - } + // Case 1: Claiming more than `withdrawableInEth` + vm.expectRevert(InsufficientBalance.selector); + operatorRewardsCollector.claimWithAmount(withdrawableAmount + 1 ether); + + // Case 2: Claiming more than `balances[msg.sender]` + vm.expectRevert(InsufficientBalance.selector); + operatorRewardsCollector.claimWithAmount(amount + 1 ether); + + // Case 3: Valid claim within limits + uint256 validClaim = 3 ether; + operatorRewardsCollector.claimWithAmount(validClaim); + assertEq(operatorRewardsCollector.balances(staderManager), amount - validClaim); vm.stopPrank(); } From acacb3541b9f5b48fae33cdd477099005f86db8b Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Mon, 9 Dec 2024 05:25:36 +0000 Subject: [PATCH 5/7] feat: add limit on keys per operator --- contracts/PermissionlessNodeRegistry.sol | 15 +++++++++++++++ .../interfaces/IPermissionlessNodeRegistry.sol | 2 ++ .../PermissionlessNodeRegistry.t.sol | 2 ++ 3 files changed, 19 insertions(+) diff --git a/contracts/PermissionlessNodeRegistry.sol b/contracts/PermissionlessNodeRegistry.sol index e1605eea..5ab3a38f 100644 --- a/contracts/PermissionlessNodeRegistry.sol +++ b/contracts/PermissionlessNodeRegistry.sol @@ -59,6 +59,7 @@ contract PermissionlessNodeRegistry is //mapping of operator address with nodeELReward vault address mapping(uint256 => address) public override nodeELRewardVaultByOperatorId; mapping(uint256 => address) public proposedRewardAddressByOperatorId; + uint256 public maxKeyPerOperator; /// @custom:oz-upgrades-unsafe-allow constructor constructor() { @@ -150,6 +151,10 @@ contract PermissionlessNodeRegistry is ) public payable override nonReentrant whenNotPaused { uint256 operatorId = onlyActiveOperator(msg.sender); uint256 keyCount = _pubkey.length; + uint256 totalKeys = getOperatorTotalKeys(operatorId); + if (totalKeys + keyCount > maxKeyPerOperator) { + revert MaxKeyLimitExceed(); + } if (keyCount != _preDepositSignature.length || keyCount != _depositSignature.length) { revert MisMatchingInputKeysSize(); } @@ -444,6 +449,16 @@ contract PermissionlessNodeRegistry is emit TransferredCollateralToPool(_amount); } + /** + * @notice update the max validator per operator value + * @dev only `MANAGER` role can update + */ + function updateMaxKeyPerOperator(uint256 _maxKeyPerOperator) external { + UtilLib.onlyManagerRole(msg.sender, staderConfig); + maxKeyPerOperator = _maxKeyPerOperator; + emit UpdateMaxKeyPerOperator(_maxKeyPerOperator); + } + /** * @param _nodeOperator @notice operator total non terminal keys within a specified validator list * @param _startIndex start index in validator queue to start with diff --git a/contracts/interfaces/IPermissionlessNodeRegistry.sol b/contracts/interfaces/IPermissionlessNodeRegistry.sol index 89b9e3ea..08bf989d 100644 --- a/contracts/interfaces/IPermissionlessNodeRegistry.sol +++ b/contracts/interfaces/IPermissionlessNodeRegistry.sol @@ -8,6 +8,7 @@ interface IPermissionlessNodeRegistry { error InSufficientBalance(); error CooldownNotComplete(); error NoChangeInState(); + error MaxKeyLimitExceed(); // Events event OnboardedOperator( @@ -21,6 +22,7 @@ interface IPermissionlessNodeRegistry { event UpdatedSocializingPoolState(uint256 operatorId, bool optedForSocializingPool, uint256 block); event TransferredCollateralToPool(uint256 amount); event ValidatorAddedViaReferral(uint256 amount, string referralId); + event UpdateMaxKeyPerOperator(uint256 maxKeyPerOperator); //Getters diff --git a/test/foundry_tests/PermissionlessNodeRegistry.t.sol b/test/foundry_tests/PermissionlessNodeRegistry.t.sol index 3ba936e1..f4e23db8 100644 --- a/test/foundry_tests/PermissionlessNodeRegistry.t.sol +++ b/test/foundry_tests/PermissionlessNodeRegistry.t.sol @@ -97,6 +97,8 @@ contract PermissionlessNodeRegistryTest is Test { staderConfig.grantRole(staderConfig.OPERATOR(), operator); vaultFactory.grantRole(vaultFactory.NODE_REGISTRY_CONTRACT(), address(nodeRegistry)); vm.stopPrank(); + vm.prank(staderManager); + nodeRegistry.updateMaxKeyPerOperator(500); } function test_JustToIncreaseCoverage() public { From d494f2ca4b8bffd0ec6d45044bb033ce12408d6a Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Tue, 10 Dec 2024 04:39:05 +0000 Subject: [PATCH 6/7] test: add test for max key limit --- .../PermissionlessNodeRegistry.t.sol | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/test/foundry_tests/PermissionlessNodeRegistry.t.sol b/test/foundry_tests/PermissionlessNodeRegistry.t.sol index f4e23db8..b02a4819 100644 --- a/test/foundry_tests/PermissionlessNodeRegistry.t.sol +++ b/test/foundry_tests/PermissionlessNodeRegistry.t.sol @@ -256,6 +256,21 @@ contract PermissionlessNodeRegistryTest is Test { assertEq(nodeRegistry.isExistingPubkey(pubkeys[0]), true); } + function test_addValidatorKeysWithLLimitExceed() public { + ( + bytes[] memory pubkeys, + bytes[] memory preDepositSignature, + bytes[] memory depositSignature + ) = getValidatorKeys(); + vm.prank(staderManager); + nodeRegistry.updateMaxKeyPerOperator(2); // Limit is 2 and trying to add 3 keys + startHoax(address(this)); + nodeRegistry.onboardNodeOperator(true, "testOP", payable(address(this))); + vm.expectRevert(IPermissionlessNodeRegistry.MaxKeyLimitExceed.selector); + nodeRegistry.addValidatorKeys{ value: 12 ether }(pubkeys, preDepositSignature, depositSignature); + vm.stopPrank(); + } + function testAddValidatorKeysNotEnoughSDCollateral() public { ( bytes[] memory pubkeys, From fabfd942fd86fe01a479bb5792834e3b1d2f7f0c Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Thu, 12 Dec 2024 05:15:06 +0000 Subject: [PATCH 7/7] fix: typo --- contracts/PermissionlessNodeRegistry.sol | 10 +++++----- contracts/interfaces/IPermissionlessNodeRegistry.sol | 2 +- test/foundry_tests/PermissionlessNodeRegistry.t.sol | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/contracts/PermissionlessNodeRegistry.sol b/contracts/PermissionlessNodeRegistry.sol index 5ab3a38f..2f7b0ba5 100644 --- a/contracts/PermissionlessNodeRegistry.sol +++ b/contracts/PermissionlessNodeRegistry.sol @@ -59,7 +59,7 @@ contract PermissionlessNodeRegistry is //mapping of operator address with nodeELReward vault address mapping(uint256 => address) public override nodeELRewardVaultByOperatorId; mapping(uint256 => address) public proposedRewardAddressByOperatorId; - uint256 public maxKeyPerOperator; + uint256 public maxKeysPerOperator; /// @custom:oz-upgrades-unsafe-allow constructor constructor() { @@ -152,7 +152,7 @@ contract PermissionlessNodeRegistry is uint256 operatorId = onlyActiveOperator(msg.sender); uint256 keyCount = _pubkey.length; uint256 totalKeys = getOperatorTotalKeys(operatorId); - if (totalKeys + keyCount > maxKeyPerOperator) { + if (totalKeys + keyCount > maxKeysPerOperator) { revert MaxKeyLimitExceed(); } if (keyCount != _preDepositSignature.length || keyCount != _depositSignature.length) { @@ -453,10 +453,10 @@ contract PermissionlessNodeRegistry is * @notice update the max validator per operator value * @dev only `MANAGER` role can update */ - function updateMaxKeyPerOperator(uint256 _maxKeyPerOperator) external { + function updateMaxKeysPerOperator(uint256 _maxKeysPerOperator) external { UtilLib.onlyManagerRole(msg.sender, staderConfig); - maxKeyPerOperator = _maxKeyPerOperator; - emit UpdateMaxKeyPerOperator(_maxKeyPerOperator); + maxKeysPerOperator = _maxKeysPerOperator; + emit UpdateMaxKeysPerOperator(_maxKeysPerOperator); } /** diff --git a/contracts/interfaces/IPermissionlessNodeRegistry.sol b/contracts/interfaces/IPermissionlessNodeRegistry.sol index 08bf989d..722acb14 100644 --- a/contracts/interfaces/IPermissionlessNodeRegistry.sol +++ b/contracts/interfaces/IPermissionlessNodeRegistry.sol @@ -22,7 +22,7 @@ interface IPermissionlessNodeRegistry { event UpdatedSocializingPoolState(uint256 operatorId, bool optedForSocializingPool, uint256 block); event TransferredCollateralToPool(uint256 amount); event ValidatorAddedViaReferral(uint256 amount, string referralId); - event UpdateMaxKeyPerOperator(uint256 maxKeyPerOperator); + event UpdateMaxKeysPerOperator(uint256 maxKeysPerOperator); //Getters diff --git a/test/foundry_tests/PermissionlessNodeRegistry.t.sol b/test/foundry_tests/PermissionlessNodeRegistry.t.sol index b02a4819..dcc4bbd8 100644 --- a/test/foundry_tests/PermissionlessNodeRegistry.t.sol +++ b/test/foundry_tests/PermissionlessNodeRegistry.t.sol @@ -98,7 +98,7 @@ contract PermissionlessNodeRegistryTest is Test { vaultFactory.grantRole(vaultFactory.NODE_REGISTRY_CONTRACT(), address(nodeRegistry)); vm.stopPrank(); vm.prank(staderManager); - nodeRegistry.updateMaxKeyPerOperator(500); + nodeRegistry.updateMaxKeysPerOperator(500); } function test_JustToIncreaseCoverage() public { @@ -263,7 +263,7 @@ contract PermissionlessNodeRegistryTest is Test { bytes[] memory depositSignature ) = getValidatorKeys(); vm.prank(staderManager); - nodeRegistry.updateMaxKeyPerOperator(2); // Limit is 2 and trying to add 3 keys + nodeRegistry.updateMaxKeysPerOperator(2); // Limit is 2 and trying to add 3 keys startHoax(address(this)); nodeRegistry.onboardNodeOperator(true, "testOP", payable(address(this))); vm.expectRevert(IPermissionlessNodeRegistry.MaxKeyLimitExceed.selector);