feat(contracts): collect validator registration fee (#970)

* Add comment

* Rename to validatorsCount

* Remove _validatorsCount

* Add _activeValidators list

* Use _activeValidators for measurement

* Add activeValidatorsCount

* Shuffle without parameters

* Use _activeValidators in calculateRoundValidators

* Add fee

* Collect and return fee

* Fix add tests

* Test registration

* Check balance

* Check resignation

* Test fee

* Check refund

* Add tests

* Improve _isActiveValidator

* Fix reentrancy

* Fee optimisation

* Fee in initializer

* Test fee

* Improve packing

* Set fee to uint128

* Update contract

* Fix deployer

* Set milestone

* Support validatorRegistrationFee option

* Calculate total supply

* Set validatorRegistrationFee

* Generate network

* style: resolve style guide violations

* regenerate networks

* fixes

* style: resolve style guide violations

---------

Co-authored-by: oXtxNt9U <120286271+oXtxNt9U@users.noreply.github.com>

Author: sebastijankuzner

contracts/src/consensus/ConsensusV1.sol

@@ -30,8 +30,9 @@ import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Own
 
 contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
     struct ValidatorData {
-        uint256 votersCount;
         uint256 voteBalance;
+        uint128 fee;
+        uint64 votersCount;
         bool isResigned;
         bytes blsPublicKey; // 96 bits
     }
@@ -63,16 +64,15 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
         address validator;
     }
 
+    event FeeUpdated(uint256 fee);
     event ValidatorRegistered(address addr, bytes blsPublicKey);
-
     event ValidatorUpdated(address addr, bytes blsPublicKey);
-
     event ValidatorResigned(address addr);
-
     event Voted(address voter, address validator);
-
     event Unvoted(address voter, address validator);
 
+    error InvalidFee();
+    error RefundFailed();
     error CallerIsNotValidator();
     error ValidatorNotRegistered();
     error ValidatorAlreadyRegistered();
@@ -96,8 +96,9 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
     mapping(address => ValidatorData) private _validatorsData;
     mapping(address => bool) private _hasValidator;
     mapping(bytes32 => bool) private _blsPublicKeys;
-    address[] private _validators;
-    uint256 private _validatorsCount; // Default 0
+    address[] private _validators; // All registered validators including resigned
+    address[] private _activeValidators; // Has valid BLS public key and is not resigned
+    mapping(address => uint256) private _activeValidatorIndex; // Points to index in _activeValidators array
     uint256 private _resignedValidatorsCount; // Default 0
 
     mapping(address => Vote) private _voters;
@@ -110,19 +111,26 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
     address private _roundValidatorsHead; // Default address(0)
     uint256 private _roundValidatorsCount; // Default 0
     uint256 private _minValidators; // Default 1
+    uint128 private _fee; // Validator registration fee: Default 0
 
     RoundValidator[][] private _rounds;
 
     // Initializers
-    function initialize() public initializer {
+    function initialize(uint128 registrationFee) public initializer {
         __Ownable_init(msg.sender);
         _minValidators = 1;
+        _fee = registrationFee;
     }
 
     // Overrides
     function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}
 
     // External functions
+    function setFee(uint128 registrationFee) external onlyOwner {
+        _fee = registrationFee;
+        emit FeeUpdated(registrationFee);
+    }
+
     function addValidator(address addr, bytes calldata blsPublicKey, bool isResigned) external onlyOwner {
         if (_rounds.length > 0) {
             revert ImportIsNotAllowed();
@@ -142,9 +150,8 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
         }
 
         ValidatorData memory validator =
-            ValidatorData({votersCount: 0, voteBalance: 0, isResigned: isResigned, blsPublicKey: blsPublicKey});
+            ValidatorData({votersCount: 0, voteBalance: 0, fee: 0, isResigned: isResigned, blsPublicKey: blsPublicKey});
 
-        _validatorsCount++;
         _hasValidator[addr] = true;
         _validatorsData[addr] = validator;
         _validators.push(addr);
@@ -153,6 +160,10 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
             _resignedValidatorsCount++;
         }
 
+        if (_canBecomeActiveValidator(addr)) {
+            _addActiveValidator(addr);
+        }
+
         emit ValidatorRegistered(addr, blsPublicKey);
     }
 
@@ -197,20 +208,29 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
         emit Voted(voter, validator);
     }
 
-    function registerValidator(bytes calldata blsPublicKey) external {
+    function registerValidator(bytes calldata blsPublicKey) external payable {
+        if (msg.value != _fee) {
+            revert InvalidFee();
+        }
+
         if (_hasValidator[msg.sender]) {
             revert ValidatorAlreadyRegistered();
         }
 
         _verifyAndRegisterBlsPublicKey(blsPublicKey);
 
-        ValidatorData memory validator =
-            ValidatorData({votersCount: 0, voteBalance: 0, isResigned: false, blsPublicKey: blsPublicKey});
+        ValidatorData memory validator = ValidatorData({
+            votersCount: 0,
+            voteBalance: 0,
+            fee: uint128(msg.value),
+            isResigned: false,
+            blsPublicKey: blsPublicKey
+        });
 
-        _validatorsCount++;
         _hasValidator[msg.sender] = true;
         _validatorsData[msg.sender] = validator;
         _validators.push(msg.sender);
+        _addActiveValidator(msg.sender);
 
         emit ValidatorRegistered(msg.sender, blsPublicKey);
     }
@@ -224,6 +244,10 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
 
         _validatorsData[msg.sender].blsPublicKey = blsPublicKey;
 
+        if (_canBecomeActiveValidator(msg.sender) && !_isActiveValidator(msg.sender)) {
+            _addActiveValidator(msg.sender);
+        }
+
         emit ValidatorUpdated(msg.sender, blsPublicKey);
     }
 
@@ -237,13 +261,25 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
             revert ValidatorAlreadyResigned();
         }
 
-        if (_validatorsCount - _resignedValidatorsCount <= _minValidators) {
+        if (_activeValidators.length <= _minValidators) {
             revert BellowMinValidators();
         }
 
         validator.isResigned = true;
         _resignedValidatorsCount += 1;
 
+        _removeActiveValidator(msg.sender);
+
+        // Refund the registration fee to the validator
+        if (validator.fee > 0) {
+            uint256 refundFee = validator.fee;
+            validator.fee = 0;
+            (bool success,) = payable(msg.sender).call{value: refundFee}("");
+            if (!success) {
+                revert RefundFailed();
+            }
+        }
+
         emit ValidatorResigned(msg.sender);
     }
 
@@ -305,18 +341,18 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
 
         _minValidators = n;
 
-        _shuffle(_validators);
+        _shuffle();
         _deleteRoundValidators();
 
         _roundValidatorsHead = address(0);
-        uint8 top = uint8(_clamp(n, 0, _validatorsCount - _resignedValidatorsCount));
+        uint8 top = uint8(_clamp(n, 0, _activeValidators.length));
 
         if (top == 0) {
             revert NoActiveValidators();
         }
 
-        for (uint256 i = 0; i < _validators.length; i++) {
-            address addr = _validators[i];
+        for (uint256 i = 0; i < _activeValidators.length; i++) {
+            address addr = _activeValidators[i];
 
             ValidatorData storage data = _validatorsData[addr];
 
@@ -375,8 +411,16 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
         return 1;
     }
 
-    function registeredValidatorsCount() external view returns (uint256) {
-        return _validatorsCount;
+    function fee() external view returns (uint256) {
+        return _fee;
+    }
+
+    function validatorsCount() external view returns (uint256) {
+        return _validators.length;
+    }
+
+    function activeValidatorsCount() external view returns (uint256) {
+        return _activeValidators.length;
     }
 
     function resignedValidatorsCount() external view returns (uint256) {
@@ -475,8 +519,8 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
     }
 
     // Internal functions
-    function _shuffle(address[] storage array) internal {
-        uint256 n = array.length;
+    function _shuffle() internal {
+        uint256 n = _activeValidators.length;
         if (n == 0) {
             return;
         }
@@ -485,10 +529,35 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
             // Get a random index between 0 and i (inclusive)
             uint256 j = uint256(keccak256(abi.encodePacked(block.timestamp, i))) % (i + 1);
 
+            if (i == j) {
+                continue; // No need to swap if indices are the same
+            }
+
+            /* Swap example
+            i = 0; j = 2;
+
+    		Initial state
+            A B C
+            A:0 B:1 C:2
+
+            Array SWAP
+            C B A
+            A:0 B:1 C:2
+
+            Index SWAP
+            C B A
+            A:2 B:1 C:0
+            */
+
             // Swap elements at index i and j
-            address temp = array[i];
-            array[i] = array[j];
-            array[j] = temp;
+            address addrA = _activeValidators[i];
+            address addrB = _activeValidators[j];
+
+            _activeValidators[i] = _activeValidators[j];
+            _activeValidators[j] = addrA;
+
+            _activeValidatorIndex[addrA] = j;
+            _activeValidatorIndex[addrB] = i;
         }
     }
 
@@ -574,6 +643,47 @@ contract ConsensusV1 is UUPSUpgradeable, OwnableUpgradeable {
         _roundValidatorsCount++;
     }
 
+    function _addActiveValidator(address addr) internal {
+        _activeValidators.push(addr);
+        _activeValidatorIndex[addr] = _activeValidators.length - 1;
+    }
+
+    function _removeActiveValidator(address addr) internal {
+        if (!_isActiveValidator(addr)) {
+            return;
+        }
+
+        uint256 index = _activeValidatorIndex[addr];
+        uint256 lastIndex = _activeValidators.length - 1;
+
+        // Copy last address to the index of the removed address. This is not swap. Last validator occurs 2 times in the array.
+        if (index != lastIndex) {
+            address lastValidator = _activeValidators[lastIndex];
+            _activeValidators[index] = lastValidator;
+            _activeValidatorIndex[lastValidator] = index;
+        }
+
+        // Remove last address
+        _activeValidators.pop();
+        delete _activeValidatorIndex[addr];
+    }
+
+    function _canBecomeActiveValidator(address addr) internal view returns (bool) {
+        ValidatorData storage data = _validatorsData[addr];
+        return !data.isResigned && data.blsPublicKey.length != 0;
+    }
+
+    function _isActiveValidator(address addr) internal view returns (bool) {
+        uint256 index = _activeValidatorIndex[addr];
+        // Support for empty array
+        if (index >= _activeValidators.length) {
+            return false;
+        }
+
+        // Check if the address at the index matches the address. Required for the case when index is 0.
+        return _activeValidators[index] == addr;
+    }
+
     function _unvote() internal returns (address) {
         Vote storage voter = _voters[msg.sender];
         if (voter.validator == address(0)) {

contracts/test/consensus/Base.sol

@@ -12,7 +12,7 @@ contract Base is Test {
     function test() public {}
 
     function setUp() public {
-        bytes memory data = abi.encode(ConsensusV1.initialize.selector);
+        bytes memory data = abi.encode(ConsensusV1.initialize.selector, 0);
         address proxy = address(new ERC1967Proxy(address(new ConsensusV1()), data));
         consensus = ConsensusV1(proxy);
     }

contracts/test/consensus/Consensus-CalculateTop.sol

@@ -94,7 +94,7 @@ contract ConsensusTest is Base {
 
     function test_consensus_sortedValidators_sameVoteCounts() public {
         vm.pauseGasMetering();
-        assertEq(consensus.registeredValidatorsCount(), 0);
+        assertEq(consensus.validatorsCount(), 0);
 
         uint256 n = 55;
         uint256 balance = 50;
@@ -140,7 +140,7 @@ contract ConsensusTest is Base {
 
     function test_consensus_200_topValidators() public {
         vm.pauseGasMetering();
-        assertEq(consensus.registeredValidatorsCount(), 0);
+        assertEq(consensus.validatorsCount(), 0);
 
         address highest = address(0);
         uint256 highestBalance = 0;

contracts/test/consensus/Consensus-Fee.sol

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: GNU GENERAL PUBLIC LICENSE
+pragma solidity ^0.8.13;
+
+import {ConsensusV1} from "@contracts/consensus/ConsensusV1.sol";
+import {Base} from "./Base.sol";
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+
+contract ConsensusTest is Base {
+    function test_default_fee() public view {
+        assertEq(consensus.fee(), 0);
+    }
+
+    function test_default_fee_custom() public {
+        uint256 initialFee = 10;
+        bytes memory data = abi.encodeWithSelector(ConsensusV1.initialize.selector, initialFee);
+        address proxy = address(new ERC1967Proxy(address(new ConsensusV1()), data));
+        ConsensusV1 consensusCustom = ConsensusV1(proxy);
+
+        assertEq(consensusCustom.fee(), initialFee);
+    }
+
+    function test_default_fee_should_be_adjustable() public {
+        assertEq(consensus.fee(), 0);
+
+        uint128 newFee = 1000;
+        vm.expectEmit(address(consensus));
+        emit ConsensusV1.FeeUpdated(newFee);
+        consensus.setFee(newFee);
+
+        assertEq(consensus.fee(), newFee);
+    }
+}

contracts/test/consensus/Consensus-GetAllValidators.sol

@@ -8,7 +8,7 @@ import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.s
 contract ConsensusTest is Base {
     function test_200_validators() public {
         vm.pauseGasMetering();
-        assertEq(consensus.registeredValidatorsCount(), 0);
+        assertEq(consensus.validatorsCount(), 0);
 
         uint256 n = 200;
         for (uint256 i = 0; i < n; i++) {