Skip to content

Commit 382a6fc

Browse files
authored
[sui 10/x] - merkle tree (#903)
* merkle tree impl * - take leftmost 20 bytes in hash - don't assign output of cursor::take_rest to _, instead just drop it * push PREFIXes (MERKLE_LEAF_PREFIX, MERKLE_NODE_PREFIX) to front instead of back * delete testXOR * test construct merkle tree depth exceeded error * invalid merkle proof test cases * comments * rename failure tests * simplification for initializing a vector * fix leafHash bug, add tests for hashLeaf and hashNode
1 parent 7c037ed commit 382a6fc

File tree

1 file changed

+381
-0
lines changed

1 file changed

+381
-0
lines changed
Lines changed: 381 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,381 @@
1+
// Implementation of a Merkle tree in Move. Supports constructing a new tree
2+
// with a given depth, as well as proving that a leaf node belongs to the tree.
3+
module pyth::merkle_tree {
4+
use std::vector::{Self};
5+
use sui::hash::{keccak256};
6+
use wormhole::bytes::{Self};
7+
use wormhole::bytes20::{Self, Bytes20, data};
8+
use wormhole::cursor::{Self, Cursor};
9+
use pyth::deserialize::{Self};
10+
11+
const MERKLE_LEAF_PREFIX: u8 = 0;
12+
const MERKLE_NODE_PREFIX: u8 = 1;
13+
const MERKLE_EMPTY_LEAF_PREFIX: u8 = 2;
14+
15+
const E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES: u64 = 1212121;
16+
17+
// take keccak256 of input data, then return 20 leftmost bytes of result
18+
fun hash(bytes: &vector<u8>): Bytes20 {
19+
let hashed_bytes = keccak256(bytes);
20+
let cursor = cursor::new(hashed_bytes);
21+
let bytes20 = bytes::take_bytes(&mut cursor, 20);
22+
cursor::take_rest(cursor);
23+
bytes20::from_bytes(bytes20)
24+
}
25+
26+
fun emptyLeafHash(): Bytes20 {
27+
let v = vector<u8>[MERKLE_EMPTY_LEAF_PREFIX];
28+
hash(&v)
29+
}
30+
31+
fun leafHash(data: &vector<u8>): Bytes20 {
32+
let v = vector<u8>[MERKLE_LEAF_PREFIX];
33+
vector::append(&mut v, *data);
34+
hash(&v)
35+
}
36+
37+
fun nodeHash(
38+
childA: Bytes20,
39+
childB: Bytes20
40+
): Bytes20 {
41+
if (greaterThan(childA, childB)) {
42+
(childA, childB) = (childB, childA);
43+
};
44+
// append data_B to data_A
45+
let data_A = bytes20::data(&childA);
46+
let data_B = bytes20::data(&childB);
47+
vector::append(&mut data_A, data_B);
48+
49+
// create a vector containing MERKLE_NODE_PREFIX and append data_A to back
50+
let v = vector::empty<u8>();
51+
vector::push_back(&mut v, MERKLE_NODE_PREFIX);
52+
vector::append(&mut v, data_A);
53+
hash(&v)
54+
}
55+
56+
// greaterThan returns whether a is strictly greater than b
57+
// note that data(&a) and data(&b) are both vectors of length 20
58+
fun greaterThan(a: Bytes20, b: Bytes20): bool{
59+
// aa and bb both have length 20
60+
let aa = data(&a);
61+
let bb = data(&b);
62+
let i = 0;
63+
while (i < 20){
64+
if (*vector::borrow(&aa, i) > *vector::borrow(&bb, i)){
65+
return true
66+
} else if (*vector::borrow(&bb, i) > *vector::borrow(&aa, i)){
67+
return false
68+
};
69+
i = i + 1;
70+
};
71+
false
72+
}
73+
74+
// The Sui Move stdlb insert function shifts v[i] and subsequent elements to the right.
75+
// We don't want this behavior, so we define our own setElement function that instead replaces the ith element.
76+
// Reference: https://github.com/MystenLabs/sui/blob/main/crates/sui-framework/packages/move-stdlib/sources/vector.move
77+
fun setElement<T: drop>(a: &mut vector<T>, value: T, index: u64){
78+
vector::push_back<T>(a, value); // push value to end
79+
vector::swap_remove(a, index); // swap value to correct position and pop last value
80+
}
81+
82+
// isProofValid returns whether a merkle proof is valid
83+
fun isProofValid(
84+
encodedProof: &mut Cursor<u8>,
85+
root: Bytes20,
86+
leafData: vector<u8>,
87+
): bool {
88+
89+
let currentDigest: Bytes20 = leafHash(&leafData);
90+
let proofSize: u8 = deserialize::deserialize_u8(encodedProof);
91+
let i: u8 = 0;
92+
while (i < proofSize){
93+
let siblingDigest: Bytes20 = bytes20::new(
94+
deserialize::deserialize_vector(encodedProof, 20)
95+
);
96+
97+
currentDigest = nodeHash(
98+
currentDigest,
99+
siblingDigest
100+
);
101+
i = i + 1;
102+
};
103+
bytes20::data(&currentDigest) == bytes20::data(&root)
104+
}
105+
106+
// constructProofs constructs a merkle tree and returns the root of the tree as
107+
// a Bytes20 as well as the vector of encoded proofs
108+
fun constructProofs(
109+
messages: &vector<vector<u8>>,
110+
depth: u8
111+
) : (Bytes20, vector<u8>) {
112+
113+
if ( 1 << depth < vector::length(messages)) {
114+
abort E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES
115+
};
116+
117+
// empty tree
118+
// The tree is structured as follows:
119+
// 1
120+
// 2 3
121+
// 4 5 6 7
122+
// ...
123+
// In this structure the parent of node x is x//2 and the children
124+
// of node x are x*2 and x*2 + 1. Also, the sibling of the node x
125+
// is x^1. The root is at index 1 and index 0 is not used.
126+
let tree = vector::empty<Bytes20>();
127+
128+
// empty leaf hash
129+
let cachedEmptyLeafHash: Bytes20 = emptyLeafHash();
130+
131+
// Instantiate tree to be a full binary tree with the appropriate depth.
132+
// Add an entry at the end for swapping
133+
let i: u64 = 0;
134+
while (i < (1 << (depth+1)) + 1){
135+
vector::push_back(&mut tree, cachedEmptyLeafHash);
136+
i = i + 1;
137+
};
138+
139+
// Fill in bottom row with leaf hashes
140+
let j: u64 = 0;
141+
while (j < vector::length(messages)){
142+
setElement<Bytes20>(&mut tree, leafHash(vector::borrow(messages, j)), (1 << depth) + j);
143+
j = j + 1;
144+
};
145+
146+
// Filling the node hashes from bottom to top
147+
let k: u8 = depth;
148+
while (k>0){
149+
let level: u8 = k-1;
150+
let levelNumNodes = 1 << level;
151+
let i: u64 = 0;
152+
while (i < levelNumNodes ){
153+
let id = (1 << level) + i;
154+
let nodeHash = nodeHash(*vector::borrow(&tree, id * 2), *vector::borrow(&tree, id * 2 + 1));
155+
setElement<Bytes20>(&mut tree, nodeHash, id);
156+
i = i + 1;
157+
};
158+
k = k - 1;
159+
};
160+
161+
let root = *vector::borrow(&tree, 1);
162+
163+
// construct proofs and create encoded proofs vector
164+
let proofs = vector::empty<u8>();
165+
let i: u64 = 0;
166+
while (i < vector::length(messages)){
167+
let cur_proof = vector::empty<u8>();
168+
vector::push_back(&mut cur_proof, depth);
169+
let idx = (1 << depth) + i;
170+
while (idx > 1) {
171+
vector::append(&mut cur_proof, bytes20::data(vector::borrow(&tree, idx ^ 1)));
172+
173+
// Jump to parent
174+
idx = idx / 2;
175+
};
176+
vector::append(&mut proofs, cur_proof);
177+
i = i + 1;
178+
};
179+
180+
(root, proofs)
181+
}
182+
183+
#[test]
184+
fun testGreaterThan(){
185+
// test 1
186+
let x = bytes20::new(x"0000000000000000000000000000000000001000");
187+
let y = bytes20::new(x"0000000000000000000000000000000000000001");
188+
let res = greaterThan(x, y);
189+
assert!(res==true, 0);
190+
res = greaterThan(y, x);
191+
assert!(res==false, 0);
192+
193+
// test 2
194+
x = bytes20::new(x"1100000000000000000000000000000000001000");
195+
y = bytes20::new(x"1100000000000000000000000000000000000001");
196+
res = greaterThan(x, y);
197+
assert!(res==true, 0);
198+
199+
// equality case
200+
x = bytes20::new(x"1100000000000000000000000000000000001001");
201+
y = bytes20::new(x"1100000000000000000000000000000000001001");
202+
res = greaterThan(x, y);
203+
assert!(res==false, 0);
204+
}
205+
206+
#[test]
207+
fun test_hash_leaf() {
208+
let data: vector<u8> = x"00640000000000000000000000000000000000000000000000000000000000000000000000000000640000000000000064000000640000000000000064000000000000006400000000000000640000000000000064";
209+
let hash = leafHash(&data);
210+
let expected = bytes20::new(x"afc6a8ac466430f35895055f8a4c951785dad5ce");
211+
assert!(hash == expected, 1);
212+
}
213+
214+
#[test]
215+
fun test_hash_node() {
216+
let h1 = bytes20::new(x"05c51b04b820c0f704e3fdd2e4fc1e70aff26dff");
217+
let h2 = bytes20::new(x"1e108841c8d21c7a5c4860c8c3499c918ea9e0ac");
218+
let hash = nodeHash(h1, h2);
219+
let expected = bytes20::new(x"2d0e4fde68184c7ce8af426a0865bd41ef84dfa4");
220+
assert!(hash == expected, 1);
221+
}
222+
223+
#[test]
224+
fun testMerkleTreeDepth1(){
225+
let messages = vector::empty<vector<u8>>();
226+
vector::push_back(&mut messages, x"1234");
227+
228+
let (root, proofs) = constructProofs(&messages, 1);
229+
230+
let proofsCursor = cursor::new(proofs);
231+
let valid = isProofValid(&mut proofsCursor, root, x"1234");
232+
assert!(valid==true, 0);
233+
234+
// destroy cursor
235+
cursor::take_rest<u8>(proofsCursor);
236+
}
237+
238+
#[test]
239+
fun testMerkleTreeDepth2(){
240+
let messages = vector::empty<vector<u8>>();
241+
vector::push_back(&mut messages, x"1234");
242+
vector::push_back(&mut messages, x"4321");
243+
vector::push_back(&mut messages, x"11");
244+
vector::push_back(&mut messages, x"22");
245+
246+
let (root, proofs) = constructProofs(&messages, 2);
247+
248+
let proofsCursor = cursor::new(proofs);
249+
assert!(isProofValid(&mut proofsCursor, root, x"1234")==true, 0);
250+
assert!(isProofValid(&mut proofsCursor, root, x"4321")==true, 0);
251+
assert!(isProofValid(&mut proofsCursor, root, x"11")==true, 0);
252+
assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0);
253+
254+
// destroy cursor
255+
cursor::take_rest<u8>(proofsCursor);
256+
}
257+
258+
#[test]
259+
fun testMerkleTreeDepth3(){
260+
let messages = vector::empty<vector<u8>>();
261+
vector::push_back(&mut messages, x"00");
262+
vector::push_back(&mut messages, x"4321");
263+
vector::push_back(&mut messages, x"444444");
264+
vector::push_back(&mut messages, x"22222222");
265+
vector::push_back(&mut messages, x"22");
266+
vector::push_back(&mut messages, x"11");
267+
vector::push_back(&mut messages, x"100000");
268+
vector::push_back(&mut messages, x"eeeeee");
269+
270+
let (root, proofs) = constructProofs(&messages, 3);
271+
272+
let proofsCursor = cursor::new(proofs);
273+
assert!(isProofValid(&mut proofsCursor, root, x"00")==true, 0);
274+
assert!(isProofValid(&mut proofsCursor, root, x"4321")==true, 0);
275+
assert!(isProofValid(&mut proofsCursor, root, x"444444")==true, 0);
276+
assert!(isProofValid(&mut proofsCursor, root, x"22222222")==true, 0);
277+
assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0);
278+
assert!(isProofValid(&mut proofsCursor, root, x"11")==true, 0);
279+
assert!(isProofValid(&mut proofsCursor, root, x"100000")==true, 0);
280+
assert!(isProofValid(&mut proofsCursor, root, x"eeeeee")==true, 0);
281+
282+
// destroy cursor
283+
cursor::take_rest<u8>(proofsCursor);
284+
}
285+
286+
#[test]
287+
fun testMerkleTreeDepth1InvalidProofs(){
288+
let messages = vector::empty<vector<u8>>();
289+
vector::push_back(&mut messages, x"1234");
290+
291+
let (root, proofs) = constructProofs(&messages, 1);
292+
293+
let proofsCursor = cursor::new(proofs);
294+
295+
// use wrong leaf data (it is not included in the tree)
296+
let valid = isProofValid(&mut proofsCursor, root, x"432222");
297+
assert!(valid==false, 0);
298+
299+
// destroy cursor
300+
cursor::take_rest<u8>(proofsCursor);
301+
}
302+
303+
#[test]
304+
fun testMerkleTreeDepth2InvalidProofs(){
305+
let messages = vector::empty<vector<u8>>();
306+
vector::push_back(&mut messages, x"1234");
307+
vector::push_back(&mut messages, x"4321");
308+
vector::push_back(&mut messages, x"11");
309+
vector::push_back(&mut messages, x"22");
310+
311+
let (root, proofs) = constructProofs(&messages, 2);
312+
313+
let proofsCursor = cursor::new(proofs);
314+
// proof fails because we used the proof of x"1234" to try to prove that x"4321" is in the tree
315+
assert!(isProofValid(&mut proofsCursor, root, x"4321")==false, 0);
316+
// proof succeeds
317+
assert!(isProofValid(&mut proofsCursor, root, x"4321")==true, 0);
318+
// proof fails because we used the proof of x"11" to try to prove that x"22" is in the tree
319+
assert!(isProofValid(&mut proofsCursor, root, x"22")==false, 0);
320+
// proof succeeds
321+
assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0);
322+
323+
// destroy cursor
324+
cursor::take_rest<u8>(proofsCursor);
325+
}
326+
327+
#[test]
328+
fun testMerkleTreeDepth3InvalidProofs(){
329+
let messages = vector::empty<vector<u8>>();
330+
vector::push_back(&mut messages, x"00");
331+
vector::push_back(&mut messages, x"4321");
332+
vector::push_back(&mut messages, x"444444");
333+
vector::push_back(&mut messages, x"22222222");
334+
vector::push_back(&mut messages, x"22");
335+
vector::push_back(&mut messages, x"11");
336+
vector::push_back(&mut messages, x"100000");
337+
vector::push_back(&mut messages, x"eeeeee");
338+
339+
let (root, proofs) = constructProofs(&messages, 3);
340+
341+
let proofsCursor = cursor::new(proofs);
342+
343+
// test various proof failure cases (because of mismatch between proof and leaf data)
344+
assert!(isProofValid(&mut proofsCursor, root, x"00")==true, 0);
345+
assert!(isProofValid(&mut proofsCursor, root, x"22")==false, 0);
346+
assert!(isProofValid(&mut proofsCursor, root, x"22222222")==false, 0);
347+
assert!(isProofValid(&mut proofsCursor, root, x"22222222")==true, 0);
348+
assert!(isProofValid(&mut proofsCursor, root, x"22")==true, 0);
349+
assert!(isProofValid(&mut proofsCursor, root, x"eeeeee")==false, 0);
350+
assert!(isProofValid(&mut proofsCursor, root, x"4321")==false, 0);
351+
assert!(isProofValid(&mut proofsCursor, root, x"eeeeee")==true, 0);
352+
353+
// destroy cursor
354+
cursor::take_rest<u8>(proofsCursor);
355+
}
356+
357+
#[test]
358+
#[expected_failure(abort_code = pyth::merkle_tree::E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES)]
359+
fun testMerkleTreeDepthExceeded1(){
360+
let messages = vector::empty<vector<u8>>();
361+
vector::push_back(&mut messages, x"00");
362+
vector::push_back(&mut messages, x"4321");
363+
vector::push_back(&mut messages, x"444444");
364+
365+
constructProofs(&messages, 1); //depth 1
366+
}
367+
368+
#[test]
369+
#[expected_failure(abort_code = pyth::merkle_tree::E_DEPTH_NOT_LARGE_ENOUGH_FOR_MESSAGES)]
370+
fun testMerkleTreeDepthExceeded2(){
371+
let messages = vector::empty<vector<u8>>();
372+
vector::push_back(&mut messages, x"00");
373+
vector::push_back(&mut messages, x"4321");
374+
vector::push_back(&mut messages, x"444444");
375+
vector::push_back(&mut messages, x"22222222");
376+
vector::push_back(&mut messages, x"22");
377+
378+
constructProofs(&messages, 2); // depth 2
379+
}
380+
381+
}

0 commit comments

Comments
 (0)