Skip to content

Commit a97a469

Browse files
AmxxAmxx
committed
add support for custom internal hash in SimpleMerkleTree
1 parent 956d681 commit a97a469

File tree

5 files changed

+145
-69
lines changed

5 files changed

+145
-69
lines changed

‎src/core.ts

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
1-
import { keccak256 } from '@ethersproject/keccak256';
21
import { BytesLike, HexString, toHex, toBytes, concat, compare } from './bytes';
2+
import { HashPairFn, keccak256SortedPair } from './hashes';
33
import { throwError } from './utils/throw-error';
44

5-
const hashPair = (a: BytesLike, b: BytesLike): HexString => keccak256(concat([a, b].sort(compare)));
6-
75
const leftChildIndex = (i: number) => 2 * i + 1;
86
const rightChildIndex = (i: number) => 2 * i + 2;
97
const parentIndex = (i: number) => i > 0 ? Math.floor((i - 1) / 2) : throwError('Root has no parent');
@@ -19,7 +17,7 @@ const checkInternalNode = (tree: unknown[], i: number) => void (isInternalNod
1917
const checkLeafNode = (tree: unknown[], i: number) => void (isLeafNode(tree, i) || throwError('Index is not a leaf'));
2018
const checkValidMerkleNode = (node: BytesLike) => void (isValidMerkleNode(node) || throwError('Merkle tree nodes must be Uint8Array of length 32'));
2119

22-
export function makeMerkleTree(leaves: BytesLike[]): HexString[] {
20+
export function makeMerkleTree(leaves: BytesLike[], hash?: HashPairFn): HexString[] {
2321
leaves.forEach(checkValidMerkleNode);
2422

2523
if (leaves.length === 0) {
@@ -32,7 +30,7 @@ export function makeMerkleTree(leaves: BytesLike[]): HexString[] {
3230
tree[tree.length - 1 - i] = toHex(leaf);
3331
}
3432
for (let i = tree.length - 1 - leaves.length; i >= 0; i--) {
35-
tree[i] = hashPair(
33+
tree[i] = (hash ?? keccak256SortedPair)(
3634
tree[leftChildIndex(i)]!,
3735
tree[rightChildIndex(i)]!,
3836
);
@@ -52,11 +50,11 @@ export function getProof(tree: BytesLike[], index: number): HexString[] {
5250
return proof.map(node => toHex(node));
5351
}
5452

55-
export function processProof(leaf: BytesLike, proof: BytesLike[]): HexString {
53+
export function processProof(leaf: BytesLike, proof: BytesLike[], hash?: HashPairFn): HexString {
5654
checkValidMerkleNode(leaf);
5755
proof.forEach(checkValidMerkleNode);
5856

59-
return toHex(proof.reduce(hashPair, leaf));
57+
return toHex(proof.reduce(hash ?? keccak256SortedPair, leaf));
6058
}
6159

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

106-
export function processMultiProof(multiproof: MultiProof<BytesLike>): HexString {
104+
export function processMultiProof(multiproof: MultiProof<BytesLike>, hash?: HashPairFn): HexString {
107105
multiproof.leaves.forEach(checkValidMerkleNode);
108106
multiproof.proof.forEach(checkValidMerkleNode);
109107

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

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

137-
export function isValidMerkleTree(tree: BytesLike[]): boolean {
135+
export function isValidMerkleTree(tree: BytesLike[], hash?: HashPairFn): boolean {
138136
for (const [i, node] of tree.entries()) {
139137
if (!isValidMerkleNode(node)) {
140138
return false;
@@ -147,7 +145,7 @@ export function isValidMerkleTree(tree: BytesLike[]): boolean {
147145
if (l < tree.length) {
148146
return false;
149147
}
150-
} else if (node !== hashPair(tree[l]!, tree[r]!)) {
148+
} else if (node !== (hash ?? keccak256SortedPair)(tree[l]!, tree[r]!)) {
151149
return false;
152150
}
153151
}

‎src/format.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type { HexString } from "./bytes";
22

33
// Dump/Load format
44
export type MerkleTreeData<T> = {
5-
format: 'standard-v1';
5+
format: 'standard-v1' | 'custom-v1';
66
tree: HexString[];
77
values: {
88
value: T;

‎src/hashes.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import { keccak256 } from '@ethersproject/keccak256';
2+
import { BytesLike, HexString, concat, compare } from './bytes';
3+
4+
export type HashPairFn = (a: BytesLike, b: BytesLike) => HexString;
5+
6+
export function keccak256SortedPair(a: BytesLike, b: BytesLike): HexString {
7+
return keccak256(concat([a, b].sort(compare)));
8+
}

‎src/simple.test.ts

Lines changed: 89 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,16 @@ import assert from 'assert/strict';
22
import { HashZero as zero } from '@ethersproject/constants';
33
import { keccak256 } from '@ethersproject/keccak256';
44
import { SimpleMerkleTree } from './simple';
5+
import { BytesLike, HexString, concat, compare } from './bytes';
6+
7+
const reverseHashPair = (a: BytesLike, b: BytesLike): HexString => keccak256(concat([a, b].sort(compare).reverse()));
58

69
describe('simple merkle tree', () => {
710
for (const opts of [
811
{},
912
{ sortLeaves: true },
1013
{ sortLeaves: false },
14+
{ hashPair: reverseHashPair },
1115
]) {
1216
describe(`with options '${JSON.stringify(opts)}'`, () => {
1317
const leaves = 'abcdef'.split('').map(c => keccak256(Buffer.from(c)));
@@ -28,7 +32,11 @@ describe('simple merkle tree', () => {
2832

2933
assert(tree.verify(id, proof1));
3034
assert(tree.verify(leaf, proof1));
31-
assert(SimpleMerkleTree.verify(tree.root, leaf, proof1));
35+
if (opts.hashPair) {
36+
assert(SimpleMerkleTree.verify(tree.root, leaf, proof1, opts.hashPair));
37+
} else {
38+
assert(SimpleMerkleTree.verify(tree.root, leaf, proof1));
39+
}
3240
}
3341
});
3442

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

3947
assert(!tree.verify(leaf, invalidProof));
40-
assert(!SimpleMerkleTree.verify(tree.root, leaf, invalidProof));
48+
if (opts.hashPair) {
49+
assert(!SimpleMerkleTree.verify(tree.root, leaf, invalidProof, opts.hashPair));
50+
} else {
51+
assert(!SimpleMerkleTree.verify(tree.root, leaf, invalidProof));
52+
}
4153
});
4254

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

5062
assert(tree.verifyMultiProof(proof1));
51-
assert(SimpleMerkleTree.verifyMultiProof(tree.root, proof1));
63+
if (opts.hashPair) {
64+
assert(SimpleMerkleTree.verifyMultiProof(tree.root, proof1, opts.hashPair));
65+
} else {
66+
assert(SimpleMerkleTree.verifyMultiProof(tree.root, proof1));
67+
}
5268
}
5369
});
5470

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

5874
assert(!tree.verifyMultiProof(multiProof));
59-
assert(!SimpleMerkleTree.verifyMultiProof(tree.root, multiProof));
75+
if (opts.hashPair) {
76+
assert(!SimpleMerkleTree.verifyMultiProof(tree.root, multiProof, opts.hashPair));
77+
} else {
78+
assert(!SimpleMerkleTree.verifyMultiProof(tree.root, multiProof));
79+
}
6080
});
6181

6282
it('renders tree representation', () => {
63-
assert.equal(
64-
tree.render(),
65-
opts.sortLeaves == false
66-
? [
67-
"0) 0x9012f1e18a87790d2e01faace75aaaca38e53df437cdce2c0552464dda4af49c",
68-
"├─ 1) 0x68203f90e9d07dc5859259d7536e87a6ba9d345f2552b5b9de2999ddce9ce1bf",
69-
"│ ├─ 3) 0xd253a52d4cb00de2895e85f2529e2976e6aaaa5c18106b68ab66813e14415669",
70-
"│ │ ├─ 7) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
71-
"│ │ └─ 8) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
72-
"│ └─ 4) 0x805b21d846b189efaeb0377d6bb0d201b3872a363e607c25088f025b0c6ae1f8",
73-
"│ ├─ 9) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
74-
"│ └─ 10) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
75-
"└─ 2) 0xf0b49bb4b0d9396e0315755ceafaa280707b32e75e6c9053f5cdf2679dcd5c6a",
76-
" ├─ 5) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
77-
" └─ 6) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
78-
].join("\n")
79-
: [
80-
"0) 0x1b404f199ea828ec5771fb30139c222d8417a82175fefad5cd42bc3a189bd8d5",
81-
"├─ 1) 0xec554bdfb01d31fa838d0830339b0e6e8a70e0d55a8f172ffa8bebbf8e8d5ba0",
82-
"│ ├─ 3) 0x434d51cfeb80272378f4c3a8fd2824561c2cad9fce556ea600d46f20550976a6",
83-
"│ │ ├─ 7) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
84-
"│ │ └─ 8) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
85-
"│ └─ 4) 0x7dea550f679f3caab547cbbc5ee1a4c978c8c039b572ba00af1baa6481b88360",
86-
"│ ├─ 9) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
87-
"│ └─ 10) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
88-
"└─ 2) 0xaf46af0745b433e1d5bed9a04b1fdf4002f67a733c20db2fca5b2af6120d9bcb",
89-
" ├─ 5) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
90-
" └─ 6) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
91-
].join("\n"),
92-
);
83+
const expected = (
84+
// standard hash + unsorted
85+
!opts.hashPair && opts.sortLeaves === false
86+
? [
87+
"0) 0x9012f1e18a87790d2e01faace75aaaca38e53df437cdce2c0552464dda4af49c",
88+
"├─ 1) 0x68203f90e9d07dc5859259d7536e87a6ba9d345f2552b5b9de2999ddce9ce1bf",
89+
"│ ├─ 3) 0xd253a52d4cb00de2895e85f2529e2976e6aaaa5c18106b68ab66813e14415669",
90+
"│ │ ├─ 7) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
91+
"│ │ └─ 8) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
92+
"│ └─ 4) 0x805b21d846b189efaeb0377d6bb0d201b3872a363e607c25088f025b0c6ae1f8",
93+
"│ ├─ 9) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
94+
"│ └─ 10) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
95+
"└─ 2) 0xf0b49bb4b0d9396e0315755ceafaa280707b32e75e6c9053f5cdf2679dcd5c6a",
96+
" ├─ 5) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
97+
" └─ 6) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
98+
]
99+
// sortLeaves = true | undefined --- standard hash + sorted
100+
: !opts.hashPair
101+
? [
102+
"0) 0x1b404f199ea828ec5771fb30139c222d8417a82175fefad5cd42bc3a189bd8d5",
103+
"├─ 1) 0xec554bdfb01d31fa838d0830339b0e6e8a70e0d55a8f172ffa8bebbf8e8d5ba0",
104+
"│ ├─ 3) 0x434d51cfeb80272378f4c3a8fd2824561c2cad9fce556ea600d46f20550976a6",
105+
"│ │ ├─ 7) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
106+
"│ │ └─ 8) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
107+
"│ └─ 4) 0x7dea550f679f3caab547cbbc5ee1a4c978c8c039b572ba00af1baa6481b88360",
108+
"│ ├─ 9) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
109+
"│ └─ 10) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
110+
"└─ 2) 0xaf46af0745b433e1d5bed9a04b1fdf4002f67a733c20db2fca5b2af6120d9bcb",
111+
" ├─ 5) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
112+
" └─ 6) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
113+
]
114+
// non standard hash
115+
: [
116+
"0) 0x8f0a1adb058c628fa4ce2e7bd26024180b888fec77087d4e5ee6890746e9c6ec",
117+
"├─ 1) 0xb9f5a6bc1b75fadcd9765163dfc8d4865d1608337a2a310ff51fecb431faaee4",
118+
"│ ├─ 3) 0x37d657e93dfbae50b18241610418794b51124af5ca872f1b56c08490cb2905ac",
119+
"│ │ ├─ 7) 0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
120+
"│ │ └─ 8) 0xa8982c89d80987fb9a510e25981ee9170206be21af3c8e0eb312ef1d3382e761",
121+
"│ └─ 4) 0xed90ef72e95e6692b91b020dc6cb5c4db9dc149a496799c4318fa8075960c48e",
122+
"│ ├─ 9) 0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
123+
"│ └─ 10) 0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
124+
"└─ 2) 0x138c55cca8f6430d75b6bbcea643a7afa8ee74c22643ad76723ecafd4fcd21d4",
125+
" ├─ 5) 0xf1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
126+
" └─ 6) 0xd1e8aeb79500496ef3dc2e57ba746a8315d048b7a664a2bf948db4fa91960483",
127+
]
128+
).join("\n");
129+
130+
assert.equal(tree.render(), expected);
93131
});
94132

95133
it('dump and load', () => {
96-
const recoveredTree = SimpleMerkleTree.load(tree.dump());
97-
134+
const recoveredTree = opts.hashPair
135+
? SimpleMerkleTree.load(tree.dump(), opts.hashPair)
136+
: SimpleMerkleTree.load(tree.dump());
98137
recoveredTree.validate();
99138
assert.deepEqual(tree, recoveredTree);
100139
});
@@ -124,6 +163,20 @@ describe('simple merkle tree', () => {
124163
);
125164
});
126165

166+
it('reject standard tree dump with a custom hash', () => {
167+
assert.throws(
168+
() => SimpleMerkleTree.load({ format: 'standard-v1'} as any, reverseHashPair),
169+
/^Error: Format 'standard-v1' does not support custom hashing functions$/,
170+
);
171+
});
172+
173+
it('reject custom tree dump without a custom hash', () => {
174+
assert.throws(
175+
() => SimpleMerkleTree.load({ format: 'custom-v1'} as any),
176+
/^Error: Format 'custom-v1' requires a hashing function$/,
177+
);
178+
});
179+
127180
it('reject malformed tree dump', () => {
128181
const loadedTree1 = SimpleMerkleTree.load({
129182
format: 'standard-v1',

0 commit comments

Comments
 (0)