add the ability to update consensus keys

contracts/src/StakeTable.sol

@@ -53,6 +53,9 @@ 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 key
+    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;
 
@@ -419,6 +422,81 @@ contract StakeTable is AbstractStakeTable {
         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 over the new BLS key
+    /// @param currBlsVK The current BLS verification 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
+    /// TODO: This function should be tested
+    function updateConsensusKeys(
+        BN254.G2Point memory currBlsVK,
+        BN254.G2Point memory newBlsVK,
+        EdOnBN254.EdOnBN254Point memory newSchnorrVK,
+        BN254.G1Point memory newBlsSig
+    ) external override {
+        bytes32 currBlsKey = _hashBlsKey(currBlsVK);
+        bytes32 newBlsKey = _hashBlsKey(newBlsVK);
+        Node memory node = nodes[currBlsKey];
+
+        // Verify that the msg.sender is the validator's staking address
+        if (msg.sender != node.account) {
+            revert Unauthenticated();
+        }
+
+        // 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();
+        }
+
+        // The staker does not provide a new key
+        if (
+            (newBlsKey == currBlsKey && EdOnBN254.isEqual(newSchnorrVK, node.schnorrVK))
+                || (
+                    newBlsKey == bytes32(0)
+                        && EdOnBN254.isEqual(newSchnorrVK, EdOnBN254.EdOnBN254Point(0, 0))
+                )
+        ) {
+            revert NoKeyChange();
+        }
+
+        // Update the node's schnorr key
+        if (
+            !EdOnBN254.isEqual(newSchnorrVK, node.schnorrVK)
+                && !EdOnBN254.isEqual(newSchnorrVK, EdOnBN254.EdOnBN254Point(0, 0))
+        ) {
+            node.schnorrVK = newSchnorrVK;
+        }
+
+        // Update the node's bls key
+        if (newBlsKey != currBlsKey && newBlsKey != bytes32(0)) {
+            // Verify that the validator can sign for that blsVK
+            bytes memory message = abi.encode(msg.sender);
+            BLSSig.verifyBlsSig(message, newBlsSig, newBlsVK);
+
+            // Update the node's bls key
+            node.blsVK = newBlsVK;
+        }
+
+        // Update the node in the stake table
+        nodes[newBlsKey] = node;
+
+        // Delete the old node from the stake table
+        delete nodes[currBlsKey];
+
+        // Emit the event
+        emit ConsensusKeysUpdated(msg.sender);
+    }
+
     /// @notice Minimum stake amount
     /// @return Minimum stake amount
     /// TODO: This value should be a variable modifiable by admin

contracts/src/interfaces/AbstractStakeTable.sol

@@ -39,6 +39,10 @@ abstract contract AbstractStakeTable {
     /// @param amount amount of the deposit
     event Deposit(bytes32 blsVKhash, uint256 amount);
 
+    /// @notice Signals a consensus key update for a validator
+    /// @param stakingAddr the address of the validator
+    event ConsensusKeysUpdated(address stakingAddr);
+
     /// @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,
     // extra layer of struct introduces overhead and more gas cost.
@@ -130,4 +134,21 @@ abstract contract AbstractStakeTable {
     /// @param blsVK The BLS verification key to withdraw
     /// @return The total amount withdrawn, equal to `Node.balance` associated with `blsVK`
     function withdrawFunds(BN254.G2Point memory blsVK) 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 currBlsVK The current BLS verification 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 currBlsVK,
+        BN254.G2Point memory newBlsVK,
+        EdOnBN254.EdOnBN254Point memory newSchnorrVK,
+        BN254.G1Point memory newBlsSig
+    ) external virtual;
 }

contracts/src/libraries/EdOnBn254.sol

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