1use crate::crh::FixedLengthCRH;
  2use algebra::bytes::ToBytes;
  3use crate::Error;
  4use std::{fmt, rc::Rc};
  5
  6
  7#[cfg(feature = "r1cs")]
  8pub mod constraints;
  9
 10pub trait MHTParameters {
 11    const HEIGHT: usize;
 12    type H: FixedLengthCRH;
 13}
 14
 15/// Stores the hashes of a particular path (in order) from leaf to root.
 16/// Our path `is_left_child()` if the boolean in `path` is true.
 17#[derive(Derivative)]
 18#[derivative(
 19    Clone(bound = "P: MHTParameters"),
 20    Debug(bound = "P: MHTParameters, <P::H as FixedLengthCRH>::Output: fmt::Debug")
 21)]
 22pub struct HashMembershipProof<P: MHTParameters> {
 23    pub(crate) path: Vec<(<P::H as FixedLengthCRH>::Output, <P::H as FixedLengthCRH>::Output)>,
 24}
 25
 26impl<P: MHTParameters> Default for HashMembershipProof<P> {
 27    fn default() -> Self {
 28        let mut path = Vec::with_capacity(P::HEIGHT as usize);
 29        for _i in 1..P::HEIGHT as usize {
 30            path.push((<P::H as FixedLengthCRH>::Output::default(), <P::H as FixedLengthCRH>::Output::default()));
 31        }
 32        Self {
 33            path,
 34        }
 35    }
 36}
 37
 38impl<P: MHTParameters> HashMembershipProof<P> {
 39    pub fn verify<L: ToBytes>(
 40        &self,
 41        parameters: &<P::H as FixedLengthCRH>::Parameters,
 42        root_hash: &<P::H as FixedLengthCRH>::Output,
 43        leaf: &L,
 44    ) -> Result<bool, Error> {
 45        if self.path.len() != (P::HEIGHT - 1) as usize {
 46            return Ok(false);
 47        }
 48        // Check that the given leaf matches the leaf in the membership proof.
 49        let mut buffer = [0u8; 128];
 50
 51        if !self.path.is_empty() {
 52            let claimed_leaf_hash = hash_leaf::<P::H, L>(parameters, leaf, &mut buffer)?;
 53
 54            // Check if leaf is one of the bottom-most siblings.
 55            if claimed_leaf_hash != self.path[0].0 && claimed_leaf_hash != self.path[0].1 {
 56                return Ok(false);
 57            };
 58
 59            let mut prev = claimed_leaf_hash;
 60            // Check levels between leaf level and root.
 61            for &(ref hash, ref sibling_hash) in &self.path {
 62                // Check if the previous hash matches the correct current hash.
 63                if &prev != hash && &prev != sibling_hash {
 64                    return Ok(false);
 65                };
 66                prev = hash_inner_node::<P::H>(parameters, hash, sibling_hash, &mut buffer)?;
 67            }
 68
 69            if root_hash != &prev {
 70                return Ok(false);
 71            }
 72            Ok(true)
 73        } else {
 74            Ok(false)
 75        }
 76    }
 77}
 78
 79pub struct MerkleHashTree<P: MHTParameters> {
 80    tree:         Vec<<P::H as FixedLengthCRH>::Output>,
 81    padding_tree: Vec<(<P::H as FixedLengthCRH>::Output, <P::H as FixedLengthCRH>::Output)>,
 82    parameters:   Rc<<P::H as FixedLengthCRH>::Parameters>,
 83    root:         Option<<P::H as FixedLengthCRH>::Output>,
 84}
 85
 86impl<P: MHTParameters> MerkleHashTree<P> {
 87    pub const HEIGHT: u8 = P::HEIGHT as u8;
 88
 89    pub fn blank(parameters: Rc<<P::H as FixedLengthCRH>::Parameters>) -> Self {
 90        MerkleHashTree {
 91            tree: Vec::new(),
 92            padding_tree: Vec::new(),
 93            root: None,
 94            parameters,
 95        }
 96    }
 97
 98    pub fn new<L: ToBytes>(parameters: Rc<<P::H as FixedLengthCRH>::Parameters>, leaves: &[L]) -> Result<Self, Error> {
 99        let new_time = start_timer!(|| "MHT::New");
100
101        let last_level_size = leaves.len().next_power_of_two();
102        let tree_size = 2 * last_level_size - 1;
103        let tree_height = tree_height(tree_size);
104        assert!(tree_height as u8 <= Self::HEIGHT);
105
106        // Initialize the merkle tree.
107        let mut tree = Vec::with_capacity(tree_size);
108        let empty_hash = hash_empty::<P::H>(&parameters)?;
109        for _ in 0..tree_size {
110            tree.push(empty_hash.clone());
111        }
112
113        // Compute the starting indices for each level of the tree.
114        let mut index = 0;
115        let mut level_indices = Vec::with_capacity(tree_height);
116        for _ in 0..tree_height {
117            level_indices.push(index);
118            index = left_child(index);
119        }
120
121        // Compute and store the hash values for each leaf.
122        let last_level_index = level_indices.pop().unwrap();
123        let mut buffer = [0u8; 128];
124        for (i, leaf) in leaves.iter().enumerate() {
125            tree[last_level_index + i] = hash_leaf::<P::H, _>(&parameters, leaf, &mut buffer)?;
126        }
127
128        // Compute the hash values for every node in the tree.
129        let mut upper_bound = last_level_index;
130        let mut buffer = [0u8; 128];
131        level_indices.reverse();
132        for &start_index in &level_indices {
133            // Iterate over the current level.
134            for current_index in start_index..upper_bound {
135                let left_index = left_child(current_index);
136                let right_index = right_child(current_index);
137
138                // Compute Hash(left || right).
139                tree[current_index] = hash_inner_node::<P::H>(
140                    &parameters,
141                    &tree[left_index],
142                    &tree[right_index],
143                    &mut buffer,
144                )?;
145            }
146            upper_bound = start_index;
147        }
148        // Finished computing actual tree.
149        // Now, we compute the dummy nodes until we hit our HEIGHT goal.
150        let mut cur_height = tree_height;
151        let mut padding_tree = Vec::new();
152        let mut cur_hash = tree[0].clone();
153        while cur_height < (Self::HEIGHT - 1) as usize {
154            cur_hash = hash_inner_node::<P::H>(&parameters, &cur_hash, &empty_hash, &mut buffer)?;
155            padding_tree.push((cur_hash.clone(), empty_hash.clone()));
156            cur_height += 1;
157        }
158
159        let root_hash = hash_inner_node::<P::H>(&parameters, &cur_hash, &empty_hash, &mut buffer)?;
160
161        end_timer!(new_time);
162
163        Ok(MerkleHashTree {
164            tree,
165            padding_tree,
166            parameters,
167            root: Some(root_hash),
168        })
169    }
170
171    #[inline]
172    pub fn root(&self) -> <P::H as FixedLengthCRH>::Output {
173        self.root.clone().unwrap()
174    }
175
176    pub fn generate_proof<L: ToBytes>(
177        &self,
178        index: usize,
179        leaf: &L,
180    ) -> Result<HashMembershipProof<P>, Error> {
181        let prove_time = start_timer!(|| "MHT::GenProof");
182        let mut path = Vec::new();
183
184        let mut buffer = [0u8; 128];
185        let leaf_hash = hash_leaf::<P::H, _>(&self.parameters, leaf, &mut buffer)?;
186        let tree_height = tree_height(self.tree.len());
187        let tree_index = convert_index_to_last_level(index, tree_height);
188        let empty_hash = hash_empty::<P::H>(&self.parameters)?;
189
190        // Check that the given index corresponds to the correct leaf.
191        if leaf_hash != self.tree[tree_index] {
192            Err(MHTError::IncorrectLeafIndex(tree_index))?
193        }
194
195        // Iterate from the leaf up to the root, storing all intermediate hash values.
196        let mut current_node = tree_index;
197        while !is_root(current_node) {
198            let sibling_node = sibling(current_node).unwrap();
199            let (curr_hash, sibling_hash) = (
200                self.tree[current_node].clone(),
201                self.tree[sibling_node].clone(),
202            );
203            if is_left_child(current_node) {
204                path.push((curr_hash, sibling_hash));
205            } else {
206                path.push((sibling_hash, curr_hash));
207            }
208            current_node = parent(current_node).unwrap();
209        }
210
211        // Store the root node. Set boolean as true for consistency with digest
212        // location.
213        assert!(path.len() < Self::HEIGHT as usize);
214        if path.len() != (Self::HEIGHT - 1) as usize {
215            path.push((self.tree[0].clone(), empty_hash));
216            for &(ref hash, ref sibling_hash) in &self.padding_tree {
217                path.push((hash.clone(), sibling_hash.clone()));
218            }
219        }
220        end_timer!(prove_time);
221        if path.len() != (Self::HEIGHT - 1) as usize {
222            Err(MHTError::IncorrectPathLength(path.len()))?
223        } else {
224            Ok(HashMembershipProof {
225                path,
226            })
227        }
228    }
229}
230
231#[derive(Debug)]
232pub enum MHTError {
233    IncorrectLeafIndex(usize),
234    IncorrectPathLength(usize),
235}
236
237impl std::fmt::Display for MHTError {
238    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
239        let msg = match self {
240            MHTError::IncorrectLeafIndex(index) => format!("incorrect leaf index: {}", index),
241            MHTError::IncorrectPathLength(len) => format!("incorrect path length: {}", len),
242        };
243        write!(f, "{}", msg)
244    }
245}
246
247impl std::error::Error for MHTError {
248    #[inline]
249    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
250        None
251    }
252}
253
254/// Returns the log2 value of the given number.
255#[inline]
256fn log2(number: usize) -> usize {
257    (number as f64).log2() as usize
258}
259
260/// Returns the height of the tree, given the size of the tree.
261#[inline]
262fn tree_height(tree_size: usize) -> usize {
263    log2(tree_size + 1)
264}
265
266/// Returns true iff the index represents the root.
267#[inline]
268fn is_root(index: usize) -> bool {
269    index == 0
270}
271
272/// Returns the index of the left child, given an index.
273#[inline]
274fn left_child(index: usize) -> usize {
275    2 * index + 1
276}
277
278/// Returns the index of the right child, given an index.
279#[inline]
280fn right_child(index: usize) -> usize {
281    2 * index + 2
282}
283
284/// Returns the index of the sibling, given an index.
285#[inline]
286fn sibling(index: usize) -> Option<usize> {
287    if index == 0 {
288        None
289    } else if is_left_child(index) {
290        Some(index + 1)
291    } else {
292        Some(index - 1)
293    }
294}
295
296/// Returns true iff the given index represents a left child.
297#[inline]
298fn is_left_child(index: usize) -> bool {
299    index % 2 == 1
300}
301
302/// Returns the index of the parent, given an index.
303#[inline]
304fn parent(index: usize) -> Option<usize> {
305    if index > 0 {
306        Some((index - 1) >> 1)
307    } else {
308        None
309    }
310}
311
312#[inline]
313fn convert_index_to_last_level(index: usize, tree_height: usize) -> usize {
314    index + (1 << (tree_height - 1)) - 1
315}
316
317/// Returns the output hash, given a left and right hash value.
318pub(crate) fn hash_inner_node<H: FixedLengthCRH>(
319    parameters: &H::Parameters,
320    left: &H::Output,
321    right: &H::Output,
322    buffer: &mut [u8],
323) -> Result<H::Output, Error> {
324    use std::io::Cursor;
325    let mut writer = Cursor::new(buffer);
326    // Construct left input.
327    left.write(&mut writer)?;
328
329    // Construct right input.
330    right.write(&mut writer)?;
331
332    let buffer = writer.into_inner();
333    H::evaluate(parameters, &buffer[..(H::INPUT_SIZE_BITS / 8)])
334}
335
336/// Returns the hash of a leaf.
337pub(crate) fn hash_leaf<H: FixedLengthCRH, L: ToBytes>(
338    parameters: &H::Parameters,
339    leaf: &L,
340    buffer: &mut [u8],
341) -> Result<H::Output, Error> {
342    use std::io::Cursor;
343    let mut writer = Cursor::new(buffer);
344    leaf.write(&mut writer)?;
345
346    let buffer = writer.into_inner();
347    H::evaluate(parameters, &buffer[..(H::INPUT_SIZE_BITS / 8)])
348}
349
350pub(crate) fn hash_empty<H: FixedLengthCRH>(
351    parameters: &H::Parameters,
352) -> Result<H::Output, Error> {
353    let empty_buffer = vec![0u8; H::INPUT_SIZE_BITS / 8];
354    H::evaluate(parameters, &empty_buffer)
355}
356
357
358#[cfg(test)]
359mod test {
360    use crate::{crh::{pedersen::*, *}, mht::*};
361    use algebra::curves::jubjub::JubJubAffine as JubJub;
362    use rand::SeedableRng;
363    use rand_xorshift::XorShiftRng;
364
365    #[derive(Clone)]
366    pub(super) struct Window4x256;
367    impl PedersenWindow for Window4x256 {
368        const WINDOW_SIZE: usize = 4;
369        const NUM_WINDOWS: usize = 256;
370    }
371
372
373    type H = PedersenCRH<JubJub, Window4x256>;
374
375    struct JubJubMHTParams;
376    
377    impl MHTParameters for JubJubMHTParams {
378        const HEIGHT: usize = 32;
379        type H = H;
380    }
381    type JubJubMHT = MerkleHashTree<JubJubMHTParams>;
382
383
384    fn generate_merkle_tree<L: ToBytes + Clone + Eq>(leaves: &[L]) -> () {
385        let mut rng = XorShiftRng::seed_from_u64(9174123u64);
386
387        let crh_parameters = Rc::new(H::setup(&mut rng).unwrap());
388        let tree = JubJubMHT::new(crh_parameters.clone(), &leaves).unwrap();
389        let root = tree.root();
390        for (i, leaf) in leaves.iter().enumerate() {
391            let proof = tree.generate_proof(i, &leaf).unwrap();
392            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
393        }
394    }
395
396    #[test]
397    fn mht_test() {
398        let mut leaves = Vec::new();
399        for i in 0..4u8 {
400            leaves.push([i, i, i, i, i, i, i, i]);
401        }
402        generate_merkle_tree(&leaves);
403        let mut leaves = Vec::new();
404        for i in 0..100u8 {
405            leaves.push([i, i, i, i, i, i, i, i]);
406        }
407        generate_merkle_tree(&leaves);
408    }
409
410    fn bad_merkle_tree_verify<L: ToBytes + Clone + Eq>(leaves: &[L]) -> () {
411        use algebra::groups::Group;
412        let mut rng = XorShiftRng::seed_from_u64(13423423u64);
413
414        let crh_parameters = Rc::new(H::setup(&mut rng).unwrap());
415        let tree = JubJubMHT::new(crh_parameters.clone(), &leaves).unwrap();
416        let root = JubJub::zero();
417        for (i, leaf) in leaves.iter().enumerate() {
418            let proof = tree.generate_proof(i, &leaf).unwrap();
419            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
420        }
421    }
422
423    #[should_panic]
424    #[test]
425    fn bad_root_mht_test() {
426        let mut leaves = Vec::new();
427        for i in 0..4u8 {
428            leaves.push([i, i, i, i, i, i, i, i]);
429        }
430        generate_merkle_tree(&leaves);
431        let mut leaves = Vec::new();
432        for i in 0..100u8 {
433            leaves.push([i, i, i, i, i, i, i, i]);
434        }
435        bad_merkle_tree_verify(&leaves);
436    }
437}