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