feat(contracts): add onlyOwner & preventOwner modifiers to Consen…

…sus.sol (#737)

* Allow calculateTopValidators

* onlyOwner on updateVoters

* Prevent owner for vote and validator registration

* Format

* Update contract

contracts/src/consensus/Consensus.sol

@@ -63,6 +63,16 @@ contract Consensus {
 		_owner = msg.sender;
 	}
 
+	modifier onlyOwner() {
+		require(msg.sender == _owner, "Caller is not the contract owner");
+		_;
+	}
+
+	modifier preventOwner() {
+		require(msg.sender != _owner, "Caller is the contract owner");
+		_;
+	}
+
 	function shuffle() internal {
 		uint256 n = _registeredValidators.length;
 		for (uint256 i = n - 1; i > 0; i--) {
@@ -87,7 +97,7 @@ contract Consensus {
 		_topValidatorsCount = 0;
 	}
 
-	function calculateTopValidators(uint8 n) external {
+	function calculateTopValidators(uint8 n) external onlyOwner {
 		shuffle();
 		deleteTopValidators();
 
@@ -102,17 +112,16 @@ contract Consensus {
 			address addr = _registeredValidators[i];
 
 			ValidatorData storage data = _registeredValidatorData[addr];
-			if(data.isResigned) {
+			if (data.isResigned) {
 				continue;
 			}
 
-			if(_head == address(0)) {
+			if (_head == address(0)) {
 				_head = addr;
 				_topValidatorsCount = 1;
 				continue;
 			}
 
-
 			if (_topValidatorsCount < top) {
 				insertTopValidator(addr, top);
 				continue;
@@ -201,15 +210,15 @@ contract Consensus {
 	}
 
 	function getAllValidators() public view returns (Validator[] memory) {
-        Validator[] memory result = new Validator[](_registeredValidators.length);
-        for (uint i = 0; i < _registeredValidators.length; i++) {
-            address addr = _registeredValidators[i];
-            ValidatorData storage data = _registeredValidatorData[addr];
-            result[i] = Validator({addr: addr, data: data});
-        }
+		Validator[] memory result = new Validator[](_registeredValidators.length);
+		for (uint i = 0; i < _registeredValidators.length; i++) {
+			address addr = _registeredValidators[i];
+			ValidatorData storage data = _registeredValidatorData[addr];
+			result[i] = Validator({addr: addr, data: data});
+		}
 
-        return result;
-    }
+		return result;
+	}
 
 	function registeredValidatorsCount() public view returns (uint256) {
 		return _registeredValidatorsCount;
@@ -223,8 +232,7 @@ contract Consensus {
 		return _calculatedTopValidators.length;
 	}
 
-	function registerValidator(bytes calldata bls12_381_public_key) external {
-		require(msg.sender != _owner, "Invalid caller");
+	function registerValidator(bytes calldata bls12_381_public_key) external preventOwner {
 		require(!_hasRegisteredValidator[msg.sender], "Validator is already registered");
 
 		bytes32 bls_public_key_hash = keccak256(bls12_381_public_key);
@@ -279,7 +287,7 @@ contract Consensus {
 		_registeredValidatorData[_validator.addr] = _validator.data;
 	}
 
-	function vote(address addr) external {
+	function vote(address addr) external preventOwner {
 		require(isValidatorRegistered(addr), "Must vote for validator");
 		require(_votes[msg.sender].validator == address(0), "Already voted");
 
@@ -306,7 +314,7 @@ contract Consensus {
 		delete _votes[msg.sender];
 	}
 
-	function updateVoters(address[] calldata voters) external {
+	function updateVoters(address[] calldata voters) external onlyOwner {
 		// TODO: limit number of voters per update?
 		for (uint i = 0; i < voters.length; i++) {
 			_updateVoter(voters[i]);

contracts/test/consensus/Consensus-CalculateTop.sol

@@ -24,6 +24,13 @@ contract ConsensusTest is Base {
 		assertEq(validators[0].addr, addr);
 	}
 
+	function test_should_allow_only_caller() public {
+		address addr = address(1);
+		vm.startPrank(addr);
+		vm.expectRevert("Caller is not the contract owner");
+		consensus.calculateTopValidators(1);
+	}
+
 	function test_should_ignore_resigned_validators() public {
 		address addr = address(1);
 

contracts/test/consensus/Consensus-GetAllValidators.sol

@@ -6,31 +6,29 @@ import {Consensus, ValidatorData, Validator} from "@contracts/consensus/Consensu
 import {Base} from "./Base.sol";
 
 contract ConsensusTest is Base {
-    Consensus public consensus;
+	Consensus public consensus;
 
-    function setUp() public {
-        consensus = new Consensus();
-    }
+	function setUp() public {
+		consensus = new Consensus();
+	}
 
+	function test_200_validators() public {
+		vm.pauseGasMetering();
+		assertEq(consensus.registeredValidatorsCount(), 0);
 
-    function test_200_validators() public {
-        vm.pauseGasMetering();
-        assertEq(consensus.registeredValidatorsCount(), 0);
+		uint256 n = 200;
+		for (uint256 i = 0; i < n; i++) {
+			address addr = address(uint160(i + 1));
 
-
-        uint256 n = 200;
-        for (uint256 i = 0; i < n; i++) {
-            address addr = address(uint160(i + 1));
-
-            vm.startPrank(addr);
+			vm.startPrank(addr);
 			consensus.registerValidator(prepareBLSKey(addr));
-            consensus.vote(addr);
-            vm.stopPrank();
-        }
+			consensus.vote(addr);
+			vm.stopPrank();
+		}
 
-        vm.resumeGasMetering();
+		vm.resumeGasMetering();
 
-        Validator[] memory validators = consensus.getAllValidators();
-        assertEq(validators.length, n);
-    }
+		Validator[] memory validators = consensus.getAllValidators();
+		assertEq(validators.length, n);
+	}
 }

contracts/test/consensus/Consensus-UpdateVoters.sol

@@ -0,0 +1,21 @@
+// 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 test_updateVoters_should_allow_only_caller() public {
+		address addr = address(1);
+		vm.startPrank(addr);
+		address[] memory voters = new address[](0);
+		vm.expectRevert("Caller is not the contract owner");
+		consensus.updateVoters(voters);
+	}
+}

contracts/test/consensus/Consensus-ValidatorRegistration.sol

@@ -33,7 +33,7 @@ contract ConsensusTest is Base {
 	}
 
 	function test_validator_registration_revert_if_caller_is_owner() public {
-		vm.expectRevert("Invalid caller");
+		vm.expectRevert("Caller is the contract owner");
 		consensus.registerValidator(prepareBLSKey(address(1)));
 	}
 

contracts/test/consensus/Consensus-ValidatorResignation.sol

@@ -36,7 +36,6 @@ contract ConsensusTest is Base {
 		consensus.resignValidator();
 		vm.stopPrank();
 
-
 		assertEq(consensus.registeredValidatorsCount(), 1);
 		validator = consensus.getValidator(addr);
 		assertEq(validator.addr, addr);

contracts/test/consensus/Consensus-Vote.sol

@@ -73,6 +73,11 @@ contract ConsensusTest is Test {
 		assertEq(voterAddr.balance, 90 ether);
 	}
 
+	function test_vote_revert_if_caller_is_owner() public {
+		vm.expectRevert("Caller is the contract owner");
+		consensus.vote(address(1));
+	}
+
 	function test_vote_allow_self_vote() public {
 		// Register validator
 		address addr = address(1);
@@ -130,8 +135,11 @@ contract ConsensusTest is Test {
 	}
 
 	function test_vote_prevent_for_unregistered_validator() public {
+		address addr = address(1);
+
+		vm.startPrank(addr);
 		vm.expectRevert("Must vote for validator");
-		consensus.vote(address(1));
+		consensus.vote(addr);
 	}
 
 	function test_vote_prevent_for_resigned_validator() public {
@@ -184,7 +192,7 @@ contract ConsensusTest is Test {
 		assertEq(voterAddr.balance, 90 ether);
 	}
 
-		function test_unvote_and_vote_in_different_blocks() public {
+	function test_unvote_and_vote_in_different_blocks() public {
 		// Register validator
 		address addr = address(1);
 		registerValidator(addr);