From d8d293267bda269c1d11f879fdf470c39f4ff7c9 Mon Sep 17 00:00:00 2001 From: Alexander Movsunov Date: Wed, 18 Dec 2024 22:20:21 +0400 Subject: [PATCH] feat: val limits for modules --- remappings.txt | 1 + src/Curator.sol | 162 +++++++++++++++++++------------ test/Curator.t.sol | 88 +++++++++++++++++ test/mocks/MockStakingRouter.sol | 47 +++++++++ 4 files changed, 237 insertions(+), 61 deletions(-) create mode 100644 remappings.txt create mode 100644 test/Curator.t.sol create mode 100644 test/mocks/MockStakingRouter.sol diff --git a/remappings.txt b/remappings.txt new file mode 100644 index 0000000..feaba2d --- /dev/null +++ b/remappings.txt @@ -0,0 +1 @@ +forge-std/=lib/forge-std/src/ diff --git a/src/Curator.sol b/src/Curator.sol index 22c08fe..3d853dc 100644 --- a/src/Curator.sol +++ b/src/Curator.sol @@ -5,85 +5,125 @@ import {IStakingRouter, StakingModule} from "../interfaces/IStakingRouter.sol"; contract Curator { event Succeeded( - address sender, - address rewardAddress, - address eoa, - uint256 moduleId, - uint256 operatorId, - uint256 keysRangeStart, - uint256 keysRangeEnd + address sender, + address rewardAddress, + address eoa, + uint256 moduleId, + uint256 operatorId, + uint256 keysRangeStart, + uint256 keysRangeEnd ); event Failed( - address sender, - address rewardAddress, - address eoa, - uint256 moduleId, - uint256 operatorId, - uint256 keysRangeStart, - uint256 keysRangeEnd + address sender, + address rewardAddress, + address eoa, + uint256 moduleId, + uint256 operatorId, + uint256 keysRangeStart, + uint256 keysRangeEnd ); event Test( - uint24 id, - address stakingModuleAddress, - uint16 stakingModuleFee, - uint16 treasuryFee, - uint16 stakeShareLimit, - uint8 status, - string name, - uint64 lastDepositAt, - uint256 lastDepositBlock, - uint256 exitedValidatorsCount, - uint16 priorityExitShareThreshold, - uint64 maxDepositsPerBlock, - uint64 minDepositBlockDistance + uint24 id, + address stakingModuleAddress, + uint16 stakingModuleFee, + uint16 treasuryFee, + uint16 stakeShareLimit, + uint8 status, + string name, + uint64 lastDepositAt, + uint256 lastDepositBlock, + uint256 exitedValidatorsCount, + uint16 priorityExitShareThreshold, + uint64 maxDepositsPerBlock, + uint64 minDepositBlockDistance ); struct RegisteredOperator { - address eoa; - uint256 moduleId; - uint256 operatorId; - uint256 keysRangeStart; - uint256 keysRangeEnd; + address eoa; + uint256 moduleId; + uint256 operatorId; + uint256 keysRangeStart; + uint256 keysRangeEnd; } address immutable public stakingRouterAddress; + address public owner; - mapping (address => RegisteredOperator) public operators; + mapping(address => RegisteredOperator) public operators; + + // Лимит валидаторов для каждого Staking Module + mapping(uint256 => uint256) public maxValidatorsForModule; + + // Модификатор для проверки прав владельца + modifier onlyOwner() { + require(msg.sender == owner, "Not the owner"); + _; + } constructor(address _stakingRouterAddress) { - stakingRouterAddress = _stakingRouterAddress; + stakingRouterAddress = _stakingRouterAddress; + owner = msg.sender; } function optIn( - address rewardAddress, - address eoa, - uint256 moduleId, - uint256 operatorId, - uint256 keysRangeStart, - uint256 keysRangeEnd + address rewardAddress, + address eoa, + uint256 moduleId, + uint256 operatorId, + uint256 keysRangeStart, + uint256 keysRangeEnd ) public { - IStakingRouter router = IStakingRouter(payable(stakingRouterAddress)); - - StakingModule memory module = router.getStakingModule(moduleId); - - emit Test( - module.id, - module.stakingModuleAddress, - module.stakingModuleFee, - module.treasuryFee, - module.stakeShareLimit, - module.status, - module.name, - module.lastDepositAt, - module.lastDepositBlock, - module.exitedValidatorsCount, - module.priorityExitShareThreshold, - module.maxDepositsPerBlock, - module.minDepositBlockDistance - ); - - ///emit Succeeded(msg.sender, msg.sender, eoa, moduleId, operatorId, keysRangeStart, keysRangeEnd); + IStakingRouter router = IStakingRouter(payable(stakingRouterAddress)); + + StakingModule memory module = router.getStakingModule(moduleId); + + // Проверяем, не превышает ли количество ключей лимит для модуля + uint256 totalKeys = keysRangeEnd - keysRangeStart + 1; + require( + totalKeys <= maxValidatorsForModule[moduleId], + "Validator limit exceeded for module" + ); + + operators[rewardAddress] = RegisteredOperator({ + eoa: eoa, + moduleId: moduleId, + operatorId: operatorId, + keysRangeStart: keysRangeStart, + keysRangeEnd: keysRangeEnd + }); + + emit Test( + module.id, + module.stakingModuleAddress, + module.stakingModuleFee, + module.treasuryFee, + module.stakeShareLimit, + module.status, + module.name, + module.lastDepositAt, + module.lastDepositBlock, + module.exitedValidatorsCount, + module.priorityExitShareThreshold, + module.maxDepositsPerBlock, + module.minDepositBlockDistance + ); + + /// emit Succeeded(msg.sender, msg.sender, eoa, moduleId, operatorId, keysRangeStart, keysRangeEnd); + } + + // Установка лимита валидаторов для модуля (только для владельца) + function setMaxValidatorsForStakingModule(uint256 moduleId, uint256 maxValidators) external onlyOwner { + require(moduleId > 0, "Invalid module ID"); + require(maxValidators > 0, "Max validators must be greater than 0"); + + maxValidatorsForModule[moduleId] = maxValidators; + } + + // Изменение владельца + function transferOwnership(address newOwner) external onlyOwner { + require(newOwner != address(0), "Invalid new owner address"); + owner = newOwner; } } diff --git a/test/Curator.t.sol b/test/Curator.t.sol new file mode 100644 index 0000000..d956439 --- /dev/null +++ b/test/Curator.t.sol @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import {Test, console} from "forge-std/Test.sol"; +import "../src/Curator.sol"; +import "./mocks/MockStakingRouter.sol"; + +contract CuratorTest is Test { + Curator public curator; + + address public owner = address(0x123); + address public user = address(0x456); + MockStakingRouter public mockRouter; + + function setUp() public { + // Развёртываем мок + mockRouter = new MockStakingRouter(); + + // Устанавливаем владельца как msg.sender + vm.prank(owner); + curator = new Curator(address(mockRouter)); + + // Настраиваем моковые данные + MockStakingRouter.MockStakingModule memory module = MockStakingRouter.MockStakingModule({ + id: 1, + stakingModuleAddress: address(0xABC), + stakingModuleFee: 10, + treasuryFee: 5, + stakeShareLimit: 50, + status: 1, + name: "Mock Module", + lastDepositAt: uint64(block.timestamp), + lastDepositBlock: block.number, + exitedValidatorsCount: 0, + priorityExitShareThreshold: 10, + maxDepositsPerBlock: 100, + minDepositBlockDistance: 1 + }); + + mockRouter.setStakingModule(1, module); + } + + function testSetMaxValidatorsForStakingModule_Success() public { + // Убедимся, что владелец может установить лимит + vm.prank(owner); + curator.setMaxValidatorsForStakingModule(1, 100); + + uint256 limit = curator.maxValidatorsForModule(1); + assertEq(limit, 100, "Max validators not set correctly"); + } + + function testSetMaxValidatorsForStakingModule_Fail_NotOwner() public { + // Попробуем установить лимит не владельцем + vm.prank(user); + vm.expectRevert("Not the owner"); + curator.setMaxValidatorsForStakingModule(1, 100); + } + + function testOptIn_Success() public { + // Установим лимит валидаторов + vm.prank(owner); + curator.setMaxValidatorsForStakingModule(1, 10); + + // Попробуем вызвать optIn с корректным диапазоном ключей + vm.prank(user); + curator.optIn(address(0xabc), address(0xdef), 1, 1, 1, 10); + + // Проверяем, что оператор зарегистрирован + (address eoa, uint256 moduleId, uint256 operatorId, uint256 keysRangeStart, uint256 keysRangeEnd) = curator.operators(address(0xabc)); + + assertEq(eoa, address(0xdef), "EOA not set correctly"); + assertEq(moduleId, 1, "Module ID not set correctly"); + assertEq(operatorId, 1, "Operator ID not set correctly"); + assertEq(keysRangeStart, 1, "Keys range start not set correctly"); + assertEq(keysRangeEnd, 10, "Keys range end not set correctly"); + } + + function testOptIn_Fail_ExceedValidatorLimit() public { + // Установим лимит валидаторов + vm.prank(owner); + curator.setMaxValidatorsForStakingModule(1, 10); + + // Попробуем вызвать optIn с диапазоном ключей, превышающим лимит + vm.prank(user); + vm.expectRevert("Validator limit exceeded for module"); + curator.optIn(address(0xabc), address(0xdef), 1, 1, 1, 20); + } +} diff --git a/test/mocks/MockStakingRouter.sol b/test/mocks/MockStakingRouter.sol new file mode 100644 index 0000000..3a5e069 --- /dev/null +++ b/test/mocks/MockStakingRouter.sol @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import "../../interfaces/IStakingRouter.sol"; + +contract MockStakingRouter is IStakingRouter { + struct MockStakingModule { + uint24 id; + address stakingModuleAddress; + uint16 stakingModuleFee; + uint16 treasuryFee; + uint16 stakeShareLimit; + uint8 status; + string name; + uint64 lastDepositAt; + uint256 lastDepositBlock; + uint256 exitedValidatorsCount; + uint16 priorityExitShareThreshold; + uint64 maxDepositsPerBlock; + uint64 minDepositBlockDistance; + } + + mapping(uint256 => MockStakingModule) public stakingModules; + + function setStakingModule(uint256 moduleId, MockStakingModule memory module) external { + stakingModules[moduleId] = module; + } + + function getStakingModule(uint256 moduleId) external view override returns (StakingModule memory) { + MockStakingModule memory mockModule = stakingModules[moduleId]; + return StakingModule({ + id: mockModule.id, + stakingModuleAddress: mockModule.stakingModuleAddress, + stakingModuleFee: mockModule.stakingModuleFee, + treasuryFee: mockModule.treasuryFee, + stakeShareLimit: mockModule.stakeShareLimit, + status: mockModule.status, + name: mockModule.name, + lastDepositAt: mockModule.lastDepositAt, + lastDepositBlock: mockModule.lastDepositBlock, + exitedValidatorsCount: mockModule.exitedValidatorsCount, + priorityExitShareThreshold: mockModule.priorityExitShareThreshold, + maxDepositsPerBlock: mockModule.maxDepositsPerBlock, + minDepositBlockDistance: mockModule.minDepositBlockDistance + }); + } +}