merkletree-rs/src/lib.rs
src/lib.rsRaw
1use std::collections::HashMap;
2
3#[macro_use]
4extern crate arrayref;
5
6extern crate tiny_keccak;
7
8extern crate rustc_hex;
9
10extern crate hex;
11
12mod utils;
13mod node;
14
15const TYPENODEEMPTY: u8 = 0;
16const TYPENODENORMAL: u8 = 1;
17const TYPENODEFINAL: u8 = 2;
18const TYPENODEVALUE: u8 = 3;
19const EMPTYNODEVALUE: [u8;32] = [0;32];
20
21pub struct TestValue {
22 bytes: Vec<u8>,
23 index_length: u32,
24}
25pub trait Value {
26 fn bytes(&self) -> &Vec<u8>;
27 fn index_length(&self) -> u32;
28 fn hi(&self) -> [u8;32];
29 fn ht(&self) -> [u8;32];
30}
31impl Value for TestValue {
32 fn bytes(&self) -> &Vec<u8> {
33 &self.bytes
34 }
35 fn index_length(&self) -> u32 {
36 self.index_length
37 }
38 fn hi(&self) -> [u8;32] {
39 utils::hash_vec(self.bytes().to_vec().split_at(self.index_length() as usize).0.to_vec())
40 }
41 fn ht(&self) -> [u8;32] {
42 utils::hash_vec(self.bytes().to_vec())
43 }
44}
45
46#[allow(dead_code)]
47pub struct Db {
48 storage: HashMap<[u8;32], Vec<u8>>,
49}
50impl Db {
51 pub fn insert(&mut self, k: [u8; 32], t: u8, il: u32, b: &mut Vec<u8>) {
52 let mut v: Vec<u8>;
53 v = [t].to_vec();
54 let il_bytes = il.to_le_bytes();
55 v.append(&mut il_bytes.to_vec()); // il_bytes are [u8;4] (4 bytes)
56 v.append(b);
57 self.storage.insert(k, v);
58 }
59 pub fn get(&self, k: &[u8;32]) -> (u8, u32, Vec<u8>) {
60 if k.to_vec() == EMPTYNODEVALUE.to_vec() {
61 return (0, 0, EMPTYNODEVALUE.to_vec());
62 }
63 match self.storage.get(k) {
64 Some(x) => {
65 let t = x[0];
66 let il_bytes: [u8; 4] = [x[1], x[2], x[3], x[4]];
67 let il = u32::from_le_bytes(il_bytes);
68 let b = &x[5..];
69 (t, il, b.to_vec())
70 },
71 None => (TYPENODEEMPTY, 0, EMPTYNODEVALUE.to_vec()),
72 }
73 }
74}
75
76pub fn new_db()-> Db {
77 Db {
78 storage: HashMap::new(),
79 }
80}
81
82pub struct MerkleTree {
83 #[allow(dead_code)]
84 root: [u8; 32],
85 #[allow(dead_code)]
86 num_levels: u32,
87 #[allow(dead_code)]
88 sto: Db,
89}
90
91
92pub fn new(num_levels: u32) -> MerkleTree {
93 MerkleTree {
94 root: EMPTYNODEVALUE,
95 num_levels,
96 sto: new_db(),
97 }
98}
99
100impl MerkleTree {
101 pub fn add(&mut self, v: &TestValue) {
102 #![allow(unused_variables)]
103 #[allow(dead_code)]
104
105 // println!("adding value: {:?}", v.bytes());
106 // add the leaf that we are adding
107 self.sto.insert(v.ht(), TYPENODEVALUE, v.index_length(), &mut v.bytes().to_vec());
108
109 let index = v.index_length() as usize;
110 let hi = v.hi();
111 let ht = v.ht();
112 let path = utils::get_path(self.num_levels, hi);
113 let mut siblings: Vec<[u8;32]> = Vec::new();
114
115
116 let mut node_hash = self.root;
117
118 for i in (0..=self.num_levels-2).rev() {
119 // get node
120 let (t, il, node_bytes) = self.sto.get(&node_hash);
121 if t == TYPENODEFINAL {
122 let hi_child = utils::hash_vec(node_bytes.to_vec().split_at(il as usize).0.to_vec());
123 let path_child = utils::get_path(self.num_levels, hi_child);
124 let pos_diff = utils::compare_paths(&path_child, &path);
125 if pos_diff == 999 { // TODO use a match here, and instead of 999 return something better
126 println!("compare paths err");
127 return;
128 }
129 let final_node_1_hash = utils::calc_hash_from_leaf_and_level(pos_diff, &path_child, utils::hash_vec(node_bytes.to_vec()));
130 self.sto.insert(final_node_1_hash, TYPENODEFINAL, il, &mut node_bytes.to_vec());
131 let final_node_2_hash = utils::calc_hash_from_leaf_and_level(pos_diff, &path, v.ht());
132 self.sto.insert(final_node_2_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
133
134 // parent node
135 let parent_node: node::TreeNode;
136 if path[pos_diff as usize] {
137 parent_node = node::TreeNode {
138 child_l: final_node_1_hash,
139 child_r: final_node_2_hash,
140 }} else {
141 parent_node = node::TreeNode {
142 child_l: final_node_2_hash,
143 child_r: final_node_1_hash,
144 }}
145 let empties = utils::get_empties_between_i_and_pos(i, pos_diff+1);
146 for empty in &empties {
147 siblings.push(*empty);
148 }
149
150 let path_from_pos_diff = utils::cut_path(&path, (pos_diff +1) as usize);
151 self.root = self.replace_leaf(path_from_pos_diff, siblings.clone(), parent_node.ht(), TYPENODENORMAL, 0, &mut parent_node.bytes().to_vec());
152 return;
153 }
154
155 let node = node::parse_node_bytes(node_bytes);
156
157 let sibling: [u8;32];
158 if path[i as usize] {
159 node_hash = node.child_l;
160 sibling = node.child_r;
161 } else {
162 sibling = node.child_l;
163 node_hash = node.child_r;
164 }
165 siblings.push(*array_ref!(sibling, 0, 32));
166 if node_hash == EMPTYNODEVALUE {
167 if i==self.num_levels-2 && siblings[siblings.len()-1]==EMPTYNODEVALUE {
168 let final_node_hash = utils::calc_hash_from_leaf_and_level(i+1, &path, v.ht());
169 self.sto.insert(final_node_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
170 self.root = final_node_hash;
171 return;
172 }
173 let final_node_hash = utils::calc_hash_from_leaf_and_level(i, &path, v.ht());
174 let path_from_i = utils::cut_path(&path, i as usize);
175 self.root = self.replace_leaf(path_from_i, siblings.clone(), final_node_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
176 return;
177 }
178 }
179 self.root = self.replace_leaf(path, siblings, v.ht(), TYPENODEVALUE, v.index_length(), &mut v.bytes().to_vec());
180 }
181
182 #[allow(dead_code)]
183 pub fn replace_leaf(&mut self, path: Vec<bool>, siblings: Vec<[u8;32]>, leaf_hash: [u8;32], node_type: u8, index_length: u32, leaf_value: &mut Vec<u8>) -> [u8;32] {
184 self.sto.insert(leaf_hash, node_type, index_length, leaf_value);
185 let mut curr_node = leaf_hash;
186
187 for i in 0..siblings.len() {
188 if path[i as usize] {
189 let node = node::TreeNode {
190 child_l: curr_node,
191 child_r: siblings[siblings.len()-1-i],
192 };
193 self.sto.insert(node.ht(), TYPENODENORMAL, 0, &mut node.bytes());
194 curr_node = node.ht();
195 } else {
196 let node = node::TreeNode {
197 child_l: siblings[siblings.len()-1-i],
198 child_r: curr_node,
199 };
200 self.sto.insert(node.ht(), TYPENODENORMAL, 0, &mut node.bytes());
201 curr_node = node.ht();
202 }
203 }
204 curr_node
205 }
206
207 #[allow(dead_code)]
208 pub fn get_value_in_pos(&self, hi: [u8;32]) -> Vec<u8> {
209 let path = utils::get_path(self.num_levels, hi);
210 let mut node_hash = self.root;
211 for i in (0..=self.num_levels-2).rev() {
212 let (t, il, node_bytes) = self.sto.get(&node_hash);
213 if t == TYPENODEFINAL {
214 let hi_node = utils::hash_vec(node_bytes.to_vec().split_at(il as usize).0.to_vec());
215 let path_node = utils::get_path(self.num_levels, hi_node);
216 let pos_diff = utils::compare_paths(&path_node, &path);
217 // if pos_diff > self.num_levels {
218 if pos_diff != 999 {
219 return EMPTYNODEVALUE.to_vec();
220 }
221 return node_bytes;
222 }
223 let node = node::parse_node_bytes(node_bytes);
224 if !path[i as usize] {
225 node_hash = node.child_l;
226 } else {
227 node_hash = node.child_r;
228 }
229 }
230 let (_t, _il, node_bytes) = self.sto.get(&node_hash);
231 node_bytes
232 }
233
234 #[allow(dead_code)]
235 pub fn generate_proof(&self, hi: [u8;32]) -> Vec<u8> {
236 let mut mp: Vec<u8> = Vec::new();
237
238 let mut empties: [u8;32] = [0;32];
239 let path = utils::get_path(self.num_levels, hi);
240
241 let mut siblings: Vec<[u8;32]> = Vec::new();
242 let mut node_hash = self.root;
243
244 for i in 0..self.num_levels {
245 let (t, il, node_bytes) = self.sto.get(&node_hash);
246 if t == TYPENODEFINAL {
247 let real_value_in_pos = self.get_value_in_pos(hi);
248 if real_value_in_pos == EMPTYNODEVALUE {
249 let leaf_hi = utils::hash_vec(node_bytes.to_vec().split_at(il as usize).0.to_vec());
250 let path_child = utils::get_path(self.num_levels, leaf_hi);
251 let pos_diff = utils::compare_paths(&path_child, &path);
252 if pos_diff == self.num_levels { // TODO use a match here, and instead of 999 return something better
253 return mp;
254 }
255 if pos_diff != self.num_levels-1-i {
256 let sibling = utils::calc_hash_from_leaf_and_level(pos_diff, &path_child, utils::hash_vec(node_bytes.to_vec()));
257 let mut new_siblings: Vec<[u8;32]> = Vec::new();
258 new_siblings.push(sibling);
259 new_siblings.append(&mut siblings);
260 siblings = new_siblings;
261 // set empties bit
262 let bit_pos = self.num_levels-2-pos_diff;
263 empties[(empties.len() as isize + (bit_pos as isize/8-1) as isize) as usize] |= 1 << (bit_pos%8);
264 }
265 }
266 break
267 }
268 let node = node::parse_node_bytes(node_bytes);
269 let sibling: [u8;32];
270 if !path[self.num_levels as usize -i as usize -2] {
271 node_hash = node.child_l;
272 sibling = node.child_r;
273 } else {
274 sibling = node.child_l;
275 node_hash = node.child_r;
276 }
277 if sibling != EMPTYNODEVALUE {
278 // set empties bit
279 empties[(empties.len() as isize + (i as isize/8-1) as isize) as usize] |= 1 << (i%8);
280 let mut new_siblings: Vec<[u8;32]> = Vec::new();
281 new_siblings.push(sibling);
282 new_siblings.append(&mut siblings);
283 siblings = new_siblings;
284 }
285 }
286 mp.append(&mut empties[..].to_vec());
287 for s in siblings {
288 mp.append(&mut s.to_vec());
289 }
290 mp
291 }
292 #[allow(dead_code)]
293 pub fn print_level(&self, parent: [u8;32], mut lvl: u32, max_level: u32) {
294 use rustc_hex::ToHex;
295
296 let mut line: String = "".to_string();
297 for _ in 0..lvl {
298 line += &" ".to_string();
299 }
300 line += &("lvl ".to_string() + &lvl.to_string());
301 line += &(" - '".to_string() + &parent.to_hex() + &"' = ".to_string());
302 let (t, _, node_bytes) = self.sto.get(&parent);
303 let mut node = node::TreeNode {
304 child_l: EMPTYNODEVALUE,
305 child_r: EMPTYNODEVALUE,
306 };
307 if t==TYPENODENORMAL {
308 node = node::parse_node_bytes(node_bytes);
309 line += &("'".to_string() + &node.child_l.to_hex() + &"' - '".to_string() + &node.child_r.to_hex() + &"'".to_string());
310 } else if t == TYPENODEVALUE {
311 //
312 } else if t == TYPENODEFINAL {
313 let hash_node_bytes = utils::hash_vec(node_bytes);
314 line += &("[final] final tree node: ".to_string() + &hash_node_bytes.to_hex()+ &"\n".to_string());
315 let (_, _, leaf_node_bytes) = self.sto.get(&hash_node_bytes);
316 for _ in 0..lvl {
317 line += " ";
318 }
319 let leaf_node_string = String::from_utf8_lossy(&leaf_node_bytes);
320 line += &("leaf value: ".to_string() + &leaf_node_string);
321 } else {
322 line += &"[EMPTY Branch]".to_string()
323 }
324 println!("{:?}", line);
325 lvl += 1;
326 if node.child_r.len()>0 && lvl<max_level && t != TYPENODEEMPTY && t != TYPENODEFINAL {
327 self.print_level(node.child_l, lvl, max_level);
328 self.print_level(node.child_r, lvl, max_level);
329 }
330 }
331 pub fn print_full_tree(&self) {
332 use rustc_hex::ToHex;
333 self.print_level(self.root, 0, self.num_levels - 1);
334 println!("root {:?}", self.root.to_hex());
335 }
336 pub fn print_levels_tree(&self, max_level: u32) {
337 use rustc_hex::ToHex;
338 self.print_level(self.root, 0, self.num_levels - 1 - max_level);
339 println!("root {:?}", self.root.to_hex());
340 }
341}
342#[allow(dead_code)]
343pub fn verify_proof(root: [u8;32], mp: Vec<u8>, hi: [u8;32], ht: [u8;32], num_levels: u32) -> bool {
344 let empties: Vec<u8>;
345 empties = mp.split_at(32).0.to_vec();
346
347 let mut siblings: Vec<[u8;32]> = Vec::new();
348 for i in (empties.len()..mp.len()).step_by(EMPTYNODEVALUE.len()) {
349 let mut sibling: [u8;32] = [0;32];
350 sibling.copy_from_slice(&mp[i..i+EMPTYNODEVALUE.len()]);
351 siblings.push(sibling);
352 }
353
354 let path = utils::get_path(num_levels, hi);
355 let mut node_hash = ht;
356 let mut sibling_used_pos = 0;
357
358 for i in (0..=num_levels-2).rev() {
359 let sibling: [u8;32];
360 if (empties[empties.len()-i as usize/8-1] & (1 << (i%8))) > 0 {
361 sibling = siblings[sibling_used_pos];
362 sibling_used_pos += 1;
363 } else {
364 sibling = EMPTYNODEVALUE;
365 }
366
367 let n: node::TreeNode;
368 if path[num_levels as usize - i as usize-2] {
369 n = node::TreeNode {
370 child_l: sibling,
371 child_r: node_hash,
372 }} else {
373 n = node::TreeNode {
374 child_l: node_hash,
375 child_r: sibling,
376 }}
377 if node_hash == EMPTYNODEVALUE && sibling == EMPTYNODEVALUE {
378 node_hash = EMPTYNODEVALUE;
379 } else {
380 node_hash = n.ht();
381 }
382 }
383 if node_hash==root {
384 return true;
385 }
386 false
387}
388
389
390#[cfg(test)]
391 mod tests {
392 use super::*;
393 use rustc_hex::ToHex;
394
395 #[test]
396 fn test_hash_vec() {
397 let a: Vec<u8> = From::from("test");
398 let h = utils::hash_vec(a);
399 assert_eq!("9c22ff5f21f0b81b113e63f7db6da94fedef11b2119b4088b89664fb9a3cb658", h.to_hex());
400 }
401
402 #[test]
403 fn test_new_mt() {
404 let mt: MerkleTree = new(140);
405 assert_eq!(140, mt.num_levels);
406 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
407 let (_t, _il, b) = mt.sto.get(&[0;32]);
408 assert_eq!(mt.root.to_vec(), b);
409 }
410
411
412 #[test]
413 fn test_tree_node() {
414 let n = node::TreeNode {
415 child_l: [1;32],
416 child_r: [2;32],
417 };
418 assert_eq!("01010101010101010101010101010101010101010101010101010101010101010202020202020202020202020202020202020202020202020202020202020202",
419 n.bytes().to_hex());
420 assert_eq!("346d8c96a2454213fcc0daff3c96ad0398148181b9fa6488f7ae2c0af5b20aa0", n.ht().to_hex());
421 }
422
423 #[test]
424 fn test_add() {
425 let mut mt: MerkleTree = new(140);
426 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
427 let val = TestValue {
428 bytes: vec![1,2,3,4,5],
429 index_length: 3,
430 };
431 mt.add(&val);
432 let (_t, _il, b) = mt.sto.get(&val.ht());
433 assert_eq!(*val.bytes(), b);
434 assert_eq!("a0e72cc948119fcb71b413cf5ada12b2b825d5133299b20a6d9325ffc3e2fbf1", mt.root.to_hex());
435 }
436 #[test]
437 fn test_add_2() {
438 let mut mt: MerkleTree = new(140);
439 let val = TestValue {
440 bytes: "this is a test leaf".as_bytes().to_vec(),
441 index_length: 15,
442 };
443 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
444 mt.add(&val);
445 let (_t, _il, b) = mt.sto.get(&val.ht());
446 assert_eq!(*val.bytes(), b);
447 assert_eq!("b4fdf8a653198f0e179ccb3af7e4fc09d76247f479d6cfc95cd92d6fda589f27", mt.root.to_hex());
448 let val2 = TestValue {
449 bytes: "this is a second test leaf".as_bytes().to_vec(),
450 index_length: 15,
451 };
452 mt.add(&val2);
453 let (_t, _il, b) = mt.sto.get(&val2.ht());
454 assert_eq!(*val2.bytes(), b);
455 assert_eq!("8ac95e9c8a6fbd40bb21de7895ee35f9c8f30ca029dbb0972c02344f49462e82", mt.root.to_hex());
456 mt.print_full_tree();
457 }
458 #[test]
459 fn test_generate_proof_and_verify_proof() {
460 let mut mt: MerkleTree = new(140);
461 let val = TestValue {
462 bytes: "this is a test leaf".as_bytes().to_vec(),
463 index_length: 15,
464 };
465 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
466 mt.add(&val);
467 let (_t, _il, b) = mt.sto.get(&val.ht());
468 assert_eq!(*val.bytes(), b);
469 assert_eq!("b4fdf8a653198f0e179ccb3af7e4fc09d76247f479d6cfc95cd92d6fda589f27", mt.root.to_hex());
470 let val2 = TestValue {
471 bytes: "this is a second test leaf".as_bytes().to_vec(),
472 index_length: 15,
473 };
474 mt.add(&val2);
475 let (_t, _il, b) = mt.sto.get(&val2.ht());
476 assert_eq!(*val2.bytes(), b);
477 assert_eq!("8ac95e9c8a6fbd40bb21de7895ee35f9c8f30ca029dbb0972c02344f49462e82", mt.root.to_hex());
478
479 let mp = mt.generate_proof(val2.hi());
480 assert_eq!("0000000000000000000000000000000000000000000000000000000000000001fd8e1a60cdb23c0c7b2cf8462c99fafd905054dccb0ed75e7c8a7d6806749b6b", mp.to_hex());
481
482 // verify
483 let v = verify_proof(mt.root, mp, val2.hi(), val2.ht(), mt.num_levels);
484 assert_eq!(true, v);
485 }
486 #[test]
487 fn test_generate_proof_empty_leaf_and_verify_proof() {
488 let mut mt: MerkleTree = new(140);
489 let val = TestValue {
490 bytes: "this is a test leaf".as_bytes().to_vec(),
491 index_length: 15,
492 };
493 mt.add(&val);
494 let val2 = TestValue {
495 bytes: "this is a second test leaf".as_bytes().to_vec(),
496 index_length: 15,
497 };
498 mt.add(&val2);
499 assert_eq!("8ac95e9c8a6fbd40bb21de7895ee35f9c8f30ca029dbb0972c02344f49462e82", mt.root.to_hex());
500
501 // proof of empty leaf
502 let val3 = TestValue {
503 bytes: "this is a third test leaf".as_bytes().to_vec(),
504 index_length: 15,
505 };
506 let mp = mt.generate_proof(val3.hi());
507 assert_eq!("000000000000000000000000000000000000000000000000000000000000000389741fa23da77c259781ad8f4331a5a7d793eef1db7e5200ddfc8e5f5ca7ce2bfd8e1a60cdb23c0c7b2cf8462c99fafd905054dccb0ed75e7c8a7d6806749b6b", mp.to_hex());
508
509 // verify that is a proof of an empty leaf (EMPTYNODEVALUE)
510 let v = verify_proof(mt.root, mp, val3.hi(), EMPTYNODEVALUE, mt.num_levels);
511 assert_eq!(true, v);
512 }
513 #[test]
514 fn test_harcoded_proofs_of_existing_leaf() {
515 // check proof of value in leaf
516 let mut root: [u8;32] = [0;32];
517 root.copy_from_slice(&hex::decode("7d7c5e8f4b3bf434f3d9d223359c4415e2764dd38de2e025fbf986e976a7ed3d").unwrap());
518 let mp = hex::decode("0000000000000000000000000000000000000000000000000000000000000002d45aada6eec346222eaa6b5d3a9260e08c9b62fcf63c72bc05df284de07e6a52").unwrap();
519 let mut hi: [u8;32] = [0;32];
520 hi.copy_from_slice(&hex::decode("786677808ba77bdd9090a969f1ef2cbd1ac5aecd9e654f340500159219106878").unwrap());
521 let mut ht: [u8;32] = [0;32];
522 ht.copy_from_slice(&hex::decode("786677808ba77bdd9090a969f1ef2cbd1ac5aecd9e654f340500159219106878").unwrap());
523 let v = verify_proof(root, mp, hi, ht, 140);
524 assert_eq!(true, v);
525 }
526 #[test]
527 fn test_harcoded_proofs_of_empty_leaf() {
528 // check proof of value in leaf
529 let mut root: [u8;32] = [0;32];
530 root.copy_from_slice(&hex::decode("8f021d00c39dcd768974ddfe0d21f5d13f7215bea28db1f1cb29842b111332e7").unwrap());
531 let mp = hex::decode("0000000000000000000000000000000000000000000000000000000000000004bf8e980d2ed328ae97f65c30c25520aeb53ff837579e392ea1464934c7c1feb9").unwrap();
532 let mut hi: [u8;32] = [0;32];
533 hi.copy_from_slice(&hex::decode("a69792a4cff51f40b7a1f7ae596c6ded4aba241646a47538898f17f2a8dff647").unwrap());
534 let v = verify_proof(root, mp, hi, EMPTYNODEVALUE, 140);
535 assert_eq!(true, v);
536 }
537 }