Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for custom internal hash in SimpleMerkleTree #35

Closed
wants to merge 6 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 18 additions & 12 deletions src/core.ts
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 { BytesLike, HexString, toHex, toBytes } 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,9 @@ 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[]): HexString[];
export function makeMerkleTree(leaves: BytesLike[], hash: HashPairFn): HexString[];
export function makeMerkleTree(leaves: BytesLike[], hash?: HashPairFn): HexString[] {
leaves.forEach(checkValidMerkleNode);

if (leaves.length === 0) {
Expand All @@ -32,7 +32,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 +52,13 @@ 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[]): HexString;
export function processProof(leaf: BytesLike, proof: BytesLike[], hash: HashPairFn): 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 +105,9 @@ export function getMultiProof(tree: BytesLike[], indices: number[]): MultiProof<
};
}

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

Expand All @@ -124,7 +128,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 +138,9 @@ export function processMultiProof(multiproof: MultiProof<BytesLike>): HexString
return toHex(stack.pop() ?? proof.shift()!);
}

export function isValidMerkleTree(tree: BytesLike[]): boolean {
export function isValidMerkleTree(tree: BytesLike[]): boolean;
export function isValidMerkleTree(tree: BytesLike[], hash: HashPairFn): 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 +153,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
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' | 'simple-v1';
format: 'standard-v1' | 'simple-v1' | 'custom-v1';
tree: HexString[];
values: {
value: T;
Expand Down
8 changes: 8 additions & 0 deletions src/hashes.ts
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)));
}
124 changes: 88 additions & 36 deletions src/simple.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,16 @@
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 @@

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 @@
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 @@
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 @@ -126,6 +165,19 @@
assert.throws(
() => SimpleMerkleTree.load({ format: 'standard-v1' } as any),
/^Error: Unknown format 'standard-v1'$/,
});

Check failure on line 168 in src/simple.test.ts

View workflow job for this annotation

GitHub Actions / test

Argument expression expected.

it('reject standard tree dump with a custom hash', () => {
assert.throws(
() => SimpleMerkleTree.load({ format: 'simple-v1'} as any, reverseHashPair),
/^Error: Format 'simple-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$/,
);
});

Expand Down
Loading
Loading