feat(contract): add unvote function and votersCount to Consensus …

…contract (#731) * Test vote * Add unvote * Rename tests * Test events * Add votersCount * Test double vote and vote for registered validator * Prevent vote for resigned validator * Test vote for self * Add contract * Add todo * Fix validator set * style: resolve style guide violations --------- Co-authored-by: sebastijankuzner <[email protected]>
ArkEcosystem · Oct 15, 2024 · 35cd52c · 35cd52c
1 parent 0637bf5
commit 35cd52c
Show file tree

Hide file tree

Showing 4 changed files with 351 additions and 328 deletions.
diff --git a/contracts/src/consensus/Consensus.sol b/contracts/src/consensus/Consensus.sol
@@ -1,6 +1,7 @@
 pragma solidity ^0.8.27;
 
 struct ValidatorData {
+	uint256 votersCount;
 	uint256 voteBalance;
 	bool isResigned;
 	bytes bls12_381_public_key; // 96 bits
@@ -206,6 +207,7 @@ contract Consensus {
 		_checkBls12_381PublicKey(bls12_381_public_key);
 
 		ValidatorData memory validator = ValidatorData({
+			votersCount: 0,
 			voteBalance: 0,
 			isResigned: false,
 			bls12_381_public_key: bls12_381_public_key
@@ -297,16 +299,32 @@ contract Consensus {
 	}
 
 	function vote(address addr) external {
-		require(isValidatorRegistered(addr), "must vote for validator");
-		require(_votes[msg.sender].validator == address(0), "TODO: already voted");
+		require(isValidatorRegistered(addr), "Must vote for validator");
+		require(_votes[msg.sender].validator == address(0), "Already voted");
+
+		ValidatorData storage validatorData = _registeredValidatorData[addr];
+		require(!validatorData.isResigned, "Must vote for unresigned validator");
 
 		_votes[msg.sender] = Vote({validator: addr, balance: msg.sender.balance});
+
 		// TODO: safe math
-		_registeredValidatorData[addr].voteBalance += msg.sender.balance;
+		validatorData.voteBalance += msg.sender.balance;
+		validatorData.votersCount += 1;
 
 		emit Voted(msg.sender, addr);
 	}
 
+	function unvote() external {
+		Vote storage voter = _votes[msg.sender];
+		require(voter.validator != address(0), "TODO: not voted");
+
+		emit Unvoted(msg.sender, voter.validator);
+
+		_registeredValidatorData[voter.validator].voteBalance -= voter.balance;
+		_registeredValidatorData[voter.validator].votersCount -= 1;
+		delete _votes[msg.sender];
+	}
+
 	function updateVoters(address[] calldata voters) external {
 		// TODO: limit number of voters per update?
 		for (uint i = 0; i < voters.length; i++) {

diff --git a/contracts/test/consensus/Consensus-Vote.sol b/contracts/test/consensus/Consensus-Vote.sol
@@ -0,0 +1,239 @@
+// SPDX-License-Identifier: GNU GENERAL PUBLIC LICENSE
+pragma solidity ^0.8.13;
+
+import {Test, console} from "@forge-std/Test.sol";
+import {Consensus, ValidatorData, Validator, Unvoted, Voted} from "@contracts/consensus/Consensus.sol";
+
+contract ConsensusTest is Test {
+	Consensus public consensus;
+
+	function setUp() public {
+		consensus = new Consensus();
+	}
+
+	function registerValidator(address addr) internal {
+		bytes32 h = keccak256(abi.encode(addr));
+		bytes memory validatorKey = new bytes(48);
+		for (uint256 j = 0; j < 32; j++) {
+			validatorKey[j] = h[j];
+		}
+
+		vm.startPrank(addr);
+		consensus.registerValidator(validatorKey);
+		vm.stopPrank();
+
+		Validator memory validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 0 ether);
+		assertEq(validator.data.votersCount, 0);
+	}
+
+	function deregisterValidator(address addr) internal {
+		vm.startPrank(addr);
+		consensus.deregisterValidator();
+		vm.stopPrank();
+	}
+
+	function test_vote() public {
+		// Register validator
+		address addr = address(1);
+		registerValidator(addr);
+
+		// Prepare voter
+		address voterAddr = address(2);
+		vm.deal(voterAddr, 100 ether);
+
+		// Vote
+		vm.startPrank(voterAddr);
+		vm.expectEmit(address(consensus));
+		emit Voted(voterAddr, addr);
+		consensus.vote(addr);
+		vm.stopPrank();
+
+		// Assert voteBalance and voter balance
+		Validator memory validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 100 ether);
+		assertEq(validator.data.votersCount, 1);
+		assertEq(voterAddr.balance, 100 ether);
+
+		// Update vote should correctly update the vote balance
+		// Let say voter has 90 eth at the end of the block
+		vm.deal(voterAddr, 90 ether);
+
+		address[] memory voters = new address[](1);
+		voters[0] = voterAddr;
+		consensus.updateVoters(voters);
+
+		// Assert voteBalance and voter balance
+		validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 90 ether);
+		assertEq(validator.data.votersCount, 1);
+		assertEq(voterAddr.balance, 90 ether);
+	}
+
+	function test_vote_allow_self_vote() public {
+		// Register validator
+		address addr = address(1);
+		registerValidator(addr);
+
+		// Prepare voter
+		address voterAddr = addr;
+		vm.deal(voterAddr, 100 ether);
+
+		// Vote
+		vm.startPrank(voterAddr);
+		vm.expectEmit(address(consensus));
+		emit Voted(voterAddr, addr);
+		consensus.vote(addr);
+		vm.stopPrank();
+
+		// Assert voteBalance and voter balance
+		Validator memory validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 100 ether);
+		assertEq(validator.data.votersCount, 1);
+		assertEq(voterAddr.balance, 100 ether);
+
+		// Update vote should correctly update the vote balance
+		// Let say voter has 90 eth at the end of the block
+		vm.deal(voterAddr, 90 ether);
+
+		address[] memory voters = new address[](1);
+		voters[0] = voterAddr;
+		consensus.updateVoters(voters);
+
+		// Assert voteBalance and voter balance
+		validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 90 ether);
+		assertEq(validator.data.votersCount, 1);
+		assertEq(voterAddr.balance, 90 ether);
+	}
+
+	function test_vote_prevent_double_vote() public {
+		// Register validator
+		address addr = address(1);
+		registerValidator(addr);
+
+		// Prepare voter
+		address voterAddr = address(2);
+
+		vm.startPrank(voterAddr);
+		vm.expectEmit(address(consensus));
+		emit Voted(voterAddr, addr);
+		consensus.vote(addr);
+
+		vm.expectRevert("Already voted");
+		consensus.vote(addr);
+	}
+
+	function test_vote_prevent_for_unregistered_validator() public {
+		vm.expectRevert("Must vote for validator");
+		consensus.vote(address(1));
+	}
+
+	function test_vote_prevent_for_resigned_validator() public {
+		// Register validator
+		address addr = address(1);
+		registerValidator(addr);
+		deregisterValidator(addr);
+
+		// Prepare voter
+		address voterAddr = address(2);
+		vm.startPrank(voterAddr);
+		vm.expectRevert("Must vote for unresigned validator");
+		consensus.vote(addr);
+	}
+
+	function test_unvote_and_vote_in_same_block() public {
+		// Register validator
+		address addr = address(1);
+		registerValidator(addr);
+
+		// Vote
+		address voterAddr = address(2);
+		vm.deal(voterAddr, 100 ether);
+		vm.startPrank(voterAddr);
+		consensus.vote(addr);
+		vm.stopPrank();
+
+		// Assert voteBalance and voter balance
+		Validator memory validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 100 ether);
+		assertEq(validator.data.votersCount, 1);
+		assertEq(voterAddr.balance, 100 ether);
+
+		// Let say voter has 90 eth after some tx
+		vm.deal(voterAddr, 90 ether);
+
+		// Unvote
+		vm.startPrank(voterAddr);
+		vm.expectEmit(address(consensus));
+		emit Unvoted(voterAddr, addr);
+		consensus.unvote();
+		vm.stopPrank();
+
+		// Assert voteBalance and voter balance
+		validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 0 ether);
+		assertEq(validator.data.votersCount, 0);
+		assertEq(voterAddr.balance, 90 ether);
+	}
+
+		function test_unvote_and_vote_in_different_blocks() public {
+		// Register validator
+		address addr = address(1);
+		registerValidator(addr);
+
+		// Vote
+		address voterAddr = address(2);
+		vm.deal(voterAddr, 100 ether);
+		vm.startPrank(voterAddr);
+		consensus.vote(addr);
+		vm.stopPrank();
+
+		// Assert voteBalance and voter balance
+		Validator memory validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 100 ether);
+		assertEq(validator.data.votersCount, 1);
+		assertEq(voterAddr.balance, 100 ether);
+
+		// Update vote should correctly update the vote balance
+		// Let say voter has 90 eth at the end of the block
+		vm.deal(voterAddr, 90 ether);
+		address[] memory voters = new address[](1);
+		voters[0] = voterAddr;
+		consensus.updateVoters(voters);
+
+		// Assert voteBalance and voter balance
+		validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 90 ether);
+		assertEq(validator.data.votersCount, 1);
+		assertEq(voterAddr.balance, 90 ether);
+
+		// Let say voter has 80 eth after some tx
+		vm.deal(voterAddr, 80 ether);
+
+		// Unvote
+		vm.startPrank(voterAddr);
+		vm.expectEmit(address(consensus));
+		emit Unvoted(voterAddr, addr);
+		consensus.unvote();
+		vm.stopPrank();
+
+		// Assert voteBalance and voter balance
+		validator = consensus.getValidator(addr);
+		assertEq(validator.addr, addr);
+		assertEq(validator.data.voteBalance, 0 ether);
+		assertEq(validator.data.votersCount, 0);
+		assertEq(voterAddr.balance, 80 ether);
+	}
+
+	// TODO: Test multiple votes
+}
diff --git a/packages/evm-consensus/source/validator-set.ts b/packages/evm-consensus/source/validator-set.ts
@@ -89,7 +89,7 @@ export class ValidatorSet implements Contracts.ValidatorSet.Service {
 
 		const validatorWallets: Contracts.State.ValidatorWallet[] = [];
 		for (const [, validator] of validators.entries()) {
-			const [address, [voteBalance, , blsPublicKey]] = validator;
+			const [address, [, voteBalance, , blsPublicKey]] = validator;
 
 			const validatorWallet: Contracts.State.ValidatorWallet = {
 				address,