2370 staketable update validator identification keys

contracts/src/StakeTable.sol

@@ -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