Skip to content

Commit

Permalink
feat: introduce unstakingBalances mapping to registry
Browse files Browse the repository at this point in the history
  • Loading branch information
@shaspitz
shaspitz committed Apr 24, 2024
1 parent a7a3c91 commit a24f978
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 30 deletions.
45 changes: 25 additions & 20 deletions contracts/contracts/ValidatorRegistry.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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;
Expand All @@ -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);
}
}
Expand All @@ -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);
}
}
Expand All @@ -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) {
Expand Down
60 changes: 50 additions & 10 deletions contracts/test/ValidatorRegistryTest.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,27 +72,29 @@ 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));
assertTrue(validatorRegistry.isStaked(user1BLSKey));
}

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();
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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;
Expand All @@ -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 {
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit a24f978

Please sign in to comment.