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