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        .with_context(|| "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            .with_context(|| "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        Self::new_from_iter(leaves.iter().cloned())
355    }
356
357    /// MKTree factory
358    pub fn new_from_iter<T: IntoIterator<Item = U>, U: Into<MKTreeNode>>(
359        leaves: T,
360    ) -> StdResult<Self> {
361        let mut inner_tree = MMR::<_, _, _>::new(0, MKTreeStore::<S>::build()?);
362        for leaf in leaves.into_iter() {
363            let leaf = Arc::new(leaf.into());
364            let inner_tree_position = inner_tree.push(leaf.clone())?;
365            inner_tree
366                .store()
367                .set_leaf_position(inner_tree_position, leaf.clone())?;
368        }
369        inner_tree.commit()?;
370
371        Ok(Self { inner_tree })
372    }
373
374    /// Append leaves to the Merkle tree
375    pub fn append<T: Into<MKTreeNode> + Clone>(&mut self, leaves: &[T]) -> StdResult<()> {
376        for leaf in leaves {
377            let leaf = Arc::new(leaf.to_owned().into());
378            let inner_tree_position = self.inner_tree.push(leaf.clone())?;
379            self.inner_tree
380                .store()
381                .set_leaf_position(inner_tree_position, leaf.clone())?;
382        }
383        self.inner_tree.commit()?;
384
385        Ok(())
386    }
387
388    /// Number of leaves in the Merkle tree
389    pub fn total_leaves(&self) -> usize {
390        self.inner_tree.store().total_leaves()
391    }
392
393    /// List of leaves with their positions in the Merkle tree
394    pub fn leaves(&self) -> Vec<MKTreeNode> {
395        self.inner_tree.store().leaves()
396    }
397
398    /// Check if the Merkle tree contains the given leaf
399    pub fn contains(&self, leaf: &MKTreeNode) -> bool {
400        self.inner_tree.store().contains_leaf(leaf)
401    }
402
403    /// Generate root of the Merkle tree
404    pub fn compute_root(&self) -> StdResult<MKTreeNode> {
405        Ok((*self
406            .inner_tree
407            .get_root()
408            .with_context(|| "Could not compute Merkle Tree root")?)
409        .clone())
410    }
411
412    /// Generate Merkle proof of memberships in the tree
413    pub fn compute_proof(&self, leaves: &[MKTreeNode]) -> StdResult<MKProof> {
414        let inner_leaves = leaves
415            .iter()
416            .map(|leaf| {
417                if let Some(leaf_position) = self.inner_tree.store().get_leaf_position(leaf) {
418                    Ok((leaf_position, Arc::new(leaf.to_owned())))
419                } else {
420                    Err(anyhow!("Leaf not found in the Merkle tree"))
421                }
422            })
423            .collect::<StdResult<Vec<_>>>()?;
424        let proof = self.inner_tree.gen_proof(
425            inner_leaves
426                .iter()
427                .map(|(leaf_position, _leaf)| *leaf_position)
428                .collect(),
429        )?;
430        Ok(MKProof {
431            inner_root: Arc::new(self.compute_root()?),
432            inner_leaves,
433            inner_proof_size: proof.mmr_size(),
434            inner_proof_items: proof.proof_items().to_vec(),
435        })
436    }
437}
438
439impl<S: MKTreeStorer> Clone for MKTree<S> {
440    fn clone(&self) -> Self {
441        // Cloning should never fail so unwrap is safe
442        Self::new(&self.leaves()).unwrap()
443    }
444}
445
446#[cfg(test)]
447mod tests {
448    use crate::test::crypto_helper::MKProofTestExtension;
449
450    use super::*;
451
452    fn generate_leaves(total_leaves: usize) -> Vec<MKTreeNode> {
453        (0..total_leaves).map(|i| format!("test-{i}").into()).collect()
454    }
455
456    #[test]
457    fn test_golden_merkle_root() {
458        let leaves = vec!["golden-1", "golden-2", "golden-3", "golden-4", "golden-5"];
459        let mktree =
460            MKTree::<MKTreeStoreInMemory>::new(&leaves).expect("MKTree creation should not fail");
461        let mkroot = mktree.compute_root().expect("MKRoot generation should not fail");
462
463        assert_eq!(
464            "3bbced153528697ecde7345a22e50115306478353619411523e804f2323fd921",
465            mkroot.to_hex()
466        );
467    }
468
469    #[test]
470    fn test_should_accept_valid_proof_generated_by_merkle_tree() {
471        let leaves = generate_leaves(10);
472        let leaves_to_verify = &[leaves[0].to_owned(), leaves[3].to_owned()];
473        let proof =
474            MKProof::from_leaves(leaves_to_verify).expect("MKProof generation should not fail");
475        proof.verify().expect("The MKProof should be valid");
476    }
477
478    #[test]
479    fn test_should_serialize_deserialize_proof() {
480        let leaves = generate_leaves(10);
481        let leaves_to_verify = &[leaves[0].to_owned(), leaves[3].to_owned()];
482        let proof =
483            MKProof::from_leaves(leaves_to_verify).expect("MKProof generation should not fail");
484
485        let serialized_proof = proof.to_bytes().expect("Serialization should not fail");
486        let deserialized_proof =
487            MKProof::from_bytes(&serialized_proof).expect("Deserialization should not fail");
488        assert_eq!(
489            proof, deserialized_proof,
490            "Deserialized proof should match the original"
491        );
492    }
493
494    #[test]
495    fn test_should_reject_invalid_proof_generated_by_merkle_tree() {
496        let leaves = generate_leaves(10);
497        let leaves_to_verify = &[leaves[0].to_owned(), leaves[3].to_owned()];
498        let mut proof =
499            MKProof::from_leaves(leaves_to_verify).expect("MKProof generation should not fail");
500        proof.inner_root = Arc::new(leaves[1].to_owned());
501        proof.verify().expect_err("The MKProof should be invalid");
502    }
503
504    #[test]
505    fn test_should_list_leaves() {
506        let leaves: Vec<MKTreeNode> = vec!["test-0".into(), "test-1".into(), "test-2".into()];
507        let mktree =
508            MKTree::<MKTreeStoreInMemory>::new(&leaves).expect("MKTree creation should not fail");
509        let leaves_retrieved = mktree.leaves();
510
511        assert_eq!(
512            leaves.iter().collect::<Vec<_>>(),
513            leaves_retrieved.iter().collect::<Vec<_>>()
514        );
515    }
516
517    #[test]
518    fn test_should_clone_and_compute_same_root() {
519        let leaves = generate_leaves(10);
520        let mktree =
521            MKTree::<MKTreeStoreInMemory>::new(&leaves).expect("MKTree creation should not fail");
522        let mktree_clone = mktree.clone();
523
524        assert_eq!(
525            mktree.compute_root().unwrap(),
526            mktree_clone.compute_root().unwrap(),
527        );
528    }
529
530    #[test]
531    fn test_should_support_append_leaves() {
532        let leaves = generate_leaves(10);
533        let leaves_creation = &leaves[..9];
534        let leaves_to_append = &leaves[9..];
535        let mut mktree = MKTree::<MKTreeStoreInMemory>::new(leaves_creation)
536            .expect("MKTree creation should not fail");
537        mktree
538            .append(leaves_to_append)
539            .expect("MKTree append leaves should not fail");
540
541        assert_eq!(10, mktree.total_leaves());
542    }
543
544    #[test]
545    fn tree_node_from_to_string() {
546        let expected_str = "my_string";
547        let expected_string = expected_str.to_string();
548        let node_str: MKTreeNode = expected_str.into();
549        let node_string: MKTreeNode = expected_string.clone().into();
550
551        assert_eq!(node_str.to_string(), expected_str);
552        assert_eq!(node_string.to_string(), expected_string);
553    }
554
555    #[test]
556    fn contains_leaves() {
557        let mut leaves_to_verify = generate_leaves(10);
558        let leaves_not_verified = leaves_to_verify.drain(3..6).collect::<Vec<_>>();
559        let proof =
560            MKProof::from_leaves(&leaves_to_verify).expect("MKProof generation should not fail");
561
562        // contains everything
563        proof.contains(&leaves_to_verify).unwrap();
564
565        // contains subpart
566        proof.contains(&leaves_to_verify[0..2]).unwrap();
567
568        // don't contains all not verified
569        proof.contains(&leaves_not_verified).unwrap_err();
570
571        // don't contains subpart of not verified
572        proof.contains(&leaves_not_verified[1..2]).unwrap_err();
573
574        // fail if part verified and part unverified
575        proof
576            .contains(&[leaves_to_verify[2].to_owned(), leaves_not_verified[0].to_owned()])
577            .unwrap_err();
578    }
579
580    #[test]
581    fn list_leaves() {
582        let leaves_to_verify = generate_leaves(10);
583        let proof =
584            MKProof::from_leaves(&leaves_to_verify).expect("MKProof generation should not fail");
585
586        let proof_leaves = proof.leaves();
587        assert_eq!(proof_leaves, leaves_to_verify);
588    }
589}
mithril_common/crypto_helper/merkle_tree.rs

mithril_common/crypto_helper/
merkle_tree.rs
