♻️ Make MerkleProofVerification Module-Friendly (#231)

### 🕓 Changelog This PR refactors the `MerkleProofVerification` contract to make it module-friendly and ready for the breaking `0.4.0` release. --------- Signed-off-by: Pascal Marco Caversaccio <pascal.caversaccio@hotmail.ch>
pcaversaccio · Apr 10, 2024 · d5a9c8d · d5a9c8d
1 parent dd58215
commit d5a9c8d
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 29 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -582,18 +582,18 @@ MathTest:testWadExp() (gas: 34590)
 MathTest:testWadExpOverflow() (gas: 11387)
 MathTest:testWadLn() (gas: 31609)
 MathTest:testWadLnNegativeValues() (gas: 11360)
-MerkleProofVerificationTest:testFuzzMultiProofVerifySingleLeaf(bytes32[],uint256) (runs: 256, μ: 1649985414, ~: 1649982220)
-MerkleProofVerificationTest:testFuzzVerify(bytes32[],uint256) (runs: 256, μ: 135979115, ~: 135975891)
-MerkleProofVerificationTest:testFuzzVerifyMultiProofMultipleLeaves(bool,bool,bool) (runs: 256, μ: 412477124, ~: 412477119)
-MerkleProofVerificationTest:testInvalidMerkleMultiProof() (gas: 412480569)
-MerkleProofVerificationTest:testInvalidMerkleProof() (gas: 33971264)
-MerkleProofVerificationTest:testInvalidMerkleProofLength() (gas: 33973431)
-MerkleProofVerificationTest:testInvalidMultiProof() (gas: 909641334)
-MerkleProofVerificationTest:testMaliciousMultiProofVerify() (gas: 303219039)
-MerkleProofVerificationTest:testMultiProofEdgeCase1() (gas: 412462703)
-MerkleProofVerificationTest:testMultiProofEdgeCase2() (gas: 412462855)
-MerkleProofVerificationTest:testMultiProofVerify() (gas: 412484548)
-MerkleProofVerificationTest:testVerify() (gas: 67941139)
+MerkleProofVerificationTest:testFuzzMultiProofVerifySingleLeaf(bytes32[],uint256) (runs: 256, μ: 262136, ~: 258922)
+MerkleProofVerificationTest:testFuzzVerify(bytes32[],uint256) (runs: 256, μ: 174897, ~: 171653)
+MerkleProofVerificationTest:testFuzzVerifyMultiProofMultipleLeaves(bool,bool,bool) (runs: 256, μ: 46263, ~: 46254)
+MerkleProofVerificationTest:testInvalidMerkleMultiProof() (gas: 49759)
+MerkleProofVerificationTest:testInvalidMerkleProof() (gas: 20183)
+MerkleProofVerificationTest:testInvalidMerkleProofLength() (gas: 22362)
+MerkleProofVerificationTest:testInvalidMultiProof() (gas: 91637)
+MerkleProofVerificationTest:testMaliciousMultiProofVerify() (gas: 36820)
+MerkleProofVerificationTest:testMultiProofEdgeCase1() (gas: 31820)
+MerkleProofVerificationTest:testMultiProofEdgeCase2() (gas: 31972)
+MerkleProofVerificationTest:testMultiProofVerify() (gas: 53775)
+MerkleProofVerificationTest:testVerify() (gas: 38997)
 MessageHashUtilsTest:testEthSignedMessageHash() (gas: 8614)
 MessageHashUtilsTest:testFuzzEthSignedMessageHash(string) (runs: 256, μ: 9321, ~: 9315)
 MessageHashUtilsTest:testFuzzToDataWithIntendedValidatorHash(address,bytes) (runs: 256, μ: 10144, ~: 10129)

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,7 @@
   - [`SignatureChecker`](https://github.com/pcaversaccio/snekmate/blob/v0.1.0/src/snekmate/utils/SignatureChecker.vy): Make `SignatureChecker` module-friendly. ([#228](https://github.com/pcaversaccio/snekmate/pull/228))
   - [`EIP712DomainSeparator`](https://github.com/pcaversaccio/snekmate/blob/v0.1.0/src/snekmate/utils/EIP712DomainSeparator.vy): Make `EIP712DomainSeparator` module-friendly. ([#229](https://github.com/pcaversaccio/snekmate/pull/229))
   - [`Math`](https://github.com/pcaversaccio/snekmate/blob/v0.1.0/src/snekmate/utils/Math.vy): Make `Math` module-friendly. ([#230](https://github.com/pcaversaccio/snekmate/pull/230))
+  - [`MerkleProofVerification`](https://github.com/pcaversaccio/snekmate/blob/v0.1.0/src/snekmate/utils/MerkleProofVerification.vy): Make `MerkleProofVerification` module-friendly. ([#231](https://github.com/pcaversaccio/snekmate/pull/231))
 - **Vyper Contract Deployer**
   - [`VyperDeployer`](https://github.com/pcaversaccio/snekmate/blob/v0.1.0/lib/utils/VyperDeployer.sol): Improve error message in the event of a Vyper compilation error. ([#219](https://github.com/pcaversaccio/snekmate/pull/219))
 

diff --git a/README.md b/README.md
@@ -73,7 +73,8 @@ src
             ├── MessageHashUtilsMock — "MessageHashUtils Module Reference Implementation"
             ├── SignatureCheckerMock — "SignatureChecker Module Reference Implementation"
             ├── EIP712DomainSeparatorMock — "EIP712DomainSeparator Module Reference Implementation"
-            └── MathMock — "Math Module Reference Implementation"
+            ├── MathMock — "Math Module Reference Implementation"
+            └── MerkleProofVerificationMock — "MerkleProofVerification Module Reference Implementation"
 ```
 
 ## 🎛 Installation

diff --git a/src/snekmate/utils/MerkleProofVerification.vy b/src/snekmate/utils/MerkleProofVerification.vy
@@ -18,7 +18,7 @@
         reinterpreted as a leaf value. OpenZeppelin's JavaScript library
         `merkle-tree` (https://github.com/OpenZeppelin/merkle-tree)
         generates Merkle trees that are safe against this attack out of
-        the box. You will find a quick start guide in the README.
+        the box. You will find a quick start guide in the `README`.
         OpenZeppelin provides some good examples of how to construct
         Merkle tree proofs correctly:
         https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/test/utils/cryptography/MerkleProof.test.js.
@@ -44,9 +44,9 @@ def __init__():
     pass
 
 
-@external
+@internal
 @pure
-def verify(proof: DynArray[bytes32, max_value(uint16)], root: bytes32, leaf: bytes32) -> bool:
+def _verify(proof: DynArray[bytes32, max_value(uint8)], root: bytes32, leaf: bytes32) -> bool:
     """
     @dev Returns `True` if it can be proved that a `leaf` is
          part of a Merkle tree defined by `root`.
@@ -63,10 +63,10 @@ def verify(proof: DynArray[bytes32, max_value(uint16)], root: bytes32, leaf: byt
     return self._process_proof(proof, leaf) == root
 
 
-@external
+@internal
 @pure
-def multi_proof_verify(proof: DynArray[bytes32, max_value(uint16)], proof_flags: DynArray[bool, max_value(uint16)],
-                       root: bytes32, leaves: DynArray[bytes32, max_value(uint16)]) -> bool:
+def _multi_proof_verify(proof: DynArray[bytes32, max_value(uint8)], proof_flags: DynArray[bool, max_value(uint8)],
+                        root: bytes32, leaves: DynArray[bytes32, max_value(uint8)]) -> bool:
     """
     @dev Returns `True` if it can be simultaneously proved that
          `leaves` are part of a Merkle tree defined by `root`
@@ -90,7 +90,7 @@ def multi_proof_verify(proof: DynArray[bytes32, max_value(uint16)], proof_flags:
 
 @internal
 @pure
-def _process_proof(proof: DynArray[bytes32, max_value(uint16)], leaf: bytes32) -> bytes32:
+def _process_proof(proof: DynArray[bytes32, max_value(uint8)], leaf: bytes32) -> bytes32:
     """
     @dev Returns the recovered hash obtained by traversing
          a Merkle tree from `leaf` using `proof`.
@@ -111,8 +111,8 @@ def _process_proof(proof: DynArray[bytes32, max_value(uint16)], leaf: bytes32) -
 
 @internal
 @pure
-def _process_multi_proof(proof: DynArray[bytes32, max_value(uint16)], proof_flags: DynArray[bool, max_value(uint16)],
-                         leaves: DynArray[bytes32, max_value(uint16)]) -> bytes32:
+def _process_multi_proof(proof: DynArray[bytes32, max_value(uint8)], proof_flags: DynArray[bool, max_value(uint8)],
+                         leaves: DynArray[bytes32, max_value(uint8)]) -> bytes32:
     """
     @dev Returns the recovered hash obtained by traversing
          a Merkle tree from `leaves` using `proof` and a
@@ -150,12 +150,12 @@ def _process_multi_proof(proof: DynArray[bytes32, max_value(uint16)], proof_flag
 
     # Checks the validity of the proof. We do not check for an
     # overflow (nor underflow) as `leaves_length`, `proof`, and
-    # `total_hashes` are bounded by the value `max_value(uint16)`
+    # `total_hashes` are bounded by the value `max_value(uint8)`
     # and therefore cannot overflow the `uint256` type when they
-    # are added together or incremented by 1.
+    # are added together or incremented by `1`.
     assert unsafe_add(leaves_length, len(proof)) == unsafe_add(total_hashes, 1), "MerkleProof: invalid multiproof"
 
-    hashes: DynArray[bytes32, max_value(uint16)] = []
+    hashes: DynArray[bytes32, max_value(uint8)] = []
     leaf_pos: uint256 = empty(uint256)
     hash_pos: uint256 = empty(uint256)
     proof_pos: uint256 = empty(uint256)
@@ -210,5 +210,5 @@ def _hash_pair(a: bytes32, b: bytes32) -> bytes32:
     """
     if (convert(a, uint256) < convert(b, uint256)):
         return keccak256(concat(a, b))
-    else:
-        return keccak256(concat(b, a))
+
+    return keccak256(concat(b, a))
diff --git a/src/snekmate/utils/mocks/MerkleProofVerificationMock.vy b/src/snekmate/utils/mocks/MerkleProofVerificationMock.vy
@@ -0,0 +1,70 @@
+# pragma version ~=0.4.0b6
+"""
+@title MerkleProofVerification Module Reference Implementation
+@custom:contract-name MerkleProofVerificationMock
+@license GNU Affero General Public License v3.0 only
+@author pcaversaccio
+"""
+
+
+# @dev We import the `MerkleProofVerification` module.
+# @notice Please note that the `MerkleProofVerification`
+# module is stateless and therefore does not require
+# the `initializes` keyword for initialisation.
+from .. import MerkleProofVerification as mp
+
+
+@deploy
+@payable
+def __init__():
+    """
+    @dev To omit the opcodes for checking the `msg.value`
+         in the creation-time EVM bytecode, the constructor
+         is declared as `payable`.
+    """
+    pass
+
+
+@external
+@pure
+def verify(proof: DynArray[bytes32, max_value(uint8)], root: bytes32, leaf: bytes32) -> bool:
+    """
+    @dev Returns `True` if it can be proved that a `leaf` is
+         part of a Merkle tree defined by `root`.
+    @notice Each pair of leaves and each pair of pre-images
+            are assumed to be sorted.
+    @param proof The 32-byte array containing sibling hashes
+           on the branch from the `leaf` to the `root` of the
+           Merkle tree.
+    @param root The 32-byte Merkle root hash.
+    @param leaf The 32-byte leaf hash.
+    @return bool The verification whether `leaf` is part of
+            a Merkle tree defined by `root`.
+    """
+    return mp._verify(proof, root, leaf)
+
+
+@external
+@pure
+def multi_proof_verify(proof: DynArray[bytes32, max_value(uint8)], proof_flags: DynArray[bool, max_value(uint8)],
+                       root: bytes32, leaves: DynArray[bytes32, max_value(uint8)]) -> bool:
+    """
+    @dev Returns `True` if it can be simultaneously proved that
+         `leaves` are part of a Merkle tree defined by `root`
+         and a given set of `proof_flags`.
+    @notice Note that not all Merkle trees allow for multiproofs.
+            See {MerkleProofVerification-_process_multi_proof} for
+            further details.
+    @param proof The 32-byte array containing sibling hashes
+           on the branches from `leaves` to the `root` of the
+           Merkle tree.
+    @param proof_flags The Boolean array of flags indicating
+           whether another value from the "main queue" (merging
+           branches) or an element from the `proof` array is used
+           to calculate the next hash.
+    @param root The 32-byte Merkle root hash.
+    @param leaves The 32-byte array containing the leaf hashes.
+    @return bool The verification whether `leaves` are simultaneously
+            part of a Merkle tree defined by `root`.
+    """
+    return mp._multi_proof_verify(proof, proof_flags, root, leaves)
diff --git a/test/utils/MerkleProofVerification.t.sol b/test/utils/MerkleProofVerification.t.sol
@@ -277,8 +277,8 @@ contract MerkleProofVerificationTest is Test {
     function setUp() public {
         merkleProofVerification = IMerkleProofVerification(
             vyperDeployer.deployContract(
-                "src/snekmate/utils/",
-                "MerkleProofVerification"
+                "src/snekmate/utils/mocks/",
+                "MerkleProofVerificationMock"
             )
         );
     }