hypercall_state_commitment/
mmr.rs

1use ckb_merkle_mountain_range::{
2    MMRStoreReadOps, MMRStoreWriteOps, Merge, Result as MmrResult, MMR,
3};
4use sha3::{Digest, Keccak256};
5use std::cell::RefCell;
6use std::collections::HashMap;
7use std::sync::RwLock;
8
9#[derive(Clone, Debug, PartialEq, Eq)]
10pub struct MmrHash(pub [u8; 32]);
11
12pub struct Keccak256Merge;
13
14impl Merge for Keccak256Merge {
15    type Item = MmrHash;
16
17    fn merge(left: &MmrHash, right: &MmrHash) -> MmrResult<MmrHash> {
18        let mut h = Keccak256::new();
19        h.update([0x01]); // domain tag: internal node
20        h.update(left.0);
21        h.update(right.0);
22        Ok(MmrHash(h.finalize().into()))
23    }
24}
25
26/// In-memory MMR store backed by a HashMap.
27pub struct MemoryMmrStore {
28    pub data: RwLock<HashMap<u64, MmrHash>>,
29}
30
31impl Default for MemoryMmrStore {
32    fn default() -> Self {
33        Self {
34            data: RwLock::new(HashMap::new()),
35        }
36    }
37}
38
39impl MMRStoreReadOps<MmrHash> for &MemoryMmrStore {
40    fn get_elem(&self, pos: u64) -> MmrResult<Option<MmrHash>> {
41        Ok(self.data.read().unwrap().get(&pos).cloned())
42    }
43}
44
45impl MMRStoreWriteOps<MmrHash> for &MemoryMmrStore {
46    fn append(&mut self, pos: u64, elems: Vec<MmrHash>) -> MmrResult<()> {
47        let mut data = self.data.write().unwrap();
48        for (i, elem) in elems.into_iter().enumerate() {
49            let p = pos + i as u64;
50            debug_assert!(
51                !data.contains_key(&p),
52                "MMR append overwrites existing position {p}"
53            );
54            data.insert(p, elem);
55        }
56        Ok(())
57    }
58}
59
60/// Metadata needed to resume an MMR at the exact append position.
61pub trait MmrMetadataStore {
62    fn load_mmr_size(&self) -> anyhow::Result<u64>;
63}
64
65pub trait SupportedMmrStore: MmrMetadataStore {}
66
67pub trait PrepareMmrStore: SupportedMmrStore {
68    fn prepare_append_many_from(
69        size: u64,
70        store: &Self,
71        data: &[[u8; 32]],
72    ) -> anyhow::Result<PreparedMmrAppend>;
73}
74
75impl MmrMetadataStore for MemoryMmrStore {
76    fn load_mmr_size(&self) -> anyhow::Result<u64> {
77        Ok(0)
78    }
79}
80
81impl SupportedMmrStore for MemoryMmrStore {}
82
83impl PrepareMmrStore for MemoryMmrStore {
84    fn prepare_append_many_from(
85        size: u64,
86        store: &Self,
87        data: &[[u8; 32]],
88    ) -> anyhow::Result<PreparedMmrAppend> {
89        prepare_append_many_with_store(size, store, data)
90    }
91}
92
93pub(crate) fn prepare_append_many_with_store<S>(
94    size: u64,
95    store: &S,
96    data: &[[u8; 32]],
97) -> anyhow::Result<PreparedMmrAppend>
98where
99    for<'a> &'a S: MMRStoreReadOps<MmrHash>,
100    for<'a> &'a StagedMmrStore<'a, S>: MMRStoreReadOps<MmrHash> + MMRStoreWriteOps<MmrHash>,
101{
102    if data.is_empty() {
103        let root = if size == 0 {
104            MmrHash([0u8; 32])
105        } else {
106            let mmr = MMR::<MmrHash, Keccak256Merge, _>::new(size, store);
107            mmr.get_root().expect("MMR root should exist")
108        };
109        return Ok(PreparedMmrAppend {
110            start_pos: size,
111            elems: Vec::new(),
112            next_size: size,
113            root,
114        });
115    }
116
117    let staged = StagedMmrStore::new(store);
118    let mut mmr = MMR::<MmrHash, Keccak256Merge, _>::new(size, &staged);
119    for item in data {
120        mmr.push(hash_mmr_leaf(item))
121            .expect("MMR push should not fail");
122    }
123    mmr.commit().expect("MMR commit should not fail");
124    let next_size = mmr.mmr_size();
125    let root = mmr.get_root().expect("MMR root should exist");
126    let writes = staged.flattened_writes()?;
127    let start_pos = writes.first().map(|(pos, _)| *pos).unwrap_or(size);
128    let elems = writes.into_iter().map(|(_, elem)| elem).collect();
129
130    Ok(PreparedMmrAppend {
131        start_pos,
132        elems,
133        next_size,
134        root,
135    })
136}
137
138/// Self-contained MMR that owns its store.
139///
140/// Wraps the ckb MMR with its backing store so callers don't have to
141/// manage store lifetimes. Exposes append + peaks_hash. Runtime code can
142/// provide a durable store; tests default to the in-memory store.
143pub struct HypercallMmr<S = MemoryMmrStore> {
144    store: S,
145    size: u64,
146}
147
148#[derive(Clone, Debug)]
149pub struct PreparedMmrAppend {
150    start_pos: u64,
151    elems: Vec<MmrHash>,
152    next_size: u64,
153    root: MmrHash,
154}
155
156impl PreparedMmrAppend {
157    pub fn root(&self) -> MmrHash {
158        self.root.clone()
159    }
160
161    pub fn is_empty(&self) -> bool {
162        self.elems.is_empty()
163    }
164}
165
166pub(crate) struct StagedMmrStore<'a, S> {
167    pub(crate) base: &'a S,
168    pub(crate) writes: RefCell<Vec<(u64, Vec<MmrHash>)>>,
169}
170
171impl<'a, S> StagedMmrStore<'a, S> {
172    fn new(base: &'a S) -> Self {
173        Self {
174            base,
175            writes: RefCell::new(Vec::new()),
176        }
177    }
178
179    fn flattened_writes(&self) -> MmrResult<Vec<(u64, MmrHash)>> {
180        let mut out = Vec::new();
181        for (start, elems) in self.writes.borrow().iter() {
182            for (index, elem) in elems.iter().enumerate() {
183                out.push((start + index as u64, elem.clone()));
184            }
185        }
186        Ok(out)
187    }
188}
189
190macro_rules! impl_staged_mmr_read {
191    ($store:ty) => {
192        impl MMRStoreReadOps<MmrHash> for &StagedMmrStore<'_, $store> {
193            fn get_elem(&self, pos: u64) -> MmrResult<Option<MmrHash>> {
194                for (start, elems) in self.writes.borrow().iter().rev() {
195                    if pos >= *start {
196                        let index = (pos - *start) as usize;
197                        if let Some(elem) = elems.get(index) {
198                            return Ok(Some(elem.clone()));
199                        }
200                    }
201                }
202                self.base.get_elem(pos)
203            }
204        }
205    };
206}
207
208impl_staged_mmr_read!(MemoryMmrStore);
209
210impl<S> MMRStoreWriteOps<MmrHash> for &StagedMmrStore<'_, S> {
211    fn append(&mut self, pos: u64, elems: Vec<MmrHash>) -> MmrResult<()> {
212        self.writes.borrow_mut().push((pos, elems));
213        Ok(())
214    }
215}
216
217impl HypercallMmr<MemoryMmrStore> {
218    pub fn new() -> Self {
219        Self {
220            store: MemoryMmrStore::default(),
221            size: 0,
222        }
223    }
224}
225
226impl<S> HypercallMmr<S>
227where
228    S: PrepareMmrStore,
229    for<'a> &'a S: MMRStoreReadOps<MmrHash> + MMRStoreWriteOps<MmrHash>,
230{
231    pub fn from_store(store: S) -> anyhow::Result<Self> {
232        let size = store.load_mmr_size()?;
233        Ok(Self { store, size })
234    }
235
236    /// Append a raw 32-byte hash as a leaf.
237    pub fn append(&mut self, data: [u8; 32]) -> anyhow::Result<()> {
238        let prepared = self.prepare_append_many(std::slice::from_ref(&data))?;
239        self.apply_prepared_append(prepared)?;
240        Ok(())
241    }
242
243    pub fn prepare_append_many(&self, data: &[[u8; 32]]) -> anyhow::Result<PreparedMmrAppend> {
244        S::prepare_append_many_from(self.size, &self.store, data)
245    }
246
247    pub fn apply_prepared_append(&mut self, prepared: PreparedMmrAppend) -> anyhow::Result<()> {
248        if prepared.is_empty() {
249            self.size = prepared.next_size;
250            return Ok(());
251        }
252
253        let mut store = &self.store;
254        store
255            .append(prepared.start_pos, prepared.elems)
256            .map_err(|error| anyhow::anyhow!("failed to apply prepared MMR append: {error:?}"))?;
257        self.size = prepared.next_size;
258        Ok(())
259    }
260
261    /// Get the peaks hash (root commitment).
262    ///
263    /// Returns a zero hash if the MMR is empty.
264    pub fn peaks_hash(&self) -> MmrHash {
265        if self.size == 0 {
266            return MmrHash([0u8; 32]);
267        }
268        let mmr = MMR::<MmrHash, Keccak256Merge, _>::new(self.size, &self.store);
269        mmr.get_root().expect("MMR root should exist")
270    }
271
272    pub fn size(&self) -> u64 {
273        self.size
274    }
275}
276
277impl Default for HypercallMmr {
278    fn default() -> Self {
279        Self::new()
280    }
281}
282
283/// Hash arbitrary data into an MMR leaf with domain separation.
284/// Domain tag 0x00 for leaves, 0x01 for internal nodes (in Merge).
285/// This prevents second-preimage attacks where a leaf hash could be
286/// mistaken for an internal node hash.
287pub fn hash_mmr_leaf(data: &[u8]) -> MmrHash {
288    let mut h = Keccak256::new();
289    h.update([0x00]); // domain tag: leaf
290    h.update(data);
291    MmrHash(h.finalize().into())
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[test]
299    fn append_and_get_root() {
300        let store = MemoryMmrStore::default();
301        let mut mmr = MMR::<MmrHash, Keccak256Merge, _>::new(0, &store);
302
303        let leaf = hash_mmr_leaf(b"command-1");
304        mmr.push(leaf.clone()).expect("push should succeed");
305        mmr.commit().expect("commit should succeed");
306
307        let root = mmr.get_root().expect("root should exist");
308        assert_eq!(root, leaf, "single-element MMR root equals the leaf");
309    }
310
311    #[test]
312    fn root_changes_on_append() {
313        let store = MemoryMmrStore::default();
314        let mut mmr = MMR::<MmrHash, Keccak256Merge, _>::new(0, &store);
315
316        mmr.push(hash_mmr_leaf(b"cmd-1")).unwrap();
317        mmr.commit().unwrap();
318        let root1 = mmr.get_root().unwrap();
319
320        mmr.push(hash_mmr_leaf(b"cmd-2")).unwrap();
321        mmr.commit().unwrap();
322        let root2 = mmr.get_root().unwrap();
323
324        assert_ne!(root1, root2);
325    }
326
327    #[test]
328    fn inclusion_proof_verifies() {
329        let store = MemoryMmrStore::default();
330        let mut mmr = MMR::<MmrHash, Keccak256Merge, _>::new(0, &store);
331
332        let positions: Vec<u64> = (0..5)
333            .map(|i| {
334                let pos = mmr
335                    .push(hash_mmr_leaf(format!("cmd-{}", i).as_bytes()))
336                    .unwrap();
337                pos
338            })
339            .collect();
340        mmr.commit().unwrap();
341
342        let root = mmr.get_root().unwrap();
343        let _mmr_size = mmr.mmr_size();
344
345        for &pos in &positions {
346            let proof = mmr.gen_proof(vec![pos]).unwrap();
347            let leaf = store.data.read().unwrap().get(&pos).cloned().unwrap();
348            assert!(
349                proof.verify(root.clone(), vec![(pos, leaf)]).unwrap(),
350                "proof for position {} should verify",
351                pos
352            );
353        }
354    }
355
356    #[test]
357    fn deterministic_roots() {
358        let store_a = MemoryMmrStore::default();
359        let store_b = MemoryMmrStore::default();
360        let mut mmr_a = MMR::<MmrHash, Keccak256Merge, _>::new(0, &store_a);
361        let mut mmr_b = MMR::<MmrHash, Keccak256Merge, _>::new(0, &store_b);
362
363        for i in 0..10 {
364            let leaf = hash_mmr_leaf(format!("command-{}", i).as_bytes());
365            mmr_a.push(leaf.clone()).unwrap();
366            mmr_b.push(leaf).unwrap();
367        }
368        mmr_a.commit().unwrap();
369        mmr_b.commit().unwrap();
370
371        assert_eq!(
372            mmr_a.get_root().unwrap(),
373            mmr_b.get_root().unwrap(),
374            "same inputs should produce same root"
375        );
376    }
377
378    #[test]
379    fn many_appends() {
380        let store = MemoryMmrStore::default();
381        let mut mmr = MMR::<MmrHash, Keccak256Merge, _>::new(0, &store);
382
383        for i in 0..1000 {
384            mmr.push(hash_mmr_leaf(format!("{}", i).as_bytes()))
385                .unwrap();
386        }
387        mmr.commit().unwrap();
388
389        let root = mmr.get_root().unwrap();
390        assert_ne!(root.0, [0u8; 32]);
391    }
392}
hypercall_state_commitment/mmr.rs

hypercall_state_commitment/
mmr.rs
