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}
293#[allow(dead_code)]
294pub fn verify_proof(root: [u8;32], mp: Vec<u8>, hi: [u8;32], ht: [u8;32], num_levels: u32) -> bool {
295 let empties: Vec<u8>;
296 empties = mp.split_at(32).0.to_vec();
297
298 let mut siblings: Vec<[u8;32]> = Vec::new();
299 for i in (empties.len()..mp.len()).step_by(EMPTYNODEVALUE.len()) {
300 let mut sibling: [u8;32] = [0;32];
301 sibling.copy_from_slice(&mp[i..i+EMPTYNODEVALUE.len()]);
302 siblings.push(sibling);
303 }
304
305 let path = utils::get_path(num_levels, hi);
306 let mut node_hash = ht;
307 let mut sibling_used_pos = 0;
308
309 for i in (0..=num_levels-2).rev() {
310 let sibling: [u8;32];
311 if (empties[empties.len()-i as usize/8-1] & (1 << (i%8))) > 0 {
312 sibling = siblings[sibling_used_pos];
313 sibling_used_pos += 1;
314 } else {
315 sibling = EMPTYNODEVALUE;
316 }
317
318 let n: node::TreeNode;
319 if path[num_levels as usize - i as usize-2] {
320 n = node::TreeNode {
321 child_l: sibling,
322 child_r: node_hash,
323 }} else {
324 n = node::TreeNode {
325 child_l: node_hash,
326 child_r: sibling,
327 }}
328 if node_hash == EMPTYNODEVALUE && sibling == EMPTYNODEVALUE {
329 node_hash = EMPTYNODEVALUE;
330 } else {
331 node_hash = n.ht();
332 }
333 }
334 if node_hash==root {
335 return true;
336 }
337 false
338}
339
340
341#[cfg(test)]
342 mod tests {
343 use super::*;
344 use rustc_hex::ToHex;
345
346 #[test]
347 fn test_hash_vec() {
348 let a: Vec<u8> = From::from("test");
349 let h = utils::hash_vec(a);
350 assert_eq!("9c22ff5f21f0b81b113e63f7db6da94fedef11b2119b4088b89664fb9a3cb658", h.to_hex());
351 }
352
353 #[test]
354 fn test_new_mt() {
355 let mt: MerkleTree = new(140);
356 assert_eq!(140, mt.num_levels);
357 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
358 let (_t, _il, b) = mt.sto.get(&[0;32]);
359 assert_eq!(mt.root.to_vec(), b);
360 }
361
362
363 #[test]
364 fn test_tree_node() {
365 let n = node::TreeNode {
366 child_l: [1;32],
367 child_r: [2;32],
368 };
369 assert_eq!("01010101010101010101010101010101010101010101010101010101010101010202020202020202020202020202020202020202020202020202020202020202",
370 n.bytes().to_hex());
371 assert_eq!("346d8c96a2454213fcc0daff3c96ad0398148181b9fa6488f7ae2c0af5b20aa0", n.ht().to_hex());
372 }
373
374 #[test]
375 fn test_add() {
376 let mut mt: MerkleTree = new(140);
377 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
378 let val = TestValue {
379 bytes: vec![1,2,3,4,5],
380 index_length: 3,
381 };
382 mt.add(&val);
383 let (_t, _il, b) = mt.sto.get(&val.ht());
384 assert_eq!(*val.bytes(), b);
385 assert_eq!("a0e72cc948119fcb71b413cf5ada12b2b825d5133299b20a6d9325ffc3e2fbf1", mt.root.to_hex());
386 }
387 #[test]
388 fn test_add_2() {
389 let mut mt: MerkleTree = new(140);
390 let val = TestValue {
391 bytes: "this is a test leaf".as_bytes().to_vec(),
392 index_length: 15,
393 };
394 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
395 mt.add(&val);
396 let (_t, _il, b) = mt.sto.get(&val.ht());
397 assert_eq!(*val.bytes(), b);
398 assert_eq!("b4fdf8a653198f0e179ccb3af7e4fc09d76247f479d6cfc95cd92d6fda589f27", mt.root.to_hex());
399 let val2 = TestValue {
400 bytes: "this is a second test leaf".as_bytes().to_vec(),
401 index_length: 15,
402 };
403 mt.add(&val2);
404 let (_t, _il, b) = mt.sto.get(&val2.ht());
405 assert_eq!(*val2.bytes(), b);
406 assert_eq!("8ac95e9c8a6fbd40bb21de7895ee35f9c8f30ca029dbb0972c02344f49462e82", mt.root.to_hex());
407 }
408 #[test]
409 fn test_generate_proof_and_verify_proof() {
410 let mut mt: MerkleTree = new(140);
411 let val = TestValue {
412 bytes: "this is a test leaf".as_bytes().to_vec(),
413 index_length: 15,
414 };
415 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
416 mt.add(&val);
417 let (_t, _il, b) = mt.sto.get(&val.ht());
418 assert_eq!(*val.bytes(), b);
419 assert_eq!("b4fdf8a653198f0e179ccb3af7e4fc09d76247f479d6cfc95cd92d6fda589f27", mt.root.to_hex());
420 let val2 = TestValue {
421 bytes: "this is a second test leaf".as_bytes().to_vec(),
422 index_length: 15,
423 };
424 mt.add(&val2);
425 let (_t, _il, b) = mt.sto.get(&val2.ht());
426 assert_eq!(*val2.bytes(), b);
427 assert_eq!("8ac95e9c8a6fbd40bb21de7895ee35f9c8f30ca029dbb0972c02344f49462e82", mt.root.to_hex());
428
429 let mp = mt.generate_proof(val2.hi());
430 assert_eq!("0000000000000000000000000000000000000000000000000000000000000001fd8e1a60cdb23c0c7b2cf8462c99fafd905054dccb0ed75e7c8a7d6806749b6b", mp.to_hex());
431
432 // verify
433 let v = verify_proof(mt.root, mp, val2.hi(), val2.ht(), mt.num_levels);
434 assert_eq!(true, v);
435 }
436 #[test]
437 fn test_generate_proof_empty_leaf_and_verify_proof() {
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 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 (EMPTYNODEVALUE)
460 let v = verify_proof(mt.root, mp, val3.hi(), EMPTYNODEVALUE, mt.num_levels);
461 assert_eq!(true, v);
462 }
463 #[test]
464 fn test_harcoded_proofs_of_existing_leaf() {
465 // check proof of value in leaf
466 let mut root: [u8;32] = [0;32];
467 root.copy_from_slice(&hex::decode("7d7c5e8f4b3bf434f3d9d223359c4415e2764dd38de2e025fbf986e976a7ed3d").unwrap());
468 let mp = hex::decode("0000000000000000000000000000000000000000000000000000000000000002d45aada6eec346222eaa6b5d3a9260e08c9b62fcf63c72bc05df284de07e6a52").unwrap();
469 let mut hi: [u8;32] = [0;32];
470 hi.copy_from_slice(&hex::decode("786677808ba77bdd9090a969f1ef2cbd1ac5aecd9e654f340500159219106878").unwrap());
471 let mut ht: [u8;32] = [0;32];
472 ht.copy_from_slice(&hex::decode("786677808ba77bdd9090a969f1ef2cbd1ac5aecd9e654f340500159219106878").unwrap());
473 let v = verify_proof(root, mp, hi, ht, 140);
474 assert_eq!(true, v);
475 }
476 #[test]
477 fn test_harcoded_proofs_of_empty_leaf() {
478 // check proof of value in leaf
479 let mut root: [u8;32] = [0;32];
480 root.copy_from_slice(&hex::decode("8f021d00c39dcd768974ddfe0d21f5d13f7215bea28db1f1cb29842b111332e7").unwrap());
481 let mp = hex::decode("0000000000000000000000000000000000000000000000000000000000000004bf8e980d2ed328ae97f65c30c25520aeb53ff837579e392ea1464934c7c1feb9").unwrap();
482 let mut hi: [u8;32] = [0;32];
483 hi.copy_from_slice(&hex::decode("a69792a4cff51f40b7a1f7ae596c6ded4aba241646a47538898f17f2a8dff647").unwrap());
484 let v = verify_proof(root, mp, hi, EMPTYNODEVALUE, 140);
485 assert_eq!(true, v);
486 }
487 }