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