From a24f978a60f76706fb39f101a8099b74e6bbd577 Mon Sep 17 00:00:00 2001 From: Shawn <44221603+shaspitz@users.noreply.github.com> Date: Wed, 24 Apr 2024 00:40:11 -0700 Subject: [PATCH] feat: introduce unstakingBalances mapping to registry --- contracts/contracts/ValidatorRegistry.sol | 45 +++++++++-------- contracts/test/ValidatorRegistryTest.sol | 60 +++++++++++++++++++---- 2 files changed, 75 insertions(+), 30 deletions(-) diff --git a/contracts/contracts/ValidatorRegistry.sol b/contracts/contracts/ValidatorRegistry.sol index 18d72e223..c0b56e375 100644 --- a/contracts/contracts/ValidatorRegistry.sol +++ b/contracts/contracts/ValidatorRegistry.sol @@ -42,6 +42,9 @@ contract ValidatorRegistry is OwnableUpgradeable, ReentrancyGuardUpgradeable { /// @dev Mapping of bls public key to block number of unstake initiation block. mapping(bytes => uint256) public unstakeBlockNums; + /// @dev Mapping of bls public key to balance of currently unstaking ether. + mapping(bytes => uint256) public unstakingBalances; + event Staked(address indexed txOriginator, bytes valBLSPubKey, uint256 amount); event Unstaked(address indexed txOriginator, bytes valBLSPubKey, uint256 amount); event StakeWithdrawn(address indexed txOriginator, bytes valBLSPubKey, uint256 amount); @@ -67,9 +70,10 @@ contract ValidatorRegistry is OwnableUpgradeable, ReentrancyGuardUpgradeable { } } - function unstake(bytes[] calldata blsPubKeys) external { + function unstake(bytes[] calldata blsPubKeys) external nonReentrant { for (uint256 i = 0; i < blsPubKeys.length; i++) { _validateBLSPubKey(blsPubKeys[i]); + require(unstakeBlockNums[blsPubKeys[i]] == 0, "Unstake already initiated for validator"); bool exists; uint256 balance; @@ -78,9 +82,13 @@ contract ValidatorRegistry is OwnableUpgradeable, ReentrancyGuardUpgradeable { require(balance >= minStake, "No staked balance over min stake"); require(stakeOriginators[blsPubKeys[i]] == msg.sender, "Not authorized to unstake validator. Must be stake originator"); - require(unstakeBlockNums[blsPubKeys[i]] == 0, "Unstake already initiated for validator"); + + bool removed = stakedBalances.remove(blsPubKeys[i]); + require(removed, "Failed to remove staked balance"); unstakeBlockNums[blsPubKeys[i]] = block.number; + unstakingBalances[blsPubKeys[i]] = balance; + emit Unstaked(msg.sender, blsPubKeys[i], balance); } } @@ -89,23 +97,19 @@ contract ValidatorRegistry is OwnableUpgradeable, ReentrancyGuardUpgradeable { for (uint256 i = 0; i < blsPubKeys.length; i++) { _validateBLSPubKey(blsPubKeys[i]); - bool exists; - uint256 balance; - (exists, balance) = stakedBalances.tryGet(blsPubKeys[i]); - require(exists, "Validator not staked"); - require(balance >= minStake, "No staked balance over min stake"); + require(unstakeBlockNums[blsPubKeys[i]] > 0, "Unstake must be initiated before withdrawal"); + require(unstakingBalances[blsPubKeys[i]] >= minStake, "No unstaking balance over min stake"); require(stakeOriginators[blsPubKeys[i]] == msg.sender , "Not authorized to withdraw stake. Must be stake originator"); - require(unstakeBlockNums[blsPubKeys[i]] > 0, "Unstake must be initiated before withdrawal"); require(block.number >= unstakeBlockNums[blsPubKeys[i]] + unstakePeriodBlocks, "withdrawal not allowed yet. Blocks requirement not met."); - bool removed = stakedBalances.remove(blsPubKeys[i]); - require(removed, "Failed to remove staked balance"); - payable(msg.sender).transfer(balance); - stakeOriginators[blsPubKeys[i]] = address(0); unstakeBlockNums[blsPubKeys[i]] = 0; + uint256 balance = unstakingBalances[blsPubKeys[i]]; + unstakingBalances[blsPubKeys[i]] = 0; + payable(msg.sender).transfer(balance); + emit StakeWithdrawn(msg.sender, blsPubKeys[i], balance); } } @@ -114,20 +118,21 @@ contract ValidatorRegistry is OwnableUpgradeable, ReentrancyGuardUpgradeable { require(valBLSPubKey.length == 48, "Invalid BLS public key length. Must be 48 bytes"); } - function isStaked(bytes calldata valBLSPubKey) external view returns (bool) { + function getStakedAmount(bytes calldata valBLSPubKey) external view returns (uint256) { bool exists; uint256 balance; (exists, balance) = stakedBalances.tryGet(valBLSPubKey); - bool adequateStake = exists && balance >= minStake; - bool notUnstaking = unstakeBlockNums[valBLSPubKey] == 0; - return adequateStake && notUnstaking; + return exists ? balance : 0; } - function getStakedAmount(bytes calldata valBLSPubKey) external view returns (uint256) { + function isStaked(bytes calldata valBLSPubKey) external view returns (bool) { bool exists; - uint256 balance; - (exists, balance) = stakedBalances.tryGet(valBLSPubKey); - return exists ? balance : 0; + (exists,) = stakedBalances.tryGet(valBLSPubKey); + return exists; + } + + function getUnstakingAmount(bytes calldata valBLSPubKey) external view returns (uint256) { + return unstakingBalances[valBLSPubKey]; } function getNumberOfStakedValidators() external view returns (uint256) { diff --git a/contracts/test/ValidatorRegistryTest.sol b/contracts/test/ValidatorRegistryTest.sol index 2ee0ea461..89c1f4ea0 100644 --- a/contracts/test/ValidatorRegistryTest.sol +++ b/contracts/test/ValidatorRegistryTest.sol @@ -72,15 +72,17 @@ contract ValidatorRegistryTest is Test { validators[1] = user2BLSKey; uint256 totalAmount = 2 ether; - vm.deal(address(this), 3 ether); - assertEq(address(this).balance, 3 ether); + vm.deal(user1, 3 ether); + assertEq(user1.balance, 3 ether); + vm.startPrank(user1); vm.expectEmit(true, true, true, true); - emit Staked(address(this), validators[0], 1 ether); - emit Staked(address(this), validators[1], 1 ether); + emit Staked(user1, user1BLSKey, 1 ether); + emit Staked(user1, user2BLSKey, 1 ether); validatorRegistry.stake{value: totalAmount}(validators); + vm.stopPrank(); - assertEq(address(this).balance, 1 ether); + assertEq(user1.balance, 1 ether); assertEq(validatorRegistry.getStakedAmount(user1BLSKey), 1 ether); assertEq(validatorRegistry.getStakedAmount(user2BLSKey), 1 ether); assertTrue(validatorRegistry.isStaked(user1BLSKey)); @@ -88,11 +90,11 @@ contract ValidatorRegistryTest is Test { } function testUnstakeInsufficientFunds() public { - vm.startPrank(user2); bytes[] memory validators = new bytes[](1); validators[0] = user2BLSKey; - assertEq(validatorRegistry.getStakedAmount(user2BLSKey), 0); + + vm.startPrank(user2); vm.expectRevert("Validator not staked"); validatorRegistry.unstake(validators); vm.stopPrank(); @@ -159,7 +161,8 @@ contract ValidatorRegistryTest is Test { assertFalse(validatorRegistry.isStaked(user1BLSKey)); assertTrue(validatorRegistry.unstakeBlockNums(user1BLSKey) > 0); - assertEq(validatorRegistry.getStakedAmount(user1BLSKey), MIN_STAKE); + assertEq(validatorRegistry.getStakedAmount(user1BLSKey), 0); + assertEq(validatorRegistry.getUnstakingAmount(user1BLSKey), MIN_STAKE); vm.startPrank(user1); vm.expectRevert("validator cannot be staked with in-progress unstake process"); @@ -208,9 +211,9 @@ contract ValidatorRegistryTest is Test { validatorRegistry.unstake(validators); vm.stopPrank(); - // still has staked balance until withdrawal, but not considered "staked" - assertEq(validatorRegistry.getStakedAmount(user1BLSKey), MIN_STAKE); + assertEq(validatorRegistry.getStakedAmount(user1BLSKey), 0); assertFalse(validatorRegistry.isStaked(user1BLSKey)); + assertEq(validatorRegistry.getUnstakingAmount(user1BLSKey), MIN_STAKE); assertEq(address(user1).balance, 8 ether); uint256 blockWaitPeriod = 11; @@ -228,6 +231,7 @@ contract ValidatorRegistryTest is Test { assertEq(validatorRegistry.getStakedAmount(user1BLSKey), 0, "User1s staked balance should be 0 after withdrawal"); assertTrue(validatorRegistry.stakeOriginators(user1BLSKey) == address(0), "User1s stake originator should be reset after withdrawal"); assertTrue(validatorRegistry.unstakeBlockNums(user1BLSKey) == 0, "User1s unstake block number should be reset after withdrawal"); + assertTrue(validatorRegistry.getUnstakingAmount(user1BLSKey) == 0, "User1s unstaking balance should be reset after withdrawal"); } function testGetStakedValidators() public { @@ -268,6 +272,42 @@ contract ValidatorRegistryTest is Test { } } + function testGetStakedValidatorsWithUnstakingInProgress() public { + testMultiStake(); + + uint256 numStakedValidators = validatorRegistry.getNumberOfStakedValidators(); + assertEq(numStakedValidators, 2); + bytes[] memory validators = validatorRegistry.getStakedValidators(0, numStakedValidators); + assertEq(validators.length, 2); + + bytes[] memory keys = new bytes[](1); + keys[0] = user1BLSKey; + + vm.startPrank(user1); + vm.expectEmit(true, true, true, true); + emit Unstaked(user1, user1BLSKey, MIN_STAKE); + validatorRegistry.unstake(keys); + assertTrue(validatorRegistry.unstakeBlockNums(user1BLSKey) > 0); + vm.stopPrank(); + + numStakedValidators = validatorRegistry.getNumberOfStakedValidators(); + assertEq(numStakedValidators, 1); + validators = validatorRegistry.getStakedValidators(0, numStakedValidators); + assertEq(validators.length, 1); + + vm.roll(500); + + vm.startPrank(user1); + vm.expectEmit(true, true, true, true); + emit StakeWithdrawn(user1, user1BLSKey, MIN_STAKE); + keys = new bytes[](1); + keys[0] = user1BLSKey; + validatorRegistry.withdraw(keys); + vm.stopPrank(); + numStakedValidators = validatorRegistry.getNumberOfStakedValidators(); + assertEq(numStakedValidators, 1); + } + // To sanity check that relevant state for an account is reset s.t. they could stake again in future function testStakingCycle() public { testUnstakeWaitThenWithdraw();