Post-audit fixes for beefy light client (paritytech#877)

* INFO-4(code duplication) * INFO-3(Redundant & operation in set and isSet) * INFO-2(Missing conditions) * INFO-1(refactor custom error) * More refactors * INFO-5(Add comments&respect MAX_SEED_LOOKAHEAD) * INFO-6 & INFO-9(optimization countSetBits) * Fix for e2e * More refactor & sync snapshot with latest main * Resolve comments * Revert Info-9(countSetBits) * Revert RANDAO_COMMIT_DELAY for quick test * Revert some refactoring * Fix! * Fix non handover case * Ensure use same initial bitfield for createFinalBitfield * More tests * Fix beefy relayer * More tests & change beefy log level * Fix missing check and more tests * More comments * Add more assert * More tests * Some polish * Improve test remove storage variable * Reset RANDAO params
Snowfork · Jul 12, 2023 · b726f59 · b726f59
1 parent 7bdcd82
commit b726f59
Show file tree

Hide file tree

Showing 10 changed files with 298 additions and 154 deletions.
diff --git a/core/packages/contracts/.gas-snapshot b/core/packages/contracts/.gas-snapshot
@@ -41,4 +41,4 @@ TokenVaultTest:testTokenTransferFailedInsufficientAllowance() (gas: 69957)
 TokenVaultTest:testWithdrawSuccessful() (gas: 121381)
 UpgradeProxyTest:testUpgrade() (gas: 35658)
 UpgradeProxyTest:testUpgradeFail() (gas: 20802)
-UpgradeProxyTest:testUpgradeFailBadOrigin() (gas: 11326)
+UpgradeProxyTest:testUpgradeFailBadOrigin() (gas: 11326)
diff --git a/core/packages/contracts/scripts/ffiWrapper.ts b/core/packages/contracts/scripts/ffiWrapper.ts
@@ -17,22 +17,17 @@ const mmrRoot = fixtureData.params.commitment.payload[0].data
 const mmrLeaf: BeefyClient.MMRLeafStruct = fixtureData.params.leaf
 const leafProofOrder = fixtureData.params.leafProofOrder
 
-const subsetSize = validatorSetSize - Math.floor((validatorSetSize - 1) / 3)
+const absentSubsetSize = Math.floor((validatorSetSize - 1) / 3)
+const subsetSize = validatorSetSize - absentSubsetSize
 const subset = createRandomSubset(validatorSetSize, subsetSize)
+const absentSubset = createRandomSubset(validatorSetSize, absentSubsetSize)
 let validatorSet: ValidatorSet
 
 if (command == "GenerateInitialSet") {
     process.stdout.write(
         `${encoder.encode(
-            [
-                "uint32",
-                "uint32",
-                "uint32",
-                "uint256[]",
-                "bytes32",
-                "bytes32"
-            ],
-            [blockNumber, validatorSetID, validatorSetSize, subset, commitHash, mmrRoot]
+            ["uint32", "uint32", "uint32", "uint256[]", "uint256[]", "bytes32", "bytes32"],
+            [blockNumber, validatorSetID, validatorSetSize, subset, absentSubset, commitHash, mmrRoot]
         )}`
     )
 } else if (command == "GenerateProofs") {
@@ -62,13 +57,7 @@ if (command == "GenerateInitialSet") {
                 "tuple(uint8 version,uint32 parentNumber,bytes32 parentHash,uint64 nextAuthoritySetID,uint32 nextAuthoritySetLen,bytes32 nextAuthoritySetRoot,bytes32 parachainHeadsRoot)",
                 "uint256",
             ],
-            [
-                validatorSet.root,
-                validatorFinalProofs,
-                mmrLeafProofs,
-                mmrLeaf,
-                leafProofOrder,
-            ]
+            [validatorSet.root, validatorFinalProofs, mmrLeafProofs, mmrLeaf, leafProofOrder]
         )}`
     )
 }
diff --git a/core/packages/contracts/src/BeefyClient.sol b/core/packages/contracts/src/BeefyClient.sol
@@ -74,7 +74,7 @@ contract BeefyClient is Ownable {
      * @param r the x component on the secp256k1 curve
      * @param s the challenge solution
      * @param index index of the validator address in the merkle tree
-     * @param addr validator address
+     * @param account validator address
      * @param proof merkle proof for the validator
      */
     struct ValidatorProof {
@@ -170,20 +170,20 @@ contract BeefyClient is Ownable {
     uint256 public immutable randaoCommitExpiration;
 
     /* Errors */
-
     error InvalidCommitment();
     error StaleCommitment();
     error InvalidValidatorProof();
     error InvalidSignature();
     error NotEnoughClaims();
     error InvalidMMRLeaf();
     error InvalidMMRLeafProof();
-    error InvalidTask();
     error InvalidBitfield();
     error WaitPeriodNotOver();
     error TicketExpired();
     error PrevRandaoAlreadyCaptured();
     error PrevRandaoNotCaptured();
+    error InvalidBitfieldLength();
+    error InvalidTicket();
     error InvalidMMRRootLength();
     error NoMMRRootInCommitment();
 
@@ -214,7 +214,6 @@ contract BeefyClient is Ownable {
      */
     function submitInitial(Commitment calldata commitment, uint256[] calldata bitfield, ValidatorProof calldata proof)
         external
-        payable
     {
         doSubmitInitial(currentValidatorSet, commitment, bitfield, proof);
     }
@@ -229,7 +228,7 @@ contract BeefyClient is Ownable {
         Commitment calldata commitment,
         uint256[] calldata bitfield,
         ValidatorProof calldata proof
-    ) external payable {
+    ) external {
         doSubmitInitial(nextValidatorSet, commitment, bitfield, proof);
     }
 
@@ -298,26 +297,13 @@ contract BeefyClient is Ownable {
     function submitFinal(Commitment calldata commitment, uint256[] calldata bitfield, ValidatorProof[] calldata proofs)
         public
     {
-        bytes32 commitmentHash = keccak256(encodeCommitment(commitment));
-        bytes32 ticketID = createTicketID(msg.sender, commitmentHash);
-        Ticket storage ticket = tickets[ticketID];
+        (bytes32 commitmentHash, bytes32 ticketID) = validate(commitment, bitfield);
 
-        if (ticket.prevRandao == 0) {
-            revert PrevRandaoNotCaptured();
-        }
+        Ticket storage ticket = tickets[ticketID];
 
-        if (commitment.validatorSetID != currentValidatorSet.id && commitment.validatorSetID != nextValidatorSet.id) {
+        if (commitment.validatorSetID != currentValidatorSet.id) {
             revert InvalidCommitment();
         }
-
-        if (commitment.blockNumber <= latestBeefyBlock) {
-            revert StaleCommitment();
-        }
-
-        if (ticket.bitfieldHash != keccak256(abi.encodePacked(bitfield))) {
-            revert InvalidBitfield();
-        }
-
         verifyCommitment(commitmentHash, bitfield, currentValidatorSet, ticket, proofs);
 
         bytes32 newMMRRoot = getFirstMMRRoot(commitment);
@@ -331,9 +317,11 @@ contract BeefyClient is Ownable {
     /**
      * @dev Submit a commitment and leaf for final verification
      * @param commitment contains the full commitment that was used for the commitmentHash
+     * @param bitfield claiming which validators have signed the commitment
      * @param proofs a struct containing the data needed to verify all validator signatures
      * @param leaf an MMR leaf provable using the MMR root in the commitment payload
      * @param leafProof an MMR leaf proof
+     * @param leafProofOrder a bitfield describing the order of each item (left vs right)
      */
     function submitFinalWithHandover(
         Commitment calldata commitment,
@@ -343,32 +331,18 @@ contract BeefyClient is Ownable {
         bytes32[] calldata leafProof,
         uint256 leafProofOrder
     ) public {
-        bytes32 commitmentHash = keccak256(encodeCommitment(commitment));
-        bytes32 ticketID = createTicketID(msg.sender, commitmentHash);
-        Ticket storage ticket = tickets[ticketID];
+        (bytes32 commitmentHash, bytes32 ticketID) = validate(commitment, bitfield);
 
-        if (ticket.prevRandao == 0) {
-            revert PrevRandaoNotCaptured();
-        }
+        Ticket storage ticket = tickets[ticketID];
 
         if (commitment.validatorSetID != nextValidatorSet.id) {
             revert InvalidCommitment();
         }
-
-        if (commitment.blockNumber <= latestBeefyBlock) {
-            revert StaleCommitment();
-        }
+        verifyCommitment(commitmentHash, bitfield, nextValidatorSet, ticket, proofs);
 
         if (leaf.nextAuthoritySetID != nextValidatorSet.id + 1) {
             revert InvalidMMRLeaf();
         }
-
-        if (ticket.bitfieldHash != keccak256(abi.encodePacked(bitfield))) {
-            revert InvalidBitfield();
-        }
-
-        verifyCommitment(commitmentHash, bitfield, nextValidatorSet, ticket, proofs);
-
         bytes32 newMMRRoot = getFirstMMRRoot(commitment);
         bool leafIsValid =
             MMRProof.verifyLeafProof(newMMRRoot, keccak256(encodeMMRLeaf(leaf)), leafProof, leafProofOrder);
@@ -392,6 +366,7 @@ contract BeefyClient is Ownable {
      * @dev Verify that the supplied MMR leaf is included in the latest verified MMR root.
      * @param leafHash contains the merkle leaf to be verified
      * @param proof contains simplified mmr proof
+     * @param proofOrder a bitfield describing the order of each item (left vs right)
      */
     function verifyMMRLeafProof(bytes32 leafHash, bytes32[] calldata proof, uint256 proofOrder)
         external
@@ -401,6 +376,41 @@ contract BeefyClient is Ownable {
         return MMRProof.verifyLeafProof(latestMMRRoot, leafHash, proof, proofOrder);
     }
 
+    /**
+     * @dev Helper to create an initial validator bitfield.
+     * @param bitsToSet contains indexes of all signed validators, should be deduplicated
+     * @param length of validator set
+     */
+    function createInitialBitfield(uint256[] calldata bitsToSet, uint256 length)
+        external
+        pure
+        returns (uint256[] memory)
+    {
+        if (length < bitsToSet.length) {
+            revert InvalidBitfieldLength();
+        }
+        return Bitfield.createBitfield(bitsToSet, length);
+    }
+
+    /**
+     * @dev Helper to create a final bitfield, with subsampled validator selections
+     * @param commitmentHash contains the commitmentHash signed by the validators
+     * @param bitfield claiming which validators have signed the commitment
+     */
+    function createFinalBitfield(bytes32 commitmentHash, uint256[] calldata bitfield)
+        external
+        view
+        returns (uint256[] memory)
+    {
+        Ticket storage ticket = tickets[createTicketID(msg.sender, commitmentHash)];
+        if (ticket.bitfieldHash != keccak256(abi.encodePacked(bitfield))) {
+            revert InvalidBitfield();
+        }
+        return Bitfield.subsample(
+            ticket.prevRandao, bitfield, minimumSignatureThreshold(ticket.validatorSetLen), ticket.validatorSetLen
+        );
+    }
+
     /* Private Functions */
 
     // Creates a unique ticket ID for a new interactive prover-verifier session
@@ -558,41 +568,44 @@ contract BeefyClient is Ownable {
      * the previous version did not require any sorting.
      *
      * @dev Checks if a validators address is a member of the merkle tree
-     * @param addr The address of the validator to check
+     * @param account The address of the validator to check
      * @param proof Merkle proof required for validation of the address
      * @return true if the validator is in the set
      */
-    function isValidatorInSet(ValidatorSet memory vset, address addr, uint256 index, bytes32[] calldata proof)
+    function isValidatorInSet(ValidatorSet memory vset, address account, uint256 index, bytes32[] calldata proof)
         internal
         pure
         returns (bool)
     {
-        bytes32 hashedLeaf = keccak256(abi.encodePacked(addr));
+        bytes32 hashedLeaf = keccak256(abi.encodePacked(account));
         return MerkleProof.verify(vset.root, hashedLeaf, index, vset.length, proof);
     }
 
-    /**
-     * @dev Helper to create an initial validator bitfield.
-     */
-    function createInitialBitfield(uint256[] calldata bitsToSet, uint256 length)
-        external
-        pure
-        returns (uint256[] memory)
-    {
-        return Bitfield.createBitfield(bitsToSet, length);
-    }
-
-    /**
-     * @dev Helper to create a final bitfield, with subsampled validator selections
-     */
-    function createFinalBitfield(bytes32 commitmentHash, uint256[] calldata bitfield)
-        external
+    // Basic checks for commitment
+    function validate(Commitment calldata commitment, uint256[] calldata bitfield)
+        internal
         view
-        returns (uint256[] memory)
+        returns (bytes32, bytes32)
     {
-        Ticket storage ticket = tickets[createTicketID(msg.sender, commitmentHash)];
-        return Bitfield.subsample(
-            ticket.prevRandao, bitfield, minimumSignatureThreshold(ticket.validatorSetLen), ticket.validatorSetLen
-        );
+        bytes32 commitmentHash = keccak256(encodeCommitment(commitment));
+        bytes32 ticketID = createTicketID(msg.sender, commitmentHash);
+        Ticket storage ticket = tickets[ticketID];
+
+        if (ticket.blockNumber == 0) {
+            revert InvalidTicket();
+        }
+
+        if (ticket.prevRandao == 0) {
+            revert PrevRandaoNotCaptured();
+        }
+
+        if (commitment.blockNumber <= latestBeefyBlock) {
+            revert StaleCommitment();
+        }
+
+        if (ticket.bitfieldHash != keccak256(abi.encodePacked(bitfield))) {
+            revert InvalidBitfield();
+        }
+        return (commitmentHash, ticketID);
     }
 }
diff --git a/core/packages/contracts/src/utils/Bitfield.sol b/core/packages/contracts/src/utils/Bitfield.sol
@@ -5,10 +5,13 @@ pragma solidity 0.8.20;
 import "./Bits.sol";
 
 library Bitfield {
+    using Bits for uint256;
+
     /**
      * @dev Constants used to efficiently calculate the hamming weight of a bitfield. See
      * https://en.wikipedia.org/wiki/Hamming_weight#Efficient_implementation for an explanation of those constants.
      */
+
     uint256 internal constant M1 = 0x5555555555555555555555555555555555555555555555555555555555555555;
     uint256 internal constant M2 = 0x3333333333333333333333333333333333333333333333333333333333333333;
     uint256 internal constant M4 = 0x0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f;
@@ -62,7 +65,7 @@ library Bitfield {
      * @dev Helper to create a bitfield.
      */
     function createBitfield(uint256[] calldata bitsToSet, uint256 length)
-        public
+        internal
         pure
         returns (uint256[] memory bitfield)
     {
@@ -71,8 +74,11 @@ library Bitfield {
 
         bitfield = new uint256[](arrayLength);
 
-        for (uint256 i = 0; i < bitsToSet.length; i++) {
+        for (uint256 i = 0; i < bitsToSet.length;) {
             set(bitfield, bitsToSet[i]);
+            unchecked {
+                ++i;
+            }
         }
 
         return bitfield;
@@ -104,20 +110,17 @@ library Bitfield {
 
     function isSet(uint256[] memory self, uint256 index) internal pure returns (bool) {
         uint256 element = index >> 8;
-        uint8 within = uint8(index & 0xFF);
-        return Bits.bit(self[element], within) == 1;
+        return self[element].bit(uint8(index)) == 1;
     }
 
     function set(uint256[] memory self, uint256 index) internal pure {
         uint256 element = index >> 8;
-        uint8 within = uint8(index & 0xFF);
-        self[element] = Bits.setBit(self[element], within);
+        self[element] = self[element].setBit(uint8(index));
     }
 
     function unset(uint256[] memory self, uint256 index) internal pure {
         uint256 element = index >> 8;
-        uint8 within = uint8(index & 0xFF);
-        self[element] = Bits.clearBit(self[element], within);
+        self[element] = self[element].clearBit(uint8(index));
     }
 
     function makeIndex(uint256 seed, uint256 iteration, uint256 length) internal pure returns (uint256 index) {