Skip to content

Commit

Permalink
2370 staketable update validator identification keys (#2373)
Browse files Browse the repository at this point in the history 
* add stake table tests

* remove stake types

* verify token allowance, balance and reprioritize verification order on registration

* set the fixed stake amount, added related tests, updated data types

* add more verification checks to the withdraw function

* updated errror types

* added TODO statements in comments to be explicit about outdated functions that need to be updated to the new spec

* a validator/node is identified by the msg.sender, blskey and schnorrkey. todo: add tests

* merged from the fixed stake branch and modified test to include the inclusion of the blsKey

* not required to get the blsSig for a withdrawal since we have the eth account the validator registered with

* add the ability to update consensus keys

* add the ability to update consensus keys

* remove vscode settings

* add test or happy path of updating consensus keys

* added unhappy path tests on based on no key changes and changed the blsVK comparison to field by field comparison so that it is more efficient and more obvious what is going on

* added comment to newly added function on StakeTable and added tests that have been commented out for now

* change the lookup node and lookup stake functions to use their ethereum address, update related implementations and add tests

* updated updateConsensusKeys function, enhance test coverage

* added todo

* updated test to use the new seed so that a new schnorr key could be generated

* Update contracts/src/StakeTable.sol

Co-authored-by: Philippe Camacho <philippe@espressosys.com>

* check for invalid bls and schnorr keys on register

* update test comments so that they are clearer

* updateConsensusKeys reverts when the keys the same as the old ones, no change occurs if both keys are zero and only one key is changed if the other is a zero key

* updating consensus keys revert if any VK is a zero point VK

* clarifies that the bls sig proves

---------

Co-authored-by: Philippe Camacho <philippe@espressosys.com>
@alysiahuggins @philippecamacho
alysiahuggins and philippecamacho authored Jan 15, 2025
1 parent 6e1dc58 commit 97308ad
Showing 4 changed files with 639 additions and 100 deletions.
205 changes: 153 additions & 52 deletions contracts/src/StakeTable.sol
View file
Original file line number Diff line number Diff line change
@@ -7,6 +7,8 @@ import { AbstractStakeTable } from "./interfaces/AbstractStakeTable.sol";
import { LightClient } from "../src/LightClient.sol";
import { EdOnBN254 } from "./libraries/EdOnBn254.sol";

using EdOnBN254 for EdOnBN254.EdOnBN254Point;

/// @title Implementation of the Stake Table interface
contract StakeTable is AbstractStakeTable {
/// Error to notify restaking is not implemented yet.
@@ -53,8 +55,17 @@ contract StakeTable is AbstractStakeTable {
// Error raised when the staker does not register with the correct stakeAmount
error InsufficientStakeAmount(uint256);

// Error raised when the staker does not provide a new schnorrVK
error InvalidSchnorrVK();

// Error raised when the staker does not provide a new blsVK
error InvalidBlsVK();

// Error raised when zero point keys are provided
error NoKeyChange();

/// Mapping from a hash of a BLS key to a node struct defined in the abstract contract.
mapping(bytes32 keyHash => Node node) public nodes;
mapping(address account => Node node) public nodes;

/// Total stake locked;
uint256 public totalStake;
@@ -103,33 +114,45 @@ contract StakeTable is AbstractStakeTable {
return keccak256(abi.encode(blsVK.x0, blsVK.x1, blsVK.y0, blsVK.y1));
}

/// @dev Compares two BLS keys for equality
/// @param a First BLS key
/// @param b Second BLS key
/// @return True if the keys are equal, false otherwise
function _isEqualBlsKey(BN254.G2Point memory a, BN254.G2Point memory b)
public
pure
returns (bool)
{
return BN254.BaseField.unwrap(a.x0) == BN254.BaseField.unwrap(b.x0)
&& BN254.BaseField.unwrap(a.x1) == BN254.BaseField.unwrap(b.x1)
&& BN254.BaseField.unwrap(a.y0) == BN254.BaseField.unwrap(b.y0)
&& BN254.BaseField.unwrap(a.y1) == BN254.BaseField.unwrap(b.y1);
}

/// TODO handle this logic more appropriately when epochs are re-introduced
/// @dev Fetches the current epoch from the light client contract.
/// @return current epoch (computed from the current block)
function currentEpoch() public pure returns (uint64) {
return 0;
}

/// @notice Look up the balance of `blsVK`
/// @param blsVK BLS public key controlled by the user.
/// @notice Look up the balance of `account`
/// @param account account controlled by the user.
/// @return Current balance owned by the user.
/// TODO modify this according to the current spec
function lookupStake(BN254.G2Point memory blsVK) external view override returns (uint256) {
Node memory node = this.lookupNode(blsVK);
function lookupStake(address account) external view override returns (uint256) {
Node memory node = this.lookupNode(account);
return node.balance;
}

/// @notice Look up the full `Node` state associated with `blsVK`
/// @dev The lookup is achieved by hashing first the four field elements of blsVK using
/// keccak256.
/// @return Node indexed by blsVK
/// TODO modify this according to the current spec
function lookupNode(BN254.G2Point memory blsVK) external view override returns (Node memory) {
return nodes[_hashBlsKey(blsVK)];
/// @notice Look up the full `Node` state associated with `account`
/// @return Node indexed by account
function lookupNode(address account) external view override returns (Node memory) {
return nodes[account];
}

/// @notice Get the next available epoch and queue size in that epoch
/// TODO modify this according to the current spec
/// TODO modify this according to the current spec
function nextRegistrationEpoch() external view override returns (uint64, uint64) {
uint64 epoch;
uint64 queueSize;
@@ -152,19 +175,22 @@ contract StakeTable is AbstractStakeTable {
// @param queueSize current size of the registration queue (after insertion of new element in
// the queue)
/// TODO modify this according to the current spec
/// TODO modify this according to the current spec
function appendRegistrationQueue(uint64 epoch, uint64 queueSize) private {
firstAvailableRegistrationEpoch = epoch;
_numPendingRegistrations = queueSize + 1;
}

/// @notice Get the number of pending registration requests in the waiting queue
/// TODO modify this according to the current spec
/// TODO modify this according to the current spec
function numPendingRegistrations() external view override returns (uint64) {
return _numPendingRegistrations;
}

/// @notice Get the next available epoch for exit and queue size in that epoch
/// TODO modify this according to the current spec
/// TODO modify this according to the current spec
function nextExitEpoch() external view override returns (uint64, uint64) {
uint64 epoch;
uint64 queueSize;
@@ -186,13 +212,15 @@ contract StakeTable is AbstractStakeTable {
// @param epoch next available exit epoch
// @param queueSize current size of the exit queue (after insertion of new element in the queue)
/// TODO modify this according to the current spec
/// TODO modify this according to the current spec
function appendExitQueue(uint64 epoch, uint64 queueSize) private {
firstAvailableExitEpoch = epoch;
_numPendingExits = queueSize + 1;
}

/// @notice Get the number of pending exit requests in the waiting queue
/// TODO modify this according to the current spec
/// TODO modify this according to the current spec
function numPendingExits() external view override returns (uint64) {
return _numPendingExits;
}
@@ -212,6 +240,7 @@ contract StakeTable is AbstractStakeTable {
/// @param node node which is assigned an exit escrow period.
/// @return Number of epochs post exit after which funds can be withdrawn.
/// TODO modify this according to the current spec
/// TODO modify this according to the current spec
function exitEscrowPeriod(Node memory node) public pure returns (uint64) {
if (node.balance > 100) {
return 10;
@@ -254,8 +283,7 @@ contract StakeTable is AbstractStakeTable {
revert InsufficientStakeAmount(amount);
}

bytes32 key = _hashBlsKey(blsVK);
Node memory node = nodes[key];
Node memory node = nodes[msg.sender];

// Verify that the node is not already registered.
if (node.account != address(0x0)) {
@@ -274,10 +302,30 @@ contract StakeTable is AbstractStakeTable {
revert InsufficientBalance(balance);
}

// Verify that blsVK is not the zero point
if (
_isEqualBlsKey(
blsVK,
BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
)
)
) {
revert InvalidBlsVK();
}

// Verify that the validator can sign for that blsVK
bytes memory message = abi.encode(msg.sender);
BLSSig.verifyBlsSig(message, blsSig, blsVK);

// Verify that the schnorrVK is non-zero
if (schnorrVK.isEqual(EdOnBN254.EdOnBN254Point(0, 0))) {
revert InvalidSchnorrVK();
}

// Find the earliest epoch at which this node can register. Usually, this will be
// currentEpoch() + 1 (the start of the next full epoch), but in periods of high churn the
// queue may fill up and it may be later. If the queue is so full that the wait time exceeds
@@ -299,28 +347,28 @@ contract StakeTable is AbstractStakeTable {
// Create an entry for the node.
node.account = msg.sender;
node.balance = fixedStakeAmount;
node.blsVK = blsVK;
node.schnorrVK = schnorrVK;
node.registerEpoch = registerEpoch;

nodes[key] = node;
nodes[msg.sender] = node;

emit Registered(key, registerEpoch, fixedStakeAmount);
emit Registered(msg.sender, registerEpoch, fixedStakeAmount);
}

/// @notice Deposit more stakes to registered keys
/// @dev TODO this implementation will be revisited later. See
/// https://github.com/EspressoSystems/espresso-sequencer/issues/806
/// @dev TODO modify this according to the current spec
/// @param blsVK The BLS verification key
/// @param amount The amount to deposit
/// @return (newBalance, effectiveEpoch) the new balance effective at a future epoch
function deposit(BN254.G2Point memory blsVK, uint256 amount)
external
override
returns (uint256, uint64)
{
bytes32 key = _hashBlsKey(blsVK);
Node memory node = nodes[key];
function deposit(uint256 amount) external override returns (uint256, uint64) {
Node memory node = nodes[msg.sender];

// if the node is not registered, revert
if (node.account == address(0)) {
revert NodeNotRegistered();
}

// The deposit must come from the node's registered account.
if (node.account != msg.sender) {
@@ -338,23 +386,27 @@ contract StakeTable is AbstractStakeTable {
revert ExitRequestInProgress();
}

nodes[key].balance += amount;
nodes[msg.sender].balance += amount;
SafeTransferLib.safeTransferFrom(ERC20(tokenAddress), msg.sender, address(this), amount);

emit Deposit(_hashBlsKey(blsVK), uint256(amount));
emit Deposit(msg.sender, uint256(amount));

uint64 effectiveEpoch = _currentEpoch + 1;

return (nodes[key].balance, effectiveEpoch);
return (nodes[msg.sender].balance, effectiveEpoch);
}

/// @notice Request to exit from the stake table, not immediately withdrawable!
///
/// @dev TODO modify this according to the current spec
/// @param blsVK The BLS verification key to exit
function requestExit(BN254.G2Point memory blsVK) external override {
bytes32 key = _hashBlsKey(blsVK);
Node memory node = nodes[key];
function requestExit() external override {
Node memory node = nodes[msg.sender];

// TODO test this behaviour when re-implementing the logic for handling epochs
// if the node is not registered, revert
if (node.account == address(0)) {
revert NodeNotRegistered();
}

// The exit request must come from the node's withdrawal account.
if (node.account != msg.sender) {
@@ -374,60 +426,109 @@ contract StakeTable is AbstractStakeTable {

// Prepare the node to exit.
(uint64 exitEpoch, uint64 queueSize) = this.nextExitEpoch();
nodes[key].exitEpoch = exitEpoch;
nodes[msg.sender].exitEpoch = exitEpoch;

appendExitQueue(exitEpoch, queueSize);

emit Exit(key, exitEpoch);
emit Exit(msg.sender, exitEpoch);
}

/// @notice Withdraw from the staking pool. Transfers occur! Only successfully exited keys can
/// withdraw past their `exitEpoch`.
///
/// @param blsVK The BLS verification key to withdraw
/// @param blsSig The BLS signature that authenticates the ethereum account this function is
/// called from the caller
/// @return The total amount withdrawn, equal to `Node.balance` associated with `blsVK`
/// TODO: This function should be tested
/// TODO modify this according to the current spec

function withdrawFunds(BN254.G2Point memory blsVK, BN254.G1Point memory blsSig)
external
override
returns (uint256)
{
bytes32 key = _hashBlsKey(blsVK);
Node memory node = nodes[key];
function withdrawFunds() external override returns (uint256) {
Node memory node = nodes[msg.sender];

// Verify that the node is already registered.
if (node.account == address(0)) {
revert NodeNotRegistered();
}

// The exit request must come from the node's withdrawal account.
if (node.account != msg.sender) {
revert Unauthenticated();
}

// Verify that the balance is greater than zero
uint256 balance = node.balance;
if (balance == 0) {
revert InsufficientStakeBalance(0);
}

// Verify that the validator can sign for that blsVK
bytes memory message = abi.encode(msg.sender);
BLSSig.verifyBlsSig(message, blsSig, blsVK);

// Verify that the exit escrow period is over.
if (currentEpoch() < node.exitEpoch + exitEscrowPeriod(node)) {
revert PrematureWithdrawal();
}

// Delete the node from the stake table.
delete nodes[key];
delete nodes[msg.sender];

// Transfer the balance to the node's account.
SafeTransferLib.safeTransfer(ERC20(tokenAddress), node.account, balance);

return balance;
}

/// @notice Update the consensus keys for a validator
/// @dev This function is used to update the consensus keys for a validator
/// @dev This function can only be called by the validator itself when it's not in the exit
/// queue
/// @dev The validator will need to give up either its old BLS key and/or old Schnorr key
/// @dev The validator will need to provide a BLS signature to prove that the account owns the
/// new BLS key
/// @param newBlsVK The new BLS verification key
/// @param newSchnorrVK The new Schnorr verification key
/// @param newBlsSig The BLS signature that the account owns the new BLS key
function updateConsensusKeys(
BN254.G2Point memory newBlsVK,
EdOnBN254.EdOnBN254Point memory newSchnorrVK,
BN254.G1Point memory newBlsSig
) external override {
Node memory node = nodes[msg.sender];

// Verify that the node is already registered.
if (node.account == address(0)) revert NodeNotRegistered();

// Verify that the node is not in the exit queue
if (node.exitEpoch != 0) revert ExitRequestInProgress();

// Verify that the keys are not the same as the old ones
if (_isEqualBlsKey(newBlsVK, node.blsVK) && newSchnorrVK.isEqual(node.schnorrVK)) {
revert NoKeyChange();
}

// Zero-point constants for verification
BN254.G2Point memory zeroBlsKey = BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
);
EdOnBN254.EdOnBN254Point memory zeroSchnorrKey = EdOnBN254.EdOnBN254Point(0, 0);

if (_isEqualBlsKey(newBlsVK, zeroBlsKey)) revert InvalidBlsVK();

if (newSchnorrVK.isEqual(zeroSchnorrKey)) revert InvalidSchnorrVK();

// Verify that the validator can sign for that newBlsVK, otherwise it inner reverts with
// BLSSigVerificationFailed
bytes memory message = abi.encode(msg.sender);
BLSSig.verifyBlsSig(message, newBlsSig, newBlsVK);

// Update the node's bls key
node.blsVK = newBlsVK;

// Update the node's schnorr key if the newSchnorrVK is not a zero point Schnorr key
node.schnorrVK = newSchnorrVK;

// Update the node in the stake table
nodes[msg.sender] = node;

// Emit the event
emit UpdatedConsensusKeys(msg.sender, node.blsVK, node.schnorrVK);
}

/// @notice Minimum stake amount
/// @return Minimum stake amount
/// TODO: This value should be a variable modifiable by admin
69 changes: 42 additions & 27 deletions contracts/src/interfaces/AbstractStakeTable.sol
View file
Original file line number Diff line number Diff line change
@@ -24,20 +24,28 @@ abstract contract AbstractStakeTable {
uint256 public totalVotingStake;

/// @notice Signals a registration of a BLS public key.
/// @param blsVKhash hash of the BLS public key that is registered.
/// @param account the address of the validator
/// @param registerEpoch epoch when the registration becomes effective.
/// @param amountDeposited amount deposited when registering the new node.
event Registered(bytes32 blsVKhash, uint64 registerEpoch, uint256 amountDeposited);
event Registered(address account, uint64 registerEpoch, uint256 amountDeposited);

/// @notice Signals an exit request has been granted.
/// @param blsVKhash hash of the BLS public key owned by the user who requested to exit.
/// @param account the address of the validator
/// @param exitEpoch epoch when the user will be allowed to withdraw its funds.
event Exit(bytes32 blsVKhash, uint64 exitEpoch);
event Exit(address account, uint64 exitEpoch);

/// @notice Signals a deposit to a BLS public key.
/// @param blsVKhash hash of the BLS public key that received the deposit.
/// @param account the address of the validator
/// @param amount amount of the deposit
event Deposit(bytes32 blsVKhash, uint256 amount);
event Deposit(address account, uint256 amount);

/// @notice Signals a consensus key update for a validator
/// @param account the address of the validator
/// @param newBlsVK the new BLS verification key
/// @param newSchnorrVK the new Schnorr verification key
event UpdatedConsensusKeys(
address account, BN254.G2Point newBlsVK, EdOnBN254.EdOnBN254Point newSchnorrVK
);

/// @dev (sadly, Solidity doesn't support type alias on non-primitive types)
// We avoid declaring another struct even if the type info helps with readability,
@@ -55,21 +63,25 @@ abstract contract AbstractStakeTable {
/// @param balance The amount of token staked.
/// @param registerEpoch The starting epoch for the validator.
/// @param exitEpoch The ending epoch for the validator.
/// @param schnorrVK The Schnorr verification key associated.
/// @param schnorrVK The Schnorr verification key associated. Used for signing the light client
/// state.
/// @param blsVK The BLS verification key associated. Used for consensus voting.
struct Node {
address account;
uint256 balance;
uint64 registerEpoch;
uint64 exitEpoch;
EdOnBN254.EdOnBN254Point schnorrVK;
BN254.G2Point blsVK;
}

// === Table State & Stats ===

/// @notice Look up the balance of `blsVK`
function lookupStake(BN254.G2Point memory blsVK) external view virtual returns (uint256);
/// @notice Look up the full `Node` state associated with `blsVK`
function lookupNode(BN254.G2Point memory blsVK) external view virtual returns (Node memory);
/// @notice Look up the balance of `account`
function lookupStake(address account) external view virtual returns (uint256);

/// @notice Look up the full `Node` state associated with `account`
function lookupNode(address account) external view virtual returns (Node memory);

// === Queuing Stats ===

@@ -89,8 +101,7 @@ abstract contract AbstractStakeTable {
/// @param blsVK The BLS verification key
/// @param schnorrVK The Schnorr verification key (as the auxiliary info)
/// @param amount The amount to register
/// @param blsSig The BLS signature that authenticates the ethereum account this function is
/// called from
/// @param blsSig The BLS signature that the caller owns the `blsVK`
/// @param validUntilEpoch The maximum epoch the sender is willing to wait to be included
/// (cannot be smaller than the current epoch)
/// @dev No validity check on `schnorrVK`, as it's assumed to be sender's responsibility,
@@ -108,28 +119,32 @@ abstract contract AbstractStakeTable {

/// @notice Deposit more stakes to registered keys
///
/// @param blsVK The BLS verification key
/// @param amount The amount to deposit
/// @return (newBalance, effectiveEpoch) the new balance effective at a future epoch
function deposit(BN254.G2Point memory blsVK, uint256 amount)
external
virtual
returns (uint256, uint64);
function deposit(uint256 amount) external virtual returns (uint256, uint64);

/// @notice Request to exit from the stake table, not immediately withdrawable!
///
/// @param blsVK The BLS verification key to exit
function requestExit(BN254.G2Point memory blsVK) external virtual;
function requestExit() external virtual;

/// @notice Withdraw from the staking pool. Transfers occur! Only successfully exited keys can
/// withdraw past their `exitEpoch`.
///
/// @param blsVK The BLS verification key to withdraw
/// @param blsSig The BLS signature that authenticates the ethereum account this function is
/// called from the caller
/// @return The total amount withdrawn, equal to `Node.balance` associated with `blsVK`
function withdrawFunds(BN254.G2Point memory blsVK, BN254.G1Point memory blsSig)
external
virtual
returns (uint256);
function withdrawFunds() external virtual returns (uint256);

/// @notice Update the consensus keys for a validator
/// @dev This function is used to update the consensus keys for a validator
/// @dev This function can only be called by the validator itself when it's not in the exit
/// queue
/// @dev The validator will need to give up either its old BLS key and/or old Schnorr key
/// @dev The validator will need to provide a BLS signature over the new BLS key
/// @param newBlsVK The new BLS verification key
/// @param newSchnorrVK The new Schnorr verification key
/// @param newBlsSig The BLS signature that the account owns the new BLS key
function updateConsensusKeys(
BN254.G2Point memory newBlsVK,
EdOnBN254.EdOnBN254Point memory newSchnorrVK,
BN254.G1Point memory newBlsSig
) external virtual;
}
9 changes: 9 additions & 0 deletions contracts/src/libraries/EdOnBn254.sol
View file
Original file line number Diff line number Diff line change
@@ -35,5 +35,14 @@ library EdOnBN254 {
return abi.encodePacked(Utils.reverseEndianness(point.x | mask));
}

/// @dev Check if two points are equal
function isEqual(EdOnBN254Point memory a, EdOnBN254Point memory b)
internal
pure
returns (bool)
{
return a.x == b.x && a.y == b.y;
}

// TODO: (alex) add `validatePoint` methods and tests
}
456 changes: 435 additions & 21 deletions contracts/test/StakeTable.t.sol
View file
Original file line number Diff line number Diff line change
@@ -25,15 +25,13 @@ import { ExampleToken } from "../src/ExampleToken.sol";
import { StakeTable as S } from "../src/StakeTable.sol";

contract StakeTable_register_Test is Test {
event Registered(bytes32, uint64, uint256);

S public stakeTable;
ExampleToken public token;
LightClientMock public lcMock;
uint256 public constant INITIAL_BALANCE = 10 ether;
address public exampleTokenCreator;

function genClientWallet(address sender)
function genClientWallet(address sender, string memory seed)
private
returns (BN254.G2Point memory, EdOnBN254.EdOnBN254Point memory, BN254.G1Point memory)
{
@@ -42,7 +40,7 @@ contract StakeTable_register_Test is Test {
cmds[0] = "diff-test";
cmds[1] = "gen-client-wallet";
cmds[2] = vm.toString(sender);
cmds[3] = "123";
cmds[3] = seed;

bytes memory result = vm.ffi(cmds);
(
@@ -85,9 +83,10 @@ contract StakeTable_register_Test is Test {
function testFuzz_RevertWhen_InvalidBLSSig(uint256 scalar) external {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

(BN254.G2Point memory blsVK, EdOnBN254.EdOnBN254Point memory schnorrVK,) =
genClientWallet(exampleTokenCreator);
genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer
vm.startPrank(exampleTokenCreator);
@@ -152,12 +151,13 @@ contract StakeTable_register_Test is Test {
function test_RevertWhen_NodeAlreadyRegistered() external {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator);
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer
vm.prank(exampleTokenCreator);
@@ -176,12 +176,13 @@ contract StakeTable_register_Test is Test {
function test_RevertWhen_NoTokenAllowanceOrBalance() external {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 10;
string memory seed = "123";

(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator);
) = genClientWallet(exampleTokenCreator, seed);

assertEq(ERC20(token).balanceOf(exampleTokenCreator), INITIAL_BALANCE);
vm.prank(exampleTokenCreator);
@@ -192,7 +193,7 @@ contract StakeTable_register_Test is Test {

// A user with 0 balance cannot register either
address newUser = makeAddr("New user with zero balance");
(blsVK, schnorrVK, sig) = genClientWallet(newUser);
(blsVK, schnorrVK, sig) = genClientWallet(newUser, seed);

vm.startPrank(newUser);
// Prepare for the token transfer by giving the StakeTable contract the required allowance
@@ -205,31 +206,34 @@ contract StakeTable_register_Test is Test {
function test_RevertWhen_WrongStakeAmount() external {
uint64 depositAmount = 5 ether;
uint64 validUntilEpoch = 10;
string memory seed = "123";

(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator);
) = genClientWallet(exampleTokenCreator, seed);

assertEq(ERC20(token).balanceOf(exampleTokenCreator), INITIAL_BALANCE);
vm.prank(exampleTokenCreator);
// The call to register is expected to fail because the depositAmount has not been approved
// The call to register is expected to fail because the correct depositAmount has not been
// approved/supplied
// and thus the stake table contract cannot lock the stake.
vm.expectRevert(abi.encodeWithSelector(S.InsufficientStakeAmount.selector, depositAmount));
stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);
}

/// @dev Tests a correct registration
function test_Registration_succeeds() external {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator);

uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer
vm.prank(exampleTokenCreator);
@@ -242,21 +246,431 @@ contract StakeTable_register_Test is Test {
totalStakeAmount = stakeTable.totalStake();
assertEq(totalStakeAmount, 0);

AbstractStakeTable.Node memory node;
node.account = exampleTokenCreator;
node.balance = depositAmount;
node.schnorrVK = schnorrVK;
node.registerEpoch = 1;

// Check event is emitted after calling successfully `register`
vm.expectEmit(false, false, false, true, address(stakeTable));
emit Registered(stakeTable._hashBlsKey(blsVK), node.registerEpoch, node.balance);
emit AbstractStakeTable.Registered(exampleTokenCreator, 1, depositAmount);
vm.prank(exampleTokenCreator);
stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// Balance after registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE - depositAmount);
totalStakeAmount = stakeTable.totalStake();
assertEq(totalStakeAmount, depositAmount);

// lookup the node and verify the data
AbstractStakeTable.Node memory node = stakeTable.lookupNode(exampleTokenCreator);
assertEq(node.account, exampleTokenCreator);
assertEq(node.balance, depositAmount);
assertEq(node.registerEpoch, 1);
assertTrue(stakeTable._isEqualBlsKey(node.blsVK, blsVK));
assertTrue(EdOnBN254.isEqual(node.schnorrVK, schnorrVK));

// lookup the stake and verify the amount
uint256 stakeAmount = stakeTable.lookupStake(exampleTokenCreator);
assertEq(stakeAmount, depositAmount);
}

/// @dev Tests a correct registration
function test_RevertWhen_InvalidBlsVK_or_InvalidSchnorrVK_on_Registration() external {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

// generate a valid blsVK and schnorrVK
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// revert when the blsVK is the zero point
BN254.G2Point memory zeroBlsVK = BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
);
vm.expectRevert(S.InvalidBlsVK.selector);
stakeTable.register(zeroBlsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// revert when the schnorrVK is the zero point
EdOnBN254.EdOnBN254Point memory zeroSchnorrVK = EdOnBN254.EdOnBN254Point(0, 0);
vm.expectRevert(S.InvalidSchnorrVK.selector);
stakeTable.register(blsVK, zeroSchnorrVK, depositAmount, sig, validUntilEpoch);

// lookup the node and verify the data but expect the node to be empty
AbstractStakeTable.Node memory node = stakeTable.lookupNode(exampleTokenCreator);
assertEq(node.account, address(0));

vm.stopPrank();
}

function test_UpdateConsensusKeys_Succeeds() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

// Check event is emitted after calling successfully `register`
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.Registered(exampleTokenCreator, 1, depositAmount);
stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// Step 2: generate a new blsVK and schnorrVK
seed = "234";
(
BN254.G2Point memory newBlsVK,
EdOnBN254.EdOnBN254Point memory newSchnorrVK,
BN254.G1Point memory newBlsSig
) = genClientWallet(exampleTokenCreator, seed);

// assert that the new blsVK and schnorrVK are not the same as the old ones
assertFalse(stakeTable._isEqualBlsKey(newBlsVK, blsVK));
assertFalse(EdOnBN254.isEqual(newSchnorrVK, schnorrVK));

// Step 3: update the consensus keys
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator, newBlsVK, newSchnorrVK);
stakeTable.updateConsensusKeys(newBlsVK, newSchnorrVK, newBlsSig);

// Step 4: verify the update
AbstractStakeTable.Node memory node = stakeTable.lookupNode(exampleTokenCreator);
assertTrue(stakeTable._isEqualBlsKey(node.blsVK, newBlsVK));
assertTrue(EdOnBN254.isEqual(node.schnorrVK, newSchnorrVK));
assertEq(node.balance, depositAmount);
assertEq(node.account, exampleTokenCreator);

vm.stopPrank();
}

function test_RevertWhen_UpdateConsensusKeysWithSameKeys() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// Step 2: update the consensus keys with the same keys
vm.expectRevert(S.NoKeyChange.selector);
stakeTable.updateConsensusKeys(blsVK, schnorrVK, sig);

vm.stopPrank();
}

function test_RevertWhen_UpdateConsensusKeysWithEmptyKeys() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// empty keys
BN254.G2Point memory emptyBlsVK = BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
);
EdOnBN254.EdOnBN254Point memory emptySchnorrVK = EdOnBN254.EdOnBN254Point(0, 0);

// Step 2: attempt to update the consensus keys with the same keys
vm.expectRevert(S.InvalidBlsVK.selector);
stakeTable.updateConsensusKeys(emptyBlsVK, emptySchnorrVK, sig);

vm.stopPrank();
}

function test_RevertWhen_UpdateConsensusKeysWithInvalidSignature() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

BN254.G1Point memory badSig =
BN254.G1Point(BN254.BaseField.wrap(0), BN254.BaseField.wrap(0));

stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// Step 2: generate a new blsVK and schnorrVK
seed = "234";
(BN254.G2Point memory newBlsVK, EdOnBN254.EdOnBN254Point memory newSchnorrVK,) =
genClientWallet(exampleTokenCreator, seed);

// Step 3: attempt to update the consensus keys with the new keys but invalid signature
vm.expectRevert(BLSSig.BLSSigVerificationFailed.selector);
stakeTable.updateConsensusKeys(newBlsVK, newSchnorrVK, badSig);

vm.stopPrank();
}

function test_RevertWhen_UpdateConsensusKeysWithZeroBlsKeyButNewSchnorrVK() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.Registered(exampleTokenCreator, 1, depositAmount);
stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// Step 2: generate an empty and new schnorrVK
seed = "234";
(, EdOnBN254.EdOnBN254Point memory newSchnorrVK,) =
genClientWallet(exampleTokenCreator, seed);

BN254.G2Point memory emptyBlsVK = BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
);

// Step 3: update the consensus keys with the new schnorr Key but zero bls key which
// indicates no change in the bls key
vm.expectRevert(S.InvalidBlsVK.selector);
stakeTable.updateConsensusKeys(emptyBlsVK, newSchnorrVK, sig);

vm.stopPrank();
}

function test_RevertWhen_UpdateConsensusKeysWithZeroSchnorrVKButNewBlsVK() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// Step 2: generate a new blsVK
seed = "234";
(BN254.G2Point memory newBlsVK,, BN254.G1Point memory newSig) =
genClientWallet(exampleTokenCreator, seed);

// Step 3: generate empty schnorrVK
EdOnBN254.EdOnBN254Point memory emptySchnorrVK = EdOnBN254.EdOnBN254Point(0, 0);

// Step 4: update the consensus keys with the new bls keys but empty schnorrVK
vm.expectRevert(S.InvalidSchnorrVK.selector);
stakeTable.updateConsensusKeys(newBlsVK, emptySchnorrVK, newSig);

vm.stopPrank();
}

function test_UpdateConsensusKeysWithSameBlsKeyButNewSchnorrVK_Succeeds() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory blsSig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.Registered(exampleTokenCreator, 1, depositAmount);
stakeTable.register(blsVK, schnorrVK, depositAmount, blsSig, validUntilEpoch);

// Step 2: generate a new schnorrVK
seed = "234";
(, EdOnBN254.EdOnBN254Point memory newSchnorrVK,) =
genClientWallet(exampleTokenCreator, seed);

// Step 3: update the consensus keys with the new schnorrVK
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator, blsVK, newSchnorrVK);
stakeTable.updateConsensusKeys(blsVK, newSchnorrVK, blsSig);

// Step 4: verify the update
AbstractStakeTable.Node memory node = stakeTable.lookupNode(exampleTokenCreator);
assertTrue(stakeTable._isEqualBlsKey(node.blsVK, blsVK)); // same as current bls vk
assertTrue(EdOnBN254.isEqual(node.schnorrVK, newSchnorrVK)); // new schnorr vk
assertEq(node.balance, depositAmount); //same balance
assertEq(node.account, exampleTokenCreator); //same account

vm.stopPrank();
}

function test_UpdateConsensusKeysWithNewBlsKeyButSameSchnorrVK_Succeeds() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";

//Step 1: generate a new blsVK and schnorrVK and register this node
(
BN254.G2Point memory blsVK,
EdOnBN254.EdOnBN254Point memory schnorrVK,
BN254.G1Point memory sig
) = genClientWallet(exampleTokenCreator, seed);

// Prepare for the token transfer by granting allowance to the contract
vm.startPrank(exampleTokenCreator);
token.approve(address(stakeTable), depositAmount);

// Balances before registration
assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);

vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.Registered(exampleTokenCreator, 1, depositAmount);
stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);

// Step 2: generate an empty and new schnorrVK
seed = "234";
(BN254.G2Point memory newBlsVK,, BN254.G1Point memory newSig) =
genClientWallet(exampleTokenCreator, seed);

// Step 3: update the consensus keys with the same bls keys but new schnorrV
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator, newBlsVK, schnorrVK);
stakeTable.updateConsensusKeys(newBlsVK, schnorrVK, newSig);

// Step 4: verify the update
AbstractStakeTable.Node memory node = stakeTable.lookupNode(exampleTokenCreator);
assertTrue(stakeTable._isEqualBlsKey(node.blsVK, newBlsVK)); // same as current bls vk
assertTrue(EdOnBN254.isEqual(node.schnorrVK, schnorrVK)); // same as current schnorr vk
assertEq(node.balance, depositAmount); //same balance
assertEq(node.account, exampleTokenCreator); //same account

vm.stopPrank();
}

function test_lookupNodeAndLookupStake_fails() public {
address randomUser = makeAddr("randomUser");

// lookup the stake for an address that is not registered and verify the amount is empty
uint256 stakeAmount = stakeTable.lookupStake(randomUser);
assertEq(stakeAmount, 0);

// lookup the node for an address that is not registered and verify the data is empty
AbstractStakeTable.Node memory node = stakeTable.lookupNode(randomUser);
assertEq(node.account, address(0));
assertEq(node.balance, 0);
assertEq(node.registerEpoch, 0);
assertTrue(
stakeTable._isEqualBlsKey(
node.blsVK,
BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
)
)
);
assertTrue(EdOnBN254.isEqual(node.schnorrVK, EdOnBN254.EdOnBN254Point(0, 0)));

// look up the stake for the zero address and verify the amount is empty
stakeAmount = stakeTable.lookupStake(address(0));
assertEq(stakeAmount, 0);

// look up the node for the zero address and verify the data is empty
node = stakeTable.lookupNode(address(0));
assertEq(node.account, address(0));
assertEq(node.balance, 0);
assertEq(node.registerEpoch, 0);
assertTrue(
stakeTable._isEqualBlsKey(
node.blsVK,
BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
)
)
);
assertTrue(EdOnBN254.isEqual(node.schnorrVK, EdOnBN254.EdOnBN254Point(0, 0)));
}
}

0 comments on commit 97308ad

Please sign in to comment.