diff --git a/account-compression/sdk/package.json b/account-compression/sdk/package.json index 20e56c7c..c7a0c045 100644 --- a/account-compression/sdk/package.json +++ b/account-compression/sdk/package.json @@ -49,6 +49,7 @@ "test:events": "start-server-and-test start-validator http://localhost:8899/health run-tests:events", "test:accounts": "start-server-and-test start-validator http://localhost:8899/health run-tests:accounts", "test:e2e": "start-server-and-test start-validator http://localhost:8899/health run-tests:e2e", + "test:merkle-tree": "jest tests/merkleTree.test.ts --detectOpenHandles", "test": "start-server-and-test start-validator http://localhost:8899/health run-tests" }, "dependencies": { @@ -56,6 +57,7 @@ "@metaplex-foundation/beet-solana": "^0.4.0", "bn.js": "^5.2.1", "borsh": "^0.7.0", + "js-sha3": "^0.8.0", "typescript-collections": "^1.3.3" }, "peerDependencies": { @@ -80,7 +82,6 @@ "gh-pages": "^4.0.0", "jest": "^29.0.1", "jest-config": "^29.0.1", - "js-sha3": "^0.8.0", "start-server-and-test": "^1.14.0", "ts-jest": "^28.0.8", "ts-jest-resolver": "^2.0.0", diff --git a/account-compression/sdk/src/merkle-tree/index.ts b/account-compression/sdk/src/merkle-tree/index.ts index 49de90b1..928e1320 100644 --- a/account-compression/sdk/src/merkle-tree/index.ts +++ b/account-compression/sdk/src/merkle-tree/index.ts @@ -18,6 +18,12 @@ export class MerkleTree { root: Buffer; depth: number; + /** + * Please use `MerkleTree.sparseMerkleTreeFromLeaves` to + * create trees instead. This method is exposed for testing purposes, + * and for those that are familiar with the MerkleTree data structure. + * @param leaves leaf nodes of the tree + */ constructor(leaves: Buffer[]) { let [nodes, finalLeaves] = buildLeaves(leaves); let seqNum = leaves.length; @@ -53,6 +59,30 @@ export class MerkleTree { this.depth = nodes.peek()!.level + 1; } + /** + * This is the recommended way to create MerkleTrees. + * If you're trying to match an on-chain MerkleTree, + * set `depth` to `{@link ConcurrentMerkleTreeAccount}.getMaxDepth()` + * + * @param leaves leaves of the tree + * @param depth number of levels in the tree + * @returns MerkleTree + */ + static sparseMerkleTreeFromLeaves( + leaves: Buffer[], + depth: number + ): MerkleTree { + const _leaves: Buffer[] = []; + for (let i = 0; i < 2 ** depth; i++) { + if (i < leaves.length) { + _leaves.push(leaves[i]); + } else { + _leaves.push(Buffer.alloc(32)); + } + } + return new MerkleTree(_leaves); + } + getRoot(): Buffer { return this.root; } @@ -127,11 +157,10 @@ export class MerkleTree { this.root = node.node; } - verify( - root: string, + static hashProof( merkleTreeProof: MerkleTreeProof, - verbose = false - ): boolean { + verbose: boolean = false + ): Buffer { const { leaf, leafIndex, proof } = merkleTreeProof; let node = new PublicKey(leaf).toBuffer(); @@ -143,12 +172,30 @@ export class MerkleTree { } if (verbose) console.log(`node ${i} ${new PublicKey(node).toString()}`); } + return node; + } + + /** + * Verifies that a root matches the proof. + * @param root Root of a MerkleTree + * @param merkleTreeProof Proof to a leaf in the MerkleTree + * @param verbose Whether to print hashed nodes + * @returns Whether the proof is valid + */ + static verify( + root: Buffer, + merkleTreeProof: MerkleTreeProof, + verbose: boolean = false + ): boolean { + const node = MerkleTree.hashProof(merkleTreeProof, verbose); const rehashed = new PublicKey(node).toString(); const received = new PublicKey(root).toString(); - if (verbose) console.log(`hashed ${rehashed} got ${received}`); if (rehashed !== received) { - throw new Error("Roots don't match!!!"); + if (verbose) + console.log(`Roots don't match! Expected ${rehashed} got ${received}`); + return false; } + if (verbose) console.log(`Hashed ${rehashed} got ${received}`); return rehashed === received; } } diff --git a/account-compression/sdk/tests/merkleTree.test.ts b/account-compression/sdk/tests/merkleTree.test.ts new file mode 100644 index 00000000..68f17706 --- /dev/null +++ b/account-compression/sdk/tests/merkleTree.test.ts @@ -0,0 +1,34 @@ +import { assert } from "chai"; +import * as crypto from "crypto"; + +import { emptyNode, MerkleTree } from "../src"; + +describe("MerkleTree tests", () => { + it("Check constructor equivalence for depth 2 tree", () => { + const leaves = [ + crypto.randomBytes(32), + crypto.randomBytes(32), + crypto.randomBytes(32), + ]; + const rawLeaves = leaves.concat(emptyNode(0)); + const merkleTreeRaw = new MerkleTree(rawLeaves); + const merkleTreeSparse = MerkleTree.sparseMerkleTreeFromLeaves(leaves, 2); + + assert(merkleTreeRaw.root.equals(merkleTreeSparse.root)); + }); + + const TEST_DEPTH = 14; + it(`Check proofs for 2^${TEST_DEPTH} tree`, () => { + const leaves: Buffer[] = []; + for (let i = 0; i < 2 ** TEST_DEPTH; i++) { + leaves.push(crypto.randomBytes(32)); + } + const merkleTree = new MerkleTree(leaves); + + // Check proofs + for (let i = 0; i < leaves.length; i++) { + const proof = merkleTree.getProof(i); + assert(MerkleTree.verify(merkleTree.getRoot(), proof)); + } + }); +});