Merge pull request #223 from artwyman/artwyman/static-imt-proof

feat(imt): add static version of IMT verifyProof
privacy-scaling-explorations · Mar 21, 2024 · a77a224 · a77a224
2 parents 56320f9 + a8a96fd
commit a77a224
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 7 deletions.
diff --git a/packages/imt/src/imt/imt.ts b/packages/imt/src/imt/imt.ts
@@ -292,11 +292,24 @@ export default class IMT {
 
     /**
      * It verifies a {@link IMTMerkleProof} to confirm that a leaf indeed
-     * belongs to the tree.
+     * belongs to a tree.  Does not verify that the node belongs to this
+     * tree in particular.  Equivalent to `IMT.verifyProof(proof, this._hash)`.
+     *
      * @param proof The Merkle tree proof.
      * @returns True if the leaf is part of the tree, and false otherwise.
      */
     public verifyProof(proof: IMTMerkleProof): boolean {
+        return IMT.verifyProof(proof, this._hash)
+    }
+
+    /**
+     * It verifies a {@link IMTMerkleProof} to confirm that a leaf indeed
+     * belongs to a tree.
+     * @param proof The Merkle tree proof.
+     * @param hash The hash function used to compute the tree nodes.
+     * @returns True if the leaf is part of the tree, and false otherwise.
+     */
+    public static verifyProof(proof: IMTMerkleProof, hash: IMTHashFunction): boolean {
         checkParameter(proof, "proof", "object")
         checkParameter(proof.root, "proof.root", "number", "string", "bigint")
         checkParameter(proof.leaf, "proof.leaf", "number", "string", "bigint")
@@ -310,7 +323,7 @@ export default class IMT {
 
             children.splice(proof.pathIndices[i], 0, node)
 
-            node = this._hash(children)
+            node = hash(children)
         }
 
         return proof.root === node

diff --git a/packages/imt/src/lean-imt/lean-imt.ts b/packages/imt/src/lean-imt/lean-imt.ts
@@ -268,11 +268,23 @@ export default class LeanIMT<N = bigint> {
 
     /**
      * It verifies a {@link LeanIMTMerkleProof} to confirm that a leaf indeed
-     * belongs to the tree.
+     * belongs to a tree.  Does not verify that the node belongs to this
+     * tree in particular.  Equivalent to
+     * `LeanIMT.verifyProof(proof, this._hash)`.
      * @param proof The Merkle tree proof.
      * @returns True if the leaf is part of the tree, and false otherwise.
      */
     public verifyProof(proof: LeanIMTMerkleProof<N>): boolean {
+        return LeanIMT.verifyProof(proof, this._hash)
+    }
+
+    /**
+     * It verifies a {@link LeanIMTMerkleProof} to confirm that a leaf indeed
+     * belongs to a tree.
+     * @param proof The Merkle tree proof.
+     * @returns True if the leaf is part of the tree, and false otherwise.
+     */
+    public static verifyProof<N>(proof: LeanIMTMerkleProof<N>, hash: LeanIMTHashFunction<N>): boolean {
         requireDefinedParameter(proof, "proof")
 
         const { root, leaf, siblings, index } = proof
@@ -289,9 +301,9 @@ export default class LeanIMT<N = bigint> {
 
         for (let i = 0; i < siblings.length; i += 1) {
             if ((index >> i) & 1) {
-                node = this._hash(siblings[i], node)
+                node = hash(siblings[i], node)
             } else {
-                node = this._hash(node, siblings[i])
+                node = hash(node, siblings[i])
             }
         }
 

diff --git a/packages/imt/tests/imt.test.ts b/packages/imt/tests/imt.test.ts
@@ -1,6 +1,6 @@
 import { poseidon2, poseidon3, poseidon5 } from "poseidon-lite"
 import { IncrementalQuinTree } from "incrementalquintree"
-import { IMT } from "../src"
+import { IMT, IMTNode } from "../src"
 
 describe("IMT", () => {
     const depth = 16
@@ -153,7 +153,7 @@ describe("IMT", () => {
                     expect(fun).toThrow("The leaf does not exist in this tree")
                 })
 
-                it("Should create a valid proof", () => {
+                it("Should create and verify valid proof", () => {
                     for (let i = 0; i < numberOfLeaves; i += 1) {
                         tree.insert(i + 1)
                     }
@@ -170,8 +170,21 @@ describe("IMT", () => {
                         }
 
                         expect(tree.verifyProof(proof)).toBeTruthy()
+                        expect(IMT.verifyProof(proof, poseidon)).toBeTruthy()
                     }
                 })
+
+                it("Should reject a proof with incorrect hash function", () => {
+                    for (let i = 0; i < numberOfLeaves; i += 1) {
+                        tree.insert(i + 1)
+                    }
+
+                    const proof = tree.createProof(numberOfLeaves - 1)
+                    function badHash(values: IMTNode[]): IMTNode {
+                        return values.reduce((a, b) => BigInt(a) + BigInt(b), BigInt(0))
+                    }
+                    expect(IMT.verifyProof(proof, badHash)).toBeFalsy()
+                })
             })
         })
     }

diff --git a/packages/imt/tests/lean-imt.test.ts b/packages/imt/tests/lean-imt.test.ts
@@ -316,8 +316,22 @@ describe("Lean IMT", () => {
             const proof = tree.generateProof(3)
 
             expect(tree.verifyProof(proof)).toBe(true)
+            expect(LeanIMT.verifyProof(proof, poseidon)).toBe(true)
         })
 
+        it("Should reject a proof with incorrect hash function", () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const proof = tree.generateProof(3)
+
+            function badHash(a: bigint, b: bigint): bigint {
+                return a + b
+            }
+            expect(LeanIMT.verifyProof(proof, badHash)).toBe(false)
+        })
+    })
+
+    describe("# import/export", () => {
         it("Should export a tree", () => {
             const tree = new LeanIMT(poseidon, leaves)