Skip to content

Commit

Permalink
contracts: only require BLS key to remove staker
Browse files Browse the repository at this point in the history 
Changes solidity code (incl. solidity tests) but no updates to bindings
and rust code yet.
  • Loading branch information
@sveitser
sveitser committed Dec 19, 2024
1 parent 3fd450f commit fc2df6c
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 29 deletions.
10 changes: 5 additions & 5 deletions contracts/src/PermissionedStakeTable.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import { EdOnBN254 } from "./libraries/EdOnBn254.sol";
* @dev An stake table mapping with owner-only access control.
*/
contract PermissionedStakeTable is Ownable {
event StakersUpdated(NodeInfo[] removed, NodeInfo[] added);
event StakersUpdated(BN254.G2Point[] removed, NodeInfo[] added);

error StakerAlreadyExists(BN254.G2Point);
error StakerNotFound(BN254.G2Point);
Expand All @@ -33,7 +33,7 @@ contract PermissionedStakeTable is Ownable {

// public methods

function update(NodeInfo[] memory stakersToRemove, NodeInfo[] memory newStakers)
function update(BN254.G2Point[] memory stakersToRemove, NodeInfo[] memory newStakers)
public
onlyOwner
{
Expand All @@ -55,12 +55,12 @@ contract PermissionedStakeTable is Ownable {
}
}

function remove(NodeInfo[] memory stakersToRemove) internal {
function remove(BN254.G2Point[] memory stakersToRemove) internal {
// TODO: revert if array empty
for (uint256 i = 0; i < stakersToRemove.length; i++) {
bytes32 stakerID = _hashBlsKey(stakersToRemove[i].blsVK);
bytes32 stakerID = _hashBlsKey(stakersToRemove[i]);
if (!stakers[stakerID]) {
revert StakerNotFound(stakersToRemove[i].blsVK);
revert StakerNotFound(stakersToRemove[i]);
}
stakers[stakerID] = false;
}
Expand Down
63 changes: 39 additions & 24 deletions contracts/test/PermissionedStakeTable.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,37 +38,53 @@ contract PermissionedStakeTableTest is Test {
return ps;
}

// Convert NodeInfo array to BLS keys array
function toBls(PermissionedStakeTable.NodeInfo[] memory nodes)
private
returns (BN254.G2Point[] memory)
{
BN254.G2Point[] memory bls = new BN254.G2Point[](nodes.length);
for (uint64 i = 0; i < nodes.length; i++) {
bls[i] = nodes[i].blsVK;
}
return bls;
}

// Empty array of NodeInfo
function emptyNodes() private returns (PermissionedStakeTable.NodeInfo[] memory nodes) {}

// Empty array of BLS keys
function emptyKeys() private returns (BN254.G2Point[] memory nodes) {}


function testInsert() public {
vm.prank(owner);
PermissionedStakeTable.NodeInfo[] memory stakers = nodes(1, 1);
PermissionedStakeTable.NodeInfo[] memory empty = nodes(1, 0);

vm.expectEmit();
emit PermissionedStakeTable.StakersUpdated(empty, stakers);
emit PermissionedStakeTable.StakersUpdated(emptyKeys(), stakers);

stakeTable.update(empty, stakers);
stakeTable.update(emptyKeys(), stakers);

assertTrue(stakeTable.isStaker(stakers[0].blsVK));
}

function testInsertMany() public {
vm.prank(owner);
PermissionedStakeTable.NodeInfo[] memory stakers = nodes(1, 10);
PermissionedStakeTable.NodeInfo[] memory empty = nodes(1, 0);

vm.expectEmit();
emit PermissionedStakeTable.StakersUpdated(empty, stakers);
emit PermissionedStakeTable.StakersUpdated(emptyKeys(), stakers);

stakeTable.update(empty, stakers);
stakeTable.update(emptyKeys(), stakers);

assertTrue(stakeTable.isStaker(stakers[0].blsVK));
}

function testInsertRevertsIfStakerExists() public {
vm.prank(owner);
PermissionedStakeTable.NodeInfo[] memory stakers = nodes(1, 1);
PermissionedStakeTable.NodeInfo[] memory empty = nodes(1, 0);
stakeTable.update(empty, stakers);
stakeTable.update(emptyKeys(), stakers);

// Try adding the same staker again
vm.expectRevert(
Expand All @@ -77,34 +93,35 @@ contract PermissionedStakeTableTest is Test {
)
);
vm.prank(owner);
stakeTable.update(empty, stakers);
stakeTable.update(emptyKeys(), stakers);
}

function testRemove() public {
PermissionedStakeTable.NodeInfo[] memory stakers = nodes(1, 1);
PermissionedStakeTable.NodeInfo[] memory empty = nodes(1, 0);
PermissionedStakeTable.NodeInfo[] memory stakersToInsert = nodes(1, 1);
BN254.G2Point[] memory keysToRemove = toBls(stakersToInsert);
vm.prank(owner);
stakeTable.update(empty, stakers);

// Insert the stakers we want to remove later.
stakeTable.update(emptyKeys(), stakersToInsert);

vm.prank(owner);

vm.expectEmit();
emit PermissionedStakeTable.StakersUpdated(stakers, empty);
emit PermissionedStakeTable.StakersUpdated(keysToRemove, emptyNodes());

stakeTable.update(stakers, empty);
stakeTable.update(keysToRemove, emptyNodes());

assertFalse(stakeTable.isStaker(stakers[0].blsVK));
assertFalse(stakeTable.isStaker(keysToRemove[0]));
}

function testRemoveRevertsIfStakerNotFound() public {
vm.prank(owner);
PermissionedStakeTable.NodeInfo[] memory stakers = nodes(1, 1);
PermissionedStakeTable.NodeInfo[] memory empty = nodes(1, 0);
BN254.G2Point[] memory keysToRemove = toBls(nodes(1, 1));
vm.expectRevert(
abi.encodeWithSelector(PermissionedStakeTable.StakerNotFound.selector, stakers[0].blsVK)
abi.encodeWithSelector(PermissionedStakeTable.StakerNotFound.selector, keysToRemove[0])
);
// Attempt to remove a non-existent staker
stakeTable.update(stakers, empty);
stakeTable.update(keysToRemove, emptyNodes());
}

function testNonOwnerCannotInsert() public {
Expand All @@ -113,17 +130,15 @@ contract PermissionedStakeTableTest is Test {
abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(2))
);
PermissionedStakeTable.NodeInfo[] memory stakers = nodes(1, 1);
PermissionedStakeTable.NodeInfo[] memory empty = nodes(1, 0);
stakeTable.update(empty, stakers);
stakeTable.update(emptyKeys(), stakers);
}

function testNonOwnerCannotRemove() public {
vm.prank(address(2));
vm.expectRevert(
abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(2))
);
PermissionedStakeTable.NodeInfo[] memory stakers = nodes(1, 1);
PermissionedStakeTable.NodeInfo[] memory empty = nodes(1, 0);
stakeTable.update(stakers, empty);
BN254.G2Point[] memory keys = toBls(nodes(1, 1));
stakeTable.update(keys, emptyNodes());
}
}

0 comments on commit fc2df6c

Please sign in to comment.