mithril_common/crypto_helper/
merkle_tree.rs

1use anyhow::{Context, anyhow};
2use blake2::{Blake2s256, Digest};
3use ckb_merkle_mountain_range::{
4    Error as MMRError, MMR, MMRStoreReadOps, MMRStoreWriteOps, Merge, MerkleProof,
5    Result as MMRResult,
6};
7use serde::{Deserialize, Serialize};
8use std::{
9    collections::{BTreeMap, HashMap},
10    fmt::Display,
11    ops::{Add, Deref},
12    sync::{Arc, RwLock},
13};
14
15use crate::{StdError, StdResult};
16
17/// Alias for a byte
18pub type Bytes = Vec<u8>;
19
20/// Alias for a Merkle tree leaf position
21pub type MKTreeLeafPosition = u64;
22
23/// A node of a Merkle tree
24#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash, Serialize, Deserialize)]
25pub struct MKTreeNode {
26    hash: Bytes,
27}
28
29impl MKTreeNode {
30    /// MKTreeNode factory
31    pub fn new(hash: Bytes) -> Self {
32        Self { hash }
33    }
34
35    /// Create a MKTreeNode from a hex representation
36    pub fn from_hex(hex: &str) -> StdResult<Self> {
37        let hash = hex::decode(hex)?;
38        Ok(Self { hash })
39    }
40
41    /// Create a hex representation of the MKTreeNode
42    pub fn to_hex(&self) -> String {
43        hex::encode(&self.hash)
44    }
45}
46
47impl Deref for MKTreeNode {
48    type Target = Bytes;
49
50    fn deref(&self) -> &Self::Target {
51        &self.hash
52    }
53}
54
55impl From<String> for MKTreeNode {
56    fn from(other: String) -> Self {
57        Self {
58            hash: other.as_str().into(),
59        }
60    }
61}
62
63impl From<&String> for MKTreeNode {
64    fn from(other: &String) -> Self {
65        Self {
66            hash: other.as_str().into(),
67        }
68    }
69}
70
71impl From<&str> for MKTreeNode {
72    fn from(other: &str) -> Self {
73        Self {
74            hash: other.as_bytes().to_vec(),
75        }
76    }
77}
78
79impl<S: MKTreeStorer> TryFrom<MKTree<S>> for MKTreeNode {
80    type Error = StdError;
81    fn try_from(other: MKTree<S>) -> Result<Self, Self::Error> {
82        other.compute_root()
83    }
84}
85
86impl<S: MKTreeStorer> TryFrom<&MKTree<S>> for MKTreeNode {
87    type Error = StdError;
88    fn try_from(other: &MKTree<S>) -> Result<Self, Self::Error> {
89        other.compute_root()
90    }
91}
92
93impl Display for MKTreeNode {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        write!(f, "{}", String::from_utf8_lossy(&self.hash))
96    }
97}
98
99impl Add for MKTreeNode {
100    type Output = MKTreeNode;
101
102    fn add(self, other: MKTreeNode) -> MKTreeNode {
103        &self + &other
104    }
105}
106
107impl Add for &MKTreeNode {
108    type Output = MKTreeNode;
109
110    fn add(self, other: &MKTreeNode) -> MKTreeNode {
111        let mut hasher = Blake2s256::new();
112        hasher.update(self.deref());
113        hasher.update(other.deref());
114        let hash_merge = hasher.finalize();
115        MKTreeNode::new(hash_merge.to_vec())
116    }
117}
118
119struct MergeMKTreeNode {}
120
121impl Merge for MergeMKTreeNode {
122    type Item = Arc<MKTreeNode>;
123
124    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> MMRResult<Self::Item> {
125        Ok(Arc::new((**lhs).clone() + (**rhs).clone()))
126    }
127}
128
129/// A Merkle proof
130#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
131pub struct MKProof {
132    inner_root: Arc<MKTreeNode>,
133    inner_leaves: Vec<(MKTreeLeafPosition, Arc<MKTreeNode>)>,
134    inner_proof_size: u64,
135    inner_proof_items: Vec<Arc<MKTreeNode>>,
136}
137
138impl MKProof {
139    /// Return a reference to its merkle root.
140    pub fn root(&self) -> &MKTreeNode {
141        &self.inner_root
142    }
143
144    /// Verification of a Merkle proof
145    pub fn verify(&self) -> StdResult<()> {
146        MerkleProof::<Arc<MKTreeNode>, MergeMKTreeNode>::new(
147            self.inner_proof_size,
148            self.inner_proof_items.clone(),
149        )
150        .verify(self.inner_root.to_owned(), self.inner_leaves.to_owned())?
151        .then_some(())
152        .ok_or(anyhow!("Invalid MKProof"))
153    }
154
155    /// Check if the proof contains the given leaves
156    pub fn contains(&self, leaves: &[MKTreeNode]) -> StdResult<()> {
157        leaves
158            .iter()
159            .all(|leaf| self.inner_leaves.iter().any(|(_, l)| l.deref() == leaf))
160            .then_some(())
161            .ok_or(anyhow!("Leaves not found in the MKProof"))
162    }
163
164    /// List the leaves of the proof
165    pub fn leaves(&self) -> Vec<MKTreeNode> {
166        self.inner_leaves
167            .iter()
168            .map(|(_, l)| (**l).clone())
169            .collect::<Vec<_>>()
170    }
171
172    /// Convert the proof to bytes
173    pub fn to_bytes(&self) -> StdResult<Bytes> {
174        bincode::serde::encode_to_vec(self, bincode::config::standard()).map_err(|e| e.into())
175    }
176
177    /// Convert the proof from bytes
178    pub fn from_bytes(bytes: &[u8]) -> StdResult<Self> {
179        let (res, _) =
180            bincode::serde::decode_from_slice::<Self, _>(bytes, bincode::config::standard())?;
181
182        Ok(res)
183    }
184}
185
186impl From<MKProof> for MKTreeNode {
187    fn from(other: MKProof) -> Self {
188        other.root().to_owned()
189    }
190}
191
192/// A Merkle tree store in memory
193#[derive(Clone)]
194pub struct MKTreeStoreInMemory {
195    inner_leaves: Arc<RwLock<HashMap<Arc<MKTreeNode>, MKTreeLeafPosition>>>,
196    inner_store: Arc<RwLock<HashMap<u64, Arc<MKTreeNode>>>>,
197}
198
199impl MKTreeStoreInMemory {
200    fn new() -> Self {
201        Self {
202            inner_leaves: Arc::new(RwLock::new(HashMap::new())),
203            inner_store: Arc::new(RwLock::new(HashMap::new())),
204        }
205    }
206}
207
208impl MKTreeLeafIndexer for MKTreeStoreInMemory {
209    fn set_leaf_position(&self, pos: MKTreeLeafPosition, node: Arc<MKTreeNode>) -> StdResult<()> {
210        let mut inner_leaves = self.inner_leaves.write().unwrap();
211        (*inner_leaves).insert(node, pos);
212
213        Ok(())
214    }
215
216    fn get_leaf_position(&self, node: &MKTreeNode) -> Option<MKTreeLeafPosition> {
217        let inner_leaves = self.inner_leaves.read().unwrap();
218        (*inner_leaves).get(node).cloned()
219    }
220
221    fn total_leaves(&self) -> usize {
222        let inner_leaves = self.inner_leaves.read().unwrap();
223        (*inner_leaves).len()
224    }
225
226    fn leaves(&self) -> Vec<MKTreeNode> {
227        let inner_leaves = self.inner_leaves.read().unwrap();
228        (*inner_leaves)
229            .iter()
230            .map(|(leaf, position)| (position, leaf))
231            .collect::<BTreeMap<_, _>>()
232            .into_values()
233            .map(|leaf| (**leaf).clone())
234            .collect()
235    }
236}
237
238impl MKTreeStorer for MKTreeStoreInMemory {
239    fn build() -> StdResult<Self> {
240        Ok(Self::new())
241    }
242
243    fn get_elem(&self, pos: u64) -> StdResult<Option<Arc<MKTreeNode>>> {
244        let inner_store = self.inner_store.read().unwrap();
245
246        Ok((*inner_store).get(&pos).cloned())
247    }
248
249    fn append(&self, pos: u64, elems: Vec<Arc<MKTreeNode>>) -> StdResult<()> {
250        let mut inner_store = self.inner_store.write().unwrap();
251        for (i, elem) in elems.into_iter().enumerate() {
252            (*inner_store).insert(pos + i as u64, elem);
253        }
254
255        Ok(())
256    }
257}
258
259/// The Merkle tree storer trait
260pub trait MKTreeStorer: Clone + Send + Sync + MKTreeLeafIndexer {
261    /// Try to create a new instance of the storer
262    fn build() -> StdResult<Self>;
263
264    /// Get the element at the given position
265    fn get_elem(&self, pos: u64) -> StdResult<Option<Arc<MKTreeNode>>>;
266
267    /// Append elements at the given position
268    fn append(&self, pos: u64, elems: Vec<Arc<MKTreeNode>>) -> StdResult<()>;
269}
270
271/// This struct exists only to implement for a [MkTreeStore] the [MMRStoreReadOps] and
272/// [MMRStoreWriteOps] from merkle_mountain_range crate without the need to reexport types
273/// from that crate.
274///
275/// Rust don't allow the following:
276/// ```ignore
277/// impl<S: MKTreeStorer> MMRStoreReadOps<Arc<MKTreeNode>> for S {}
278/// ```
279/// Since it disallows implementations of traits for arbitrary types which are not defined in
280/// the same crate as the trait itself (see [E0117](https://doc.rust-lang.org/error_codes/E0117.html)).
281struct MKTreeStore<S: MKTreeStorer> {
282    storer: Box<S>,
283}
284
285impl<S: MKTreeStorer> MKTreeStore<S> {
286    fn build() -> StdResult<Self> {
287        let storer = Box::new(S::build()?);
288        Ok(Self { storer })
289    }
290}
291
292impl<S: MKTreeStorer> MMRStoreReadOps<Arc<MKTreeNode>> for MKTreeStore<S> {
293    fn get_elem(&self, pos: u64) -> MMRResult<Option<Arc<MKTreeNode>>> {
294        self.storer
295            .get_elem(pos)
296            .map_err(|e| MMRError::StoreError(e.to_string()))
297    }
298}
299
300impl<S: MKTreeStorer> MMRStoreWriteOps<Arc<MKTreeNode>> for MKTreeStore<S> {
301    fn append(&mut self, pos: u64, elems: Vec<Arc<MKTreeNode>>) -> MMRResult<()> {
302        self.storer
303            .append(pos, elems)
304            .map_err(|e| MMRError::StoreError(e.to_string()))
305    }
306}
307
308impl<S: MKTreeStorer> MKTreeLeafIndexer for MKTreeStore<S> {
309    fn set_leaf_position(&self, pos: MKTreeLeafPosition, leaf: Arc<MKTreeNode>) -> StdResult<()> {
310        self.storer.set_leaf_position(pos, leaf)
311    }
312
313    fn get_leaf_position(&self, leaf: &MKTreeNode) -> Option<MKTreeLeafPosition> {
314        self.storer.get_leaf_position(leaf)
315    }
316
317    fn total_leaves(&self) -> usize {
318        self.storer.total_leaves()
319    }
320
321    fn leaves(&self) -> Vec<MKTreeNode> {
322        self.storer.leaves()
323    }
324}
325
326/// The Merkle tree leaves indexer trait
327pub trait MKTreeLeafIndexer {
328    /// Get the position of the leaf in the Merkle tree
329    fn set_leaf_position(&self, pos: MKTreeLeafPosition, leaf: Arc<MKTreeNode>) -> StdResult<()>;
330
331    /// Get the position of the leaf in the Merkle tree
332    fn get_leaf_position(&self, leaf: &MKTreeNode) -> Option<MKTreeLeafPosition>;
333
334    /// Number of leaves in the Merkle tree
335    fn total_leaves(&self) -> usize;
336
337    /// List of leaves with their positions in the Merkle tree
338    fn leaves(&self) -> Vec<MKTreeNode>;
339
340    /// Check if the Merkle tree contains the given leaf
341    fn contains_leaf(&self, leaf: &MKTreeNode) -> bool {
342        self.get_leaf_position(leaf).is_some()
343    }
344}
345
346/// A Merkle tree
347pub struct MKTree<S: MKTreeStorer> {
348    inner_tree: MMR<Arc<MKTreeNode>, MergeMKTreeNode, MKTreeStore<S>>,
349}
350
351impl<S: MKTreeStorer> MKTree<S> {
352    /// MKTree factory
353    pub fn new<T: Into<MKTreeNode> + Clone>(leaves: &[T]) -> StdResult<Self> {
354        let mut inner_tree = MMR::<_, _, _>::new(0, MKTreeStore::<S>::build()?);
355        for leaf in leaves {
356            let leaf = Arc::new(leaf.to_owned().into());
357            let inner_tree_position = inner_tree.push(leaf.clone())?;
358            inner_tree
359                .store()
360                .set_leaf_position(inner_tree_position, leaf.clone())?;
361        }
362        inner_tree.commit()?;
363
364        Ok(Self { inner_tree })
365    }
366
367    /// Append leaves to the Merkle tree
368    pub fn append<T: Into<MKTreeNode> + Clone>(&mut self, leaves: &[T]) -> StdResult<()> {
369        for leaf in leaves {
370            let leaf = Arc::new(leaf.to_owned().into());
371            let inner_tree_position = self.inner_tree.push(leaf.clone())?;
372            self.inner_tree
373                .store()
374                .set_leaf_position(inner_tree_position, leaf.clone())?;
375        }
376        self.inner_tree.commit()?;
377
378        Ok(())
379    }
380
381    /// Number of leaves in the Merkle tree
382    pub fn total_leaves(&self) -> usize {
383        self.inner_tree.store().total_leaves()
384    }
385
386    /// List of leaves with their positions in the Merkle tree
387    pub fn leaves(&self) -> Vec<MKTreeNode> {
388        self.inner_tree.store().leaves()
389    }
390
391    /// Check if the Merkle tree contains the given leaf
392    pub fn contains(&self, leaf: &MKTreeNode) -> bool {
393        self.inner_tree.store().contains_leaf(leaf)
394    }
395
396    /// Generate root of the Merkle tree
397    pub fn compute_root(&self) -> StdResult<MKTreeNode> {
398        Ok((*self
399            .inner_tree
400            .get_root()
401            .with_context(|| "Could not compute Merkle Tree root")?)
402        .clone())
403    }
404
405    /// Generate Merkle proof of memberships in the tree
406    pub fn compute_proof(&self, leaves: &[MKTreeNode]) -> StdResult<MKProof> {
407        let inner_leaves = leaves
408            .iter()
409            .map(|leaf| {
410                if let Some(leaf_position) = self.inner_tree.store().get_leaf_position(leaf) {
411                    Ok((leaf_position, Arc::new(leaf.to_owned())))
412                } else {
413                    Err(anyhow!("Leaf not found in the Merkle tree"))
414                }
415            })
416            .collect::<StdResult<Vec<_>>>()?;
417        let proof = self.inner_tree.gen_proof(
418            inner_leaves
419                .iter()
420                .map(|(leaf_position, _leaf)| *leaf_position)
421                .collect(),
422        )?;
423        Ok(MKProof {
424            inner_root: Arc::new(self.compute_root()?),
425            inner_leaves,
426            inner_proof_size: proof.mmr_size(),
427            inner_proof_items: proof.proof_items().to_vec(),
428        })
429    }
430}
431
432impl<S: MKTreeStorer> Clone for MKTree<S> {
433    fn clone(&self) -> Self {
434        // Cloning should never fail so unwrap is safe
435        Self::new(&self.leaves()).unwrap()
436    }
437}
438
439#[cfg(test)]
440mod tests {
441    use crate::test::crypto_helper::MKProofTestExtension;
442
443    use super::*;
444
445    fn generate_leaves(total_leaves: usize) -> Vec<MKTreeNode> {
446        (0..total_leaves).map(|i| format!("test-{i}").into()).collect()
447    }
448
449    #[test]
450    fn test_golden_merkle_root() {
451        let leaves = vec!["golden-1", "golden-2", "golden-3", "golden-4", "golden-5"];
452        let mktree =
453            MKTree::<MKTreeStoreInMemory>::new(&leaves).expect("MKTree creation should not fail");
454        let mkroot = mktree.compute_root().expect("MKRoot generation should not fail");
455
456        assert_eq!(
457            "3bbced153528697ecde7345a22e50115306478353619411523e804f2323fd921",
458            mkroot.to_hex()
459        );
460    }
461
462    #[test]
463    fn test_should_accept_valid_proof_generated_by_merkle_tree() {
464        let leaves = generate_leaves(10);
465        let leaves_to_verify = &[leaves[0].to_owned(), leaves[3].to_owned()];
466        let proof =
467            MKProof::from_leaves(leaves_to_verify).expect("MKProof generation should not fail");
468        proof.verify().expect("The MKProof should be valid");
469    }
470
471    #[test]
472    fn test_should_serialize_deserialize_proof() {
473        let leaves = generate_leaves(10);
474        let leaves_to_verify = &[leaves[0].to_owned(), leaves[3].to_owned()];
475        let proof =
476            MKProof::from_leaves(leaves_to_verify).expect("MKProof generation should not fail");
477
478        let serialized_proof = proof.to_bytes().expect("Serialization should not fail");
479        let deserialized_proof =
480            MKProof::from_bytes(&serialized_proof).expect("Deserialization should not fail");
481        assert_eq!(
482            proof, deserialized_proof,
483            "Deserialized proof should match the original"
484        );
485    }
486
487    #[test]
488    fn test_should_reject_invalid_proof_generated_by_merkle_tree() {
489        let leaves = generate_leaves(10);
490        let leaves_to_verify = &[leaves[0].to_owned(), leaves[3].to_owned()];
491        let mut proof =
492            MKProof::from_leaves(leaves_to_verify).expect("MKProof generation should not fail");
493        proof.inner_root = Arc::new(leaves[1].to_owned());
494        proof.verify().expect_err("The MKProof should be invalid");
495    }
496
497    #[test]
498    fn test_should_list_leaves() {
499        let leaves: Vec<MKTreeNode> = vec!["test-0".into(), "test-1".into(), "test-2".into()];
500        let mktree =
501            MKTree::<MKTreeStoreInMemory>::new(&leaves).expect("MKTree creation should not fail");
502        let leaves_retrieved = mktree.leaves();
503
504        assert_eq!(
505            leaves.iter().collect::<Vec<_>>(),
506            leaves_retrieved.iter().collect::<Vec<_>>()
507        );
508    }
509
510    #[test]
511    fn test_should_clone_and_compute_same_root() {
512        let leaves = generate_leaves(10);
513        let mktree =
514            MKTree::<MKTreeStoreInMemory>::new(&leaves).expect("MKTree creation should not fail");
515        let mktree_clone = mktree.clone();
516
517        assert_eq!(
518            mktree.compute_root().unwrap(),
519            mktree_clone.compute_root().unwrap(),
520        );
521    }
522
523    #[test]
524    fn test_should_support_append_leaves() {
525        let leaves = generate_leaves(10);
526        let leaves_creation = &leaves[..9];
527        let leaves_to_append = &leaves[9..];
528        let mut mktree = MKTree::<MKTreeStoreInMemory>::new(leaves_creation)
529            .expect("MKTree creation should not fail");
530        mktree
531            .append(leaves_to_append)
532            .expect("MKTree append leaves should not fail");
533
534        assert_eq!(10, mktree.total_leaves());
535    }
536
537    #[test]
538    fn tree_node_from_to_string() {
539        let expected_str = "my_string";
540        let expected_string = expected_str.to_string();
541        let node_str: MKTreeNode = expected_str.into();
542        let node_string: MKTreeNode = expected_string.clone().into();
543
544        assert_eq!(node_str.to_string(), expected_str);
545        assert_eq!(node_string.to_string(), expected_string);
546    }
547
548    #[test]
549    fn contains_leaves() {
550        let mut leaves_to_verify = generate_leaves(10);
551        let leaves_not_verified = leaves_to_verify.drain(3..6).collect::<Vec<_>>();
552        let proof =
553            MKProof::from_leaves(&leaves_to_verify).expect("MKProof generation should not fail");
554
555        // contains everything
556        proof.contains(&leaves_to_verify).unwrap();
557
558        // contains subpart
559        proof.contains(&leaves_to_verify[0..2]).unwrap();
560
561        // don't contains all not verified
562        proof.contains(&leaves_not_verified).unwrap_err();
563
564        // don't contains subpart of not verified
565        proof.contains(&leaves_not_verified[1..2]).unwrap_err();
566
567        // fail if part verified and part unverified
568        proof
569            .contains(&[leaves_to_verify[2].to_owned(), leaves_not_verified[0].to_owned()])
570            .unwrap_err();
571    }
572
573    #[test]
574    fn list_leaves() {
575        let leaves_to_verify = generate_leaves(10);
576        let proof =
577            MKProof::from_leaves(&leaves_to_verify).expect("MKProof generation should not fail");
578
579        let proof_leaves = proof.leaves();
580        assert_eq!(proof_leaves, leaves_to_verify);
581    }
582}
mithril_common/crypto_helper/merkle_tree.rs

mithril_common/crypto_helper/
merkle_tree.rs
