Skip to content

Commit

Permalink
add support for custom internal hash in SimpleMerkleTree
Browse files Browse the repository at this point in the history
  • Loading branch information
@Amxx
Amxx committed Feb 22, 2024
1 parent 956d681 commit a97a469
Show file tree
Hide file tree
Showing 5 changed files with 145 additions and 69 deletions.
20 changes: 9 additions & 11 deletions src/core.ts
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
import { keccak256 } from '@ethersproject/keccak256';
import { BytesLike, HexString, toHex, toBytes, concat, compare } from './bytes';
import { HashPairFn, keccak256SortedPair } from './hashes';
import { throwError } from './utils/throw-error';

const hashPair = (a: BytesLike, b: BytesLike): HexString => keccak256(concat([a, b].sort(compare)));

const leftChildIndex = (i: number) => 2 * i + 1;
const rightChildIndex = (i: number) => 2 * i + 2;
const parentIndex = (i: number) => i > 0 ? Math.floor((i - 1) / 2) : throwError('Root has no parent');
Expand All @@ -19,7 +17,7 @@ const checkInternalNode = (tree: unknown[], i: number) => void (isInternalNod
const checkLeafNode = (tree: unknown[], i: number) => void (isLeafNode(tree, i) || throwError('Index is not a leaf'));
const checkValidMerkleNode = (node: BytesLike) => void (isValidMerkleNode(node) || throwError('Merkle tree nodes must be Uint8Array of length 32'));

export function makeMerkleTree(leaves: BytesLike[]): HexString[] {
export function makeMerkleTree(leaves: BytesLike[], hash?: HashPairFn): HexString[] {
leaves.forEach(checkValidMerkleNode);

if (leaves.length === 0) {
Expand All @@ -32,7 +30,7 @@ export function makeMerkleTree(leaves: BytesLike[]): HexString[] {
tree[tree.length - 1 - i] = toHex(leaf);
}
for (let i = tree.length - 1 - leaves.length; i >= 0; i--) {
tree[i] = hashPair(
tree[i] = (hash ?? keccak256SortedPair)(
tree[leftChildIndex(i)]!,
tree[rightChildIndex(i)]!,
);
Expand All @@ -52,11 +50,11 @@ export function getProof(tree: BytesLike[], index: number): HexString[] {
return proof.map(node => toHex(node));
}

export function processProof(leaf: BytesLike, proof: BytesLike[]): HexString {
export function processProof(leaf: BytesLike, proof: BytesLike[], hash?: HashPairFn): HexString {
checkValidMerkleNode(leaf);
proof.forEach(checkValidMerkleNode);

return toHex(proof.reduce(hashPair, leaf));
return toHex(proof.reduce(hash ?? keccak256SortedPair, leaf));
}

export interface MultiProof<T, L = T> {
Expand Down Expand Up @@ -103,7 +101,7 @@ export function getMultiProof(tree: BytesLike[], indices: number[]): MultiProof<
};
}

export function processMultiProof(multiproof: MultiProof<BytesLike>): HexString {
export function processMultiProof(multiproof: MultiProof<BytesLike>, hash?: HashPairFn): HexString {
multiproof.leaves.forEach(checkValidMerkleNode);
multiproof.proof.forEach(checkValidMerkleNode);

Expand All @@ -124,7 +122,7 @@ export function processMultiProof(multiproof: MultiProof<BytesLike>): HexString
if (a === undefined || b === undefined) {
throw new Error('Broken invariant');
}
stack.push(hashPair(a, b));
stack.push((hash ?? keccak256SortedPair)(a, b));
}

if (stack.length + proof.length !== 1) {
Expand All @@ -134,7 +132,7 @@ export function processMultiProof(multiproof: MultiProof<BytesLike>): HexString
return toHex(stack.pop() ?? proof.shift()!);
}

export function isValidMerkleTree(tree: BytesLike[]): boolean {
export function isValidMerkleTree(tree: BytesLike[], hash?: HashPairFn): boolean {
for (const [i, node] of tree.entries()) {
if (!isValidMerkleNode(node)) {
return false;
Expand All @@ -147,7 +145,7 @@ export function isValidMerkleTree(tree: BytesLike[]): boolean {
if (l < tree.length) {
return false;
}
} else if (node !== hashPair(tree[l]!, tree[r]!)) {
} else if (node !== (hash ?? keccak256SortedPair)(tree[l]!, tree[r]!)) {
return false;
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/format.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import type { HexString } from "./bytes";

// Dump/Load format
export type MerkleTreeData<T> = {
format: 'standard-v1';
format: 'standard-v1' | 'custom-v1';

Check warning on line 5 in src/format.ts

View check run for this annotation

Codecov / codecov/patch

src/format.ts#L5 

Added line #L5 was not covered by tests
tree: HexString[];
values: {
value: T;
Expand Down
8 changes: 8 additions & 0 deletions src/hashes.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
import { keccak256 } from '@ethersproject/keccak256';
import { BytesLike, HexString, concat, compare } from './bytes';

export type HashPairFn = (a: BytesLike, b: BytesLike) => HexString;

export function keccak256SortedPair(a: BytesLike, b: BytesLike): HexString {
return keccak256(concat([a, b].sort(compare)));
}
125 changes: 89 additions & 36 deletions src/simple.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,16 @@ import assert from 'assert/strict';
import { HashZero as zero } from '@ethersproject/constants';
import { keccak256 } from '@ethersproject/keccak256';
import { SimpleMerkleTree } from './simple';
import { BytesLike, HexString, concat, compare } from './bytes';

const reverseHashPair = (a: BytesLike, b: BytesLike): HexString => keccak256(concat([a, b].sort(compare).reverse()));

describe('simple merkle tree', () => {
for (const opts of [
{},
{ sortLeaves: true },
{ sortLeaves: false },
{ hashPair: reverseHashPair },
]) {
describe(`with options '${JSON.stringify(opts)}'`, () => {
const leaves = 'abcdef'.split('').map(c => keccak256(Buffer.from(c)));
Expand All @@ -28,7 +32,11 @@ describe('simple merkle tree', () => {

assert(tree.verify(id, proof1));
assert(tree.verify(leaf, proof1));
assert(SimpleMerkleTree.verify(tree.root, leaf, proof1));
if (opts.hashPair) {
assert(SimpleMerkleTree.verify(tree.root, leaf, proof1, opts.hashPair));
} else {
assert(SimpleMerkleTree.verify(tree.root, leaf, proof1));
}
}
});

Expand All @@ -37,7 +45,11 @@ describe('simple merkle tree', () => {
const invalidProof = otherTree.getProof(leaf);

assert(!tree.verify(leaf, invalidProof));
assert(!SimpleMerkleTree.verify(tree.root, leaf, invalidProof));
if (opts.hashPair) {
assert(!SimpleMerkleTree.verify(tree.root, leaf, invalidProof, opts.hashPair));
} else {
assert(!SimpleMerkleTree.verify(tree.root, leaf, invalidProof));
}
});

it('generates valid multiproofs', () => {
Expand All @@ -48,53 +60,80 @@ describe('simple merkle tree', () => {
assert.deepEqual(proof1, proof2);

assert(tree.verifyMultiProof(proof1));
assert(SimpleMerkleTree.verifyMultiProof(tree.root, proof1));
if (opts.hashPair) {
assert(SimpleMerkleTree.verifyMultiProof(tree.root, proof1, opts.hashPair));
} else {
assert(SimpleMerkleTree.verifyMultiProof(tree.root, proof1));
}
}
});

it('rejects invalid multiproofs', () => {
const multiProof = otherTree.getMultiProof(leaves.slice(0, 3));

assert(!tree.verifyMultiProof(multiProof));
assert(!SimpleMerkleTree.verifyMultiProof(tree.root, multiProof));
if (opts.hashPair) {
assert(!SimpleMerkleTree.verifyMultiProof(tree.root, multiProof, opts.hashPair));
} else {
assert(!SimpleMerkleTree.verifyMultiProof(tree.root, multiProof));
}
});

it('renders tree representation', () => {
assert.equal(
tree.render(),
opts.sortLeaves == false
? [
"0) 0x9012f1e18a87790d2e01faace75aaaca38e53df437cdce2c0552464dda4af49c",
"├─ 1) 0x68203f90e9d07dc5859259d7536e87a6ba9d345f2552b5b9de2999ddce9ce1bf",
"│ ├─ 3) 0xd253a52d4cb00de2895e85f2529e2976e6aaaa5c18106b68ab66813e14415669",
"│ │ ├─ 7) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
"│ │ └─ 8) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
"│ └─ 4) 0x805b21d846b189efaeb0377d6bb0d201b3872a363e607c25088f025b0c6ae1f8",
"│ ├─ 9) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
"│ └─ 10) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
"└─ 2) 0xf0b49bb4b0d9396e0315755ceafaa280707b32e75e6c9053f5cdf2679dcd5c6a",
" ├─ 5) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
" └─ 6) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
].join("\n")
: [
"0) 0x1b404f199ea828ec5771fb30139c222d8417a82175fefad5cd42bc3a189bd8d5",
"├─ 1) 0xec554bdfb01d31fa838d0830339b0e6e8a70e0d55a8f172ffa8bebbf8e8d5ba0",
"│ ├─ 3) 0x434d51cfeb80272378f4c3a8fd2824561c2cad9fce556ea600d46f20550976a6",
"│ │ ├─ 7) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
"│ │ └─ 8) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
"│ └─ 4) 0x7dea550f679f3caab547cbbc5ee1a4c978c8c039b572ba00af1baa6481b88360",
"│ ├─ 9) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
"│ └─ 10) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
"└─ 2) 0xaf46af0745b433e1d5bed9a04b1fdf4002f67a733c20db2fca5b2af6120d9bcb",
" ├─ 5) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
" └─ 6) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
].join("\n"),
);
const expected = (
// standard hash + unsorted
!opts.hashPair && opts.sortLeaves === false
? [
"0) 0x9012f1e18a87790d2e01faace75aaaca38e53df437cdce2c0552464dda4af49c",
"├─ 1) 0x68203f90e9d07dc5859259d7536e87a6ba9d345f2552b5b9de2999ddce9ce1bf",
"│ ├─ 3) 0xd253a52d4cb00de2895e85f2529e2976e6aaaa5c18106b68ab66813e14415669",
"│ │ ├─ 7) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
"│ │ └─ 8) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
"│ └─ 4) 0x805b21d846b189efaeb0377d6bb0d201b3872a363e607c25088f025b0c6ae1f8",
"│ ├─ 9) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
"│ └─ 10) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
"└─ 2) 0xf0b49bb4b0d9396e0315755ceafaa280707b32e75e6c9053f5cdf2679dcd5c6a",
" ├─ 5) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
" └─ 6) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
]
// sortLeaves = true | undefined --- standard hash + sorted
: !opts.hashPair
? [
"0) 0x1b404f199ea828ec5771fb30139c222d8417a82175fefad5cd42bc3a189bd8d5",
"├─ 1) 0xec554bdfb01d31fa838d0830339b0e6e8a70e0d55a8f172ffa8bebbf8e8d5ba0",
"│ ├─ 3) 0x434d51cfeb80272378f4c3a8fd2824561c2cad9fce556ea600d46f20550976a6",
"│ │ ├─ 7) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
"│ │ └─ 8) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
"│ └─ 4) 0x7dea550f679f3caab547cbbc5ee1a4c978c8c039b572ba00af1baa6481b88360",
"│ ├─ 9) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
"│ └─ 10) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
"└─ 2) 0xaf46af0745b433e1d5bed9a04b1fdf4002f67a733c20db2fca5b2af6120d9bcb",
" ├─ 5) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
" └─ 6) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
]
// non standard hash
: [
"0) 0x8f0a1adb058c628fa4ce2e7bd26024180b888fec77087d4e5ee6890746e9c6ec",
"├─ 1) 0xb9f5a6bc1b75fadcd9765163dfc8d4865d1608337a2a310ff51fecb431faaee4",
"│ ├─ 3) 0x37d657e93dfbae50b18241610418794b51124af5ca872f1b56c08490cb2905ac",
"│ │ ├─ 7) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
"│ │ └─ 8) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
"│ └─ 4) 0xed90ef72e95e6692b91b020dc6cb5c4db9dc149a496799c4318fa8075960c48e",
"│ ├─ 9) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
"│ └─ 10) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
"└─ 2) 0x138c55cca8f6430d75b6bbcea643a7afa8ee74c22643ad76723ecafd4fcd21d4",
" ├─ 5) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
" └─ 6) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
]
).join("\n");

assert.equal(tree.render(), expected);
});

it('dump and load', () => {
const recoveredTree = SimpleMerkleTree.load(tree.dump());

const recoveredTree = opts.hashPair
? SimpleMerkleTree.load(tree.dump(), opts.hashPair)
: SimpleMerkleTree.load(tree.dump());
recoveredTree.validate();
assert.deepEqual(tree, recoveredTree);
});
Expand Down Expand Up @@ -124,6 +163,20 @@ describe('simple merkle tree', () => {
);
});

it('reject standard tree dump with a custom hash', () => {
assert.throws(
() => SimpleMerkleTree.load({ format: 'standard-v1'} as any, reverseHashPair),
/^Error: Format 'standard-v1' does not support custom hashing functions$/,
);
});

it('reject custom tree dump without a custom hash', () => {
assert.throws(
() => SimpleMerkleTree.load({ format: 'custom-v1'} as any),
/^Error: Format 'custom-v1' requires a hashing function$/,
);
});

it('reject malformed tree dump', () => {
const loadedTree1 = SimpleMerkleTree.load({
format: 'standard-v1',
Expand Down
Loading

0 comments on commit a97a469

Please sign in to comment.