From 713558d8f8cb140442b39c8184d4da0dc2d1c7ab Mon Sep 17 00:00:00 2001 From: Alysia Tech Date: Fri, 6 Dec 2024 13:51:18 -0500 Subject: [PATCH] 2368 stake table registration with fixed stake (#2365) * 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 --- contracts/src/StakeTable.sol | 152 +++++++--- .../src/interfaces/AbstractStakeTable.sol | 29 +- contracts/test/StakeTable.t.sol | 262 ++++++++++++++++++ 3 files changed, 385 insertions(+), 58 deletions(-) create mode 100644 contracts/test/StakeTable.t.sol diff --git a/contracts/src/StakeTable.sol b/contracts/src/StakeTable.sol index 6a809e0c7a..1bff92ea27 100644 --- a/contracts/src/StakeTable.sol +++ b/contracts/src/StakeTable.sol @@ -21,6 +21,9 @@ contract StakeTable is AbstractStakeTable { /// account. error NodeAlreadyRegistered(); + /// Error raised when a user tries to withdraw funds from a node that is not registered. + error NodeNotRegistered(); + /// Error raised when a user tries to make a deposit or request an exit but does not control the /// node public key. error Unauthenticated(); @@ -37,14 +40,24 @@ contract StakeTable is AbstractStakeTable { // Error raised when a user tries to withdraw funds before the exit escrow period is over. error PrematureWithdrawal(); + // Error raised when this contract does not have the sufficient allowance on the stake ERC20 + // token + error InsufficientAllowance(uint256, uint256); + + // Error raised when the staker does not have the sufficient balance on the stake ERC20 token + error InsufficientBalance(uint256); + + // Error raised when the staker does not have the sufficient stake balance to withdraw + error InsufficientStakeBalance(uint256); + + // Error raised when the staker does not register with the correct stakeAmount + error InsufficientStakeAmount(uint256); + /// Mapping from a hash of a BLS key to a node struct defined in the abstract contract. mapping(bytes32 keyHash => Node node) public nodes; - /// Total native stake locked for the latest stake table (HEAD). - uint256 public totalNativeStake; - - /// Total restaked stake locked for the latest stake table (HEAD). - uint256 public totalRestakedStake; + /// Total stake locked; + uint256 public totalStake; /// Address of the native token contract. address public tokenAddress; @@ -97,18 +110,11 @@ contract StakeTable is AbstractStakeTable { return 0; } - /// @notice Total stakes of the registered keys in the latest stake table (Head). - /// @dev Given that the current implementation does not support restaking, the second value of - /// the output is set to 0. - /// @return The total stake for native token and restaked token respectively. - function totalStake() external view override returns (uint256, uint256) { - return (totalNativeStake, totalRestakedStake); - } - /// @notice Look up the balance of `blsVK` /// @param blsVK BLS public key controlled by the user. /// @return Current balance owned by the user. - function lookupStake(BN254.G2Point memory blsVK) external view override returns (uint64) { + /// 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); return node.balance; } @@ -117,11 +123,13 @@ contract StakeTable is AbstractStakeTable { /// @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 Get the next available epoch and queue size in that epoch + /// TODO modify this according to the current spec function nextRegistrationEpoch() external view override returns (uint64, uint64) { uint64 epoch; uint64 queueSize; @@ -143,17 +151,20 @@ contract StakeTable is AbstractStakeTable { // @param epoch next available registration epoch // @param queueSize current size of the registration queue (after insertion of new element in // the queue) + /// 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 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 function nextExitEpoch() external view override returns (uint64, uint64) { uint64 epoch; uint64 queueSize; @@ -174,12 +185,14 @@ contract StakeTable is AbstractStakeTable { // @notice Update the exit queue // @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 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 function numPendingExits() external view override returns (uint64) { return _numPendingExits; } @@ -198,6 +211,7 @@ contract StakeTable is AbstractStakeTable { /// withdraw. /// @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 function exitEscrowPeriod(Node memory node) public pure returns (uint64) { if (node.balance > 100) { return 10; @@ -211,29 +225,56 @@ contract StakeTable is AbstractStakeTable { /// @param blsVK The BLS verification key /// @param schnorrVK The Schnorr verification key (as the auxiliary info) /// @param amount The amount to register - /// @param stakeType The type of staking (native or restaking) /// @param blsSig The BLS signature that authenticates the ethereum account this function is /// called from /// @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, - /// the contract only treat it as auxiliary info submitted by `blsVK`. - /// @dev `blsSig` field is necessary to prevent "rogue public-key attack". + /// @dev The function will revert if the sender does not have the correct stake amount. + /// @dev The function will revert if the sender does not have the correct allowance. + /// @dev The function will revert if the sender does not have the correct balance. + /// @dev The function will revert if the sender does not have the correct BLS signature. + /// `blsSig` field is necessary to prevent "rogue public-key attack". /// The signature is over the caller address of the function to ensure that each message is /// unique. + /// @dev No validity check on `schnorrVK`, as it's assumed to be sender's responsibility, + /// the contract only treat it as auxiliary info submitted by `blsVK`. + /// @dev The function will revert if the sender does not have the correct registration epoch. function register( BN254.G2Point memory blsVK, EdOnBN254.EdOnBN254Point memory schnorrVK, - uint64 amount, - StakeType stakeType, + uint256 amount, BN254.G1Point memory blsSig, uint64 validUntilEpoch ) external override { - if (stakeType != StakeType.Native) { - revert RestakingNotImplemented(); + uint256 fixedStakeAmount = minStakeAmount(); + + // Verify that the sender amount is the minStakeAmount + if (amount < fixedStakeAmount) { + revert InsufficientStakeAmount(amount); } + bytes32 key = _hashBlsKey(blsVK); + Node memory node = nodes[key]; + + // Verify that the node is not already registered. + if (node.account != address(0x0)) { + revert NodeAlreadyRegistered(); + } + + // Verify that this contract has permissions to access the validator's stake token. + uint256 allowance = ERC20(tokenAddress).allowance(msg.sender, address(this)); + if (allowance < fixedStakeAmount) { + revert InsufficientAllowance(allowance, fixedStakeAmount); + } + + // Verify that the validator has the balance for this stake token. + uint256 balance = ERC20(tokenAddress).balanceOf(msg.sender); + if (balance < fixedStakeAmount) { + revert InsufficientBalance(balance); + } + + // Verify that the validator can sign for that blsVK bytes memory message = abi.encode(msg.sender); BLSSig.verifyBlsSig(message, blsSig, blsVK); @@ -247,42 +288,36 @@ contract StakeTable is AbstractStakeTable { } appendRegistrationQueue(registerEpoch, queueSize); - bytes32 key = _hashBlsKey(blsVK); - Node memory node = nodes[key]; + // Transfer the stake amount of ERC20 tokens from the sender to this contract. + SafeTransferLib.safeTransferFrom( + ERC20(tokenAddress), msg.sender, address(this), fixedStakeAmount + ); - // The node must not already be registered. - if (node.account != address(0x0)) { - revert NodeAlreadyRegistered(); - } + // Update the total staked amount + totalStake += fixedStakeAmount; // Create an entry for the node. node.account = msg.sender; - node.balance = amount; - node.stakeType = stakeType; + node.balance = fixedStakeAmount; node.schnorrVK = schnorrVK; node.registerEpoch = registerEpoch; nodes[key] = node; - // Lock the deposited tokens in this contract. - if (stakeType == StakeType.Native) { - totalNativeStake += amount; - SafeTransferLib.safeTransferFrom(ERC20(tokenAddress), msg.sender, address(this), amount); - } // Other case will be implemented when we support restaking - - emit Registered(key, registerEpoch, stakeType, amount); + emit Registered(key, 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, uint64 amount) + function deposit(BN254.G2Point memory blsVK, uint256 amount) external override - returns (uint64, uint64) + returns (uint256, uint64) { bytes32 key = _hashBlsKey(blsVK); Node memory node = nodes[key]; @@ -315,6 +350,7 @@ contract StakeTable is AbstractStakeTable { /// @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); @@ -349,19 +385,53 @@ contract StakeTable is AbstractStakeTable { /// 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) external override returns (uint64) { + /// 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]; + // Verify that the node is already registered. + if (node.account == address(0)) { + revert NodeNotRegistered(); + } + + // 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(); } - uint64 balance = node.balance; + + // Delete the node from the stake table. delete nodes[key]; + // Transfer the balance to the node's account. SafeTransferLib.safeTransfer(ERC20(tokenAddress), node.account, balance); return balance; } + + /// @notice Minimum stake amount + /// @return Minimum stake amount + /// TODO: This value should be a variable modifiable by admin + function minStakeAmount() public pure returns (uint256) { + return 10 ether; + } } diff --git a/contracts/src/interfaces/AbstractStakeTable.sol b/contracts/src/interfaces/AbstractStakeTable.sol index d5a37ba099..62ff0fc5e3 100644 --- a/contracts/src/interfaces/AbstractStakeTable.sol +++ b/contracts/src/interfaces/AbstractStakeTable.sol @@ -26,11 +26,8 @@ abstract contract AbstractStakeTable { /// @notice Signals a registration of a BLS public key. /// @param blsVKhash hash of the BLS public key that is registered. /// @param registerEpoch epoch when the registration becomes effective. - /// @param stakeType native or restake token. /// @param amountDeposited amount deposited when registering the new node. - event Registered( - bytes32 blsVKhash, uint64 registerEpoch, StakeType stakeType, uint256 amountDeposited - ); + event Registered(bytes32 blsVKhash, 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. @@ -55,15 +52,13 @@ abstract contract AbstractStakeTable { /// @notice Represents a HotShot validator node /// In the dual-staking model, a HotShot validator could have multiple `Node` entries. /// @param account The Ethereum account of the validator. - /// @param stakeType The type of token staked. /// @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. struct Node { address account; - StakeType stakeType; - uint64 balance; + uint256 balance; uint64 registerEpoch; uint64 exitEpoch; EdOnBN254.EdOnBN254Point schnorrVK; @@ -71,11 +66,8 @@ abstract contract AbstractStakeTable { // === Table State & Stats === - /// @notice Total stakes of the registered keys in the latest stake table (Head). - /// @return The total stake for native token and restaked token respectively. - function totalStake() external view virtual returns (uint256, uint256); /// @notice Look up the balance of `blsVK` - function lookupStake(BN254.G2Point memory blsVK) external view virtual returns (uint64); + 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); @@ -97,7 +89,6 @@ 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 stakeType The type of staking (native or restaking) /// @param blsSig The BLS signature that authenticates the ethereum account this function is /// called from /// @param validUntilEpoch The maximum epoch the sender is willing to wait to be included @@ -110,8 +101,7 @@ abstract contract AbstractStakeTable { function register( BN254.G2Point memory blsVK, EdOnBN254.EdOnBN254Point memory schnorrVK, - uint64 amount, - StakeType stakeType, + uint256 amount, BN254.G1Point memory blsSig, uint64 validUntilEpoch ) external virtual; @@ -121,10 +111,10 @@ abstract contract AbstractStakeTable { /// @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, uint64 amount) + function deposit(BN254.G2Point memory blsVK, uint256 amount) external virtual - returns (uint64, uint64); + returns (uint256, uint64); /// @notice Request to exit from the stake table, not immediately withdrawable! /// @@ -135,6 +125,11 @@ abstract contract AbstractStakeTable { /// 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) external virtual returns (uint64); + function withdrawFunds(BN254.G2Point memory blsVK, BN254.G1Point memory blsSig) + external + virtual + returns (uint256); } diff --git a/contracts/test/StakeTable.t.sol b/contracts/test/StakeTable.t.sol new file mode 100644 index 0000000000..53154850f6 --- /dev/null +++ b/contracts/test/StakeTable.t.sol @@ -0,0 +1,262 @@ +// SPDX-License-Identifier: Unlicensed + +/* solhint-disable contract-name-camelcase, func-name-mixedcase, one-contract-per-file */ + +pragma solidity ^0.8.0; + +// Libraries +import "forge-std/Test.sol"; +// import {console} from "forge-std/console.sol"; + +using stdStorage for StdStorage; + +import { ERC20 } from "solmate/utils/SafeTransferLib.sol"; +import { BN254 } from "bn254/BN254.sol"; +import { BLSSig } from "../src/libraries/BLSSig.sol"; +import { EdOnBN254 } from "../src/libraries/EdOnBn254.sol"; +import { AbstractStakeTable } from "../src/interfaces/AbstractStakeTable.sol"; +import { LightClient } from "../src/LightClient.sol"; +import { LightClientMock } from "../test/mocks/LightClientMock.sol"; + +// Token contract +import { ExampleToken } from "../src/ExampleToken.sol"; + +// Target contract +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) + private + returns (BN254.G2Point memory, EdOnBN254.EdOnBN254Point memory, BN254.G1Point memory) + { + // Generate a BLS signature and other values using rust code + string[] memory cmds = new string[](4); + cmds[0] = "diff-test"; + cmds[1] = "gen-client-wallet"; + cmds[2] = vm.toString(sender); + cmds[3] = "123"; + + bytes memory result = vm.ffi(cmds); + ( + BN254.G1Point memory blsSig, + BN254.G2Point memory blsVK, + uint256 schnorrVKx, + uint256 schnorrVKy, + ) = abi.decode(result, (BN254.G1Point, BN254.G2Point, uint256, uint256, address)); + + return ( + blsVK, // blsVK + EdOnBN254.EdOnBN254Point(schnorrVKx, schnorrVKy), // schnorrVK + blsSig // sig + ); + } + + function setUp() public { + exampleTokenCreator = makeAddr("tokenCreator"); + vm.prank(exampleTokenCreator); + token = new ExampleToken(INITIAL_BALANCE); + + string[] memory cmds = new string[](3); + cmds[0] = "diff-test"; + cmds[1] = "mock-genesis"; + cmds[2] = "5"; + + bytes memory result = vm.ffi(cmds); + ( + LightClientMock.LightClientState memory state, + LightClientMock.StakeTableState memory stakeState + ) = abi.decode(result, (LightClient.LightClientState, LightClient.StakeTableState)); + LightClientMock.LightClientState memory genesis = state; + LightClientMock.StakeTableState memory genesisStakeTableState = stakeState; + + lcMock = new LightClientMock(genesis, genesisStakeTableState, 864000); + address lightClientAddress = address(lcMock); + stakeTable = new S(address(token), lightClientAddress, 10); + } + + function testFuzz_RevertWhen_InvalidBLSSig(uint256 scalar) external { + uint64 depositAmount = 10 ether; + uint64 validUntilEpoch = 5; + + (BN254.G2Point memory blsVK, EdOnBN254.EdOnBN254Point memory schnorrVK,) = + genClientWallet(exampleTokenCreator); + + // Prepare for the token transfer + vm.startPrank(exampleTokenCreator); + token.approve(address(stakeTable), depositAmount); + + // Ensure the scalar is valid + // Note: Apparently BN254.scalarMul is not well defined when the scalar is 0 + scalar = bound(scalar, 1, BN254.R_MOD - 1); + BN254.validateScalarField(BN254.ScalarField.wrap(scalar)); + BN254.G1Point memory badSig = BN254.scalarMul(BN254.P1(), BN254.ScalarField.wrap(scalar)); + BN254.validateG1Point(badSig); + + // Failed signature verification + vm.expectRevert(BLSSig.BLSSigVerificationFailed.selector); + stakeTable.register(blsVK, schnorrVK, depositAmount, badSig, validUntilEpoch); + vm.stopPrank(); + } + + // commenting out epoch related tests for now + // function testFuzz_RevertWhen_InvalidNextRegistrationEpoch(uint64 rand) external { + // LCMock.setCurrentEpoch(3); + // uint64 currentEpoch = stakeTable.currentEpoch(); + + // uint64 depositAmount = 10 ether; + // vm.prank(exampleTokenCreator); + // token.approve(address(stakeTable), depositAmount); + + // ( + // BN254.G2Point memory blsVK, + // EdOnBN254.EdOnBN254Point memory schnorrVK, + // BN254.G1Point memory sig + // ) = genClientWallet(exampleTokenCreator); + + // // Invalid next registration epoch + // uint64 validUntilEpoch = uint64(bound(rand, 0, currentEpoch - 1)); + // vm.prank(exampleTokenCreator); + // vm.expectRevert( + // abi.encodeWithSelector( + // S.InvalidNextRegistrationEpoch.selector, currentEpoch + 1, validUntilEpoch + // ) + // ); + // stakeTable.register( + // blsVK, + // schnorrVK, + // depositAmount, + // sig, + // validUntilEpoch + // ); + + // // Valid next registration epoch + // validUntilEpoch = uint64(bound(rand, currentEpoch + 1, type(uint64).max)); + // vm.prank(exampleTokenCreator); + // stakeTable.register( + // blsVK, + // schnorrVK, + // depositAmount, + // sig, + // validUntilEpoch + // ); + // } + + function test_RevertWhen_NodeAlreadyRegistered() external { + uint64 depositAmount = 10 ether; + uint64 validUntilEpoch = 5; + + ( + BN254.G2Point memory blsVK, + EdOnBN254.EdOnBN254Point memory schnorrVK, + BN254.G1Point memory sig + ) = genClientWallet(exampleTokenCreator); + + // Prepare for the token transfer + vm.prank(exampleTokenCreator); + token.approve(address(stakeTable), depositAmount); + + // Successful call to register + vm.prank(exampleTokenCreator); + stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch); + + // The node is already registered + vm.prank(exampleTokenCreator); + vm.expectRevert(S.NodeAlreadyRegistered.selector); + stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch); + } + + function test_RevertWhen_NoTokenAllowanceOrBalance() external { + uint64 depositAmount = 10 ether; + uint64 validUntilEpoch = 10; + + ( + BN254.G2Point memory blsVK, + EdOnBN254.EdOnBN254Point memory schnorrVK, + BN254.G1Point memory sig + ) = genClientWallet(exampleTokenCreator); + + 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 + // and thus the stake table contract cannot lock the stake. + vm.expectRevert(abi.encodeWithSelector(S.InsufficientAllowance.selector, 0, depositAmount)); + stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch); + + // A user with 0 balance cannot register either + address newUser = makeAddr("New user with zero balance"); + (blsVK, schnorrVK, sig) = genClientWallet(newUser); + + vm.startPrank(newUser); + // Prepare for the token transfer by giving the StakeTable contract the required allowance + token.approve(address(stakeTable), depositAmount); + vm.expectRevert(abi.encodeWithSelector(S.InsufficientBalance.selector, 0)); + stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch); + vm.stopPrank(); + } + + function test_RevertWhen_WrongStakeAmount() external { + uint64 depositAmount = 5 ether; + uint64 validUntilEpoch = 10; + + ( + BN254.G2Point memory blsVK, + EdOnBN254.EdOnBN254Point memory schnorrVK, + BN254.G1Point memory sig + ) = genClientWallet(exampleTokenCreator); + + 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 + // 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 { + ( + BN254.G2Point memory blsVK, + EdOnBN254.EdOnBN254Point memory schnorrVK, + BN254.G1Point memory sig + ) = genClientWallet(exampleTokenCreator); + + uint64 depositAmount = 10 ether; + uint64 validUntilEpoch = 5; + + // Prepare for the token transfer + vm.prank(exampleTokenCreator); + token.approve(address(stakeTable), depositAmount); + + // Balances before registration + assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE); + + uint256 totalStakeAmount; + 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); + 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); + } +}