Skip to content

Commit

Permalink
updateConsensusKeys reverts when the keys the same as the old ones, n…
Browse files Browse the repository at this point in the history 
…o change occurs if both keys are zero and only one key is changed if the other is a zero key
  • Loading branch information
@alysiahuggins
alysiahuggins committed Jan 8, 2025
1 parent 48db5c4 commit 5f3a988
Show file tree
Hide file tree
Showing 3 changed files with 129 additions and 76 deletions.
73 changes: 25 additions & 48 deletions contracts/src/StakeTable.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,15 +55,15 @@ 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();

// 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(address account => Node node) public nodes;

Expand Down Expand Up @@ -479,7 +479,6 @@ contract StakeTable is AbstractStakeTable {
/// @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: consider emitting the changed keys
function updateConsensusKeys(
BN254.G2Point memory newBlsVK,
EdOnBN254.EdOnBN254Point memory newSchnorrVK,
Expand All @@ -488,66 +487,44 @@ contract StakeTable is AbstractStakeTable {
Node memory node = nodes[msg.sender];

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

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

// The staker does not provide a key change (both keys are the same or both keys are zero)
if (
(
_isEqualBlsKey(newBlsVK, node.blsVK)
&& EdOnBN254.isEqual(newSchnorrVK, node.schnorrVK)
)
|| (
_isEqualBlsKey(
newBlsVK,
BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
)
) && EdOnBN254.isEqual(newSchnorrVK, EdOnBN254.EdOnBN254Point(0, 0))
)
) {
// Verify that the keys are not the same as the old ones
if (_isEqualBlsKey(newBlsVK, node.blsVK) && newSchnorrVK.isEqual(node.schnorrVK)) {
revert NoKeyChange();
}

// Update the node's schnorr key once it's nonzero
if (!newSchnorrVK.isEqual(EdOnBN254.EdOnBN254Point(0, 0))) {
node.schnorrVK = newSchnorrVK;
}
// 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);

// Update the node's bls key once it's nonzero
if (
!_isEqualBlsKey(
newBlsVK,
BN254.G2Point(
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0),
BN254.BaseField.wrap(0)
)
)
) {
// Verify that the validator can sign for that blsVK
// Update the node's bls key if the newBlsVK is not a zero point BLS key
if (!_isEqualBlsKey(newBlsVK, zeroBlsKey)) {
// 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
if (!newSchnorrVK.isEqual(zeroSchnorrKey)) {
node.schnorrVK = newSchnorrVK;
}

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

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

/// @notice Minimum stake amount
Expand Down
6 changes: 5 additions & 1 deletion contracts/src/interfaces/AbstractStakeTable.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,11 @@ abstract contract AbstractStakeTable {

/// @notice Signals a consensus key update for a validator
/// @param account the address of the validator
event UpdatedConsensusKeys(address account);
/// @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,
Expand Down
126 changes: 99 additions & 27 deletions contracts/test/StakeTable.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -309,7 +309,7 @@ contract StakeTable_register_Test is Test {
vm.stopPrank();
}

function test_UpdateConsensusKeys_succeeds() public {
function test_UpdateConsensusKeys_Succeeds() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";
Expand Down Expand Up @@ -346,6 +346,8 @@ contract StakeTable_register_Test is Test {
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
Expand Down Expand Up @@ -379,14 +381,14 @@ contract StakeTable_register_Test is Test {

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

// Step 2: attempt to update the consensus keys with the same keys
// 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 {
function test_UpdateConsensusKeysWithEmptyKeys_CausesNoChange() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";
Expand Down Expand Up @@ -417,7 +419,8 @@ contract StakeTable_register_Test is Test {
EdOnBN254.EdOnBN254Point memory emptySchnorrVK = EdOnBN254.EdOnBN254Point(0, 0);

// Step 2: attempt to update the consensus keys with the same keys
vm.expectRevert(S.NoKeyChange.selector);
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator, blsVK, schnorrVK);
stakeTable.updateConsensusKeys(emptyBlsVK, emptySchnorrVK, sig);

vm.stopPrank();
Expand Down Expand Up @@ -459,9 +462,7 @@ contract StakeTable_register_Test is Test {
vm.stopPrank();
}

function test_UpdateConsensusKeysWithInvalidBlsKeyInfoButOnlySchnorrVKChanged_Succeeds()
public
{
function test_UpdateConsensusKeysWithZeroBlsKeyButNewSchnorrVK_Succeeds() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";
Expand Down Expand Up @@ -499,7 +500,7 @@ contract StakeTable_register_Test is Test {
// Step 3: update the consensus keys with the new schnorr Key but zero bls key which
// indicates no change in the bls key
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator);
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator, blsVK, newSchnorrVK);
stakeTable.updateConsensusKeys(emptyBlsVK, newSchnorrVK, sig);

// Step 4: verify the update
Expand All @@ -509,26 +510,10 @@ contract StakeTable_register_Test is Test {
assertEq(node.balance, depositAmount); //same balance
assertEq(node.account, exampleTokenCreator); //same account

// Step 5: update the consensus keys with the same bls keys but new schnorrVK
seed = "235";
(, EdOnBN254.EdOnBN254Point memory newSchnorrVK2,) =
genClientWallet(exampleTokenCreator, seed);

vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator);
stakeTable.updateConsensusKeys(blsVK, newSchnorrVK2, sig);

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

vm.stopPrank();
}

function test_UpdateConsensusKeysWithInvalidSchnorrVKButNewBlsVK_Succeeds() public {
function test_UpdateConsensusKeysWithZeroSchnorrVKButNewBlsVK_Succeeds() public {
uint64 depositAmount = 10 ether;
uint64 validUntilEpoch = 5;
string memory seed = "123";
Expand Down Expand Up @@ -556,7 +541,7 @@ contract StakeTable_register_Test is Test {

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

// Step 4: verify the update
Expand All @@ -574,7 +559,7 @@ contract StakeTable_register_Test is Test {
(BN254.G2Point memory newBlsVK2,, BN254.G1Point memory newSig2) =
genClientWallet(exampleTokenCreator, seed);
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator);
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator, newBlsVK2, schnorrVK);
stakeTable.updateConsensusKeys(newBlsVK2, emptySchnorrVK, newSig2);

// Step 7: verify the update
Expand All @@ -587,6 +572,93 @@ contract StakeTable_register_Test is Test {
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 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 schnorrVK and new blsVK
EdOnBN254.EdOnBN254Point memory emptySchnorrVK = EdOnBN254.EdOnBN254Point(0, 0);
seed = "234";
(BN254.G2Point memory newBlsVK,, BN254.G1Point memory newSig) =
genClientWallet(exampleTokenCreator, seed);

// Step 3: update the consensus keys with the new bls keys but empty schnorrVK
vm.expectEmit(false, false, false, true, address(stakeTable));
emit AbstractStakeTable.UpdatedConsensusKeys(exampleTokenCreator, newBlsVK, schnorrVK);
stakeTable.updateConsensusKeys(newBlsVK, emptySchnorrVK, 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_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");

Expand Down

0 comments on commit 5f3a988

Please sign in to comment.