Merge branch 'rkhalil/tidy-host' of github.com:risc0/zeth into rkhali…

…l/derive-tx-exec

host/src/operations/rollups.rs

@@ -30,7 +30,9 @@ use zeth_lib::{
     },
 };
 use zeth_primitives::{
-    block::Header, transactions::optimism::OptimismTxEssence, tree::MerkleMountainRange,
+    block::Header,
+    transactions::optimism::OptimismTxEssence,
+    tree::{MerkleMountainRange, MerkleProof},
 };
 
 use crate::{
@@ -402,7 +404,7 @@ pub async fn compose_derived_rollup_blocks(
             compose_image_id: OP_COMPOSE_ID,
             operation: ComposeInputOperation::LIFT {
                 derivation: derive_output,
-                eth_tail_proof: MerkleMountainRange::proof(&sibling_map, eth_tail_hash),
+                eth_tail_proof: MerkleProof::new(&sibling_map, eth_tail_hash),
             },
             eth_chain_merkle_root: eth_chain_root,
         };

lib/src/optimism/composition.rs

@@ -16,7 +16,11 @@ use anyhow::bail;
 #[cfg(target_os = "zkvm")]
 use risc0_zkvm::{guest::env, serde::to_vec, sha::Digest};
 use serde::{Deserialize, Serialize};
-use zeth_primitives::{block::Header, tree::MerkleMountainRange, BlockHash, BlockNumber};
+use zeth_primitives::{
+    block::Header,
+    tree::{MerkleMountainRange, MerkleProof},
+    BlockHash, BlockNumber,
+};
 
 use crate::optimism::DeriveOutput;
 
@@ -39,7 +43,7 @@ pub enum ComposeInputOperation {
     },
     LIFT {
         derivation: DeriveOutput,
-        eth_tail_proof: MerkleMountainRange,
+        eth_tail_proof: MerkleProof,
     },
     JOIN {
         left: ComposeOutput,
@@ -147,15 +151,10 @@ impl ComposeInput {
                         .expect("Failed to lift derivation receipt");
                 }
                 // Verify inclusion of ethereum tail in Merkle root
-                assert_eq!(
-                    self.eth_chain_merkle_root,
+                assert!(
                     eth_tail_proof
-                        .root(None)
-                        .expect("No proof included for ethereum tail")
-                );
-                assert_eq!(
-                    eth_tail_proof.roots.first().unwrap().unwrap(),
-                    derive_output.eth_tail.1 .0
+                        .verify(&self.eth_chain_merkle_root, &derive_output.eth_tail.1 .0),
+                    "Invalid ethereum tail inclusion proof!"
                 );
                 // Create output
                 ComposeOutput {

primitives/src/tree.rs

@@ -17,21 +17,33 @@ use std::collections::HashMap;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 
-#[derive(Debug, Clone, Deserialize, Serialize, Default)]
-pub struct MerkleMountainRange {
-    pub roots: Vec<Option<[u8; 32]>>,
+pub type Hash = [u8; 32];
+pub type SiblingMap = HashMap<Hash, Hash>;
+
+/// Returns the hash of two sibling nodes by appending them together while placing
+/// the lexicographically smaller node first.
+#[inline]
+fn branch_hash(a: &Hash, b: &Hash) -> Hash {
+    let mut hasher = Sha256::new();
+    if a < b {
+        hasher.update(a);
+        hasher.update(b);
+    } else {
+        hasher.update(b);
+        hasher.update(a);
+    };
+    hasher.finalize().into()
 }
 
+#[derive(Debug, Clone, Deserialize, Serialize, Default)]
+pub struct MerkleMountainRange(pub Vec<Option<Hash>>);
+
 impl MerkleMountainRange {
     /// Appends a new leaf to the mountain range, bubbling up all the changes required to
     /// the ancestor tree roots. The optional `_sibling_map` parameter can be
     /// specified in order to record all visited siblings.
-    pub fn append_leaf(
-        &mut self,
-        mut value: [u8; 32],
-        mut _sibling_map: Option<&mut HashMap<[u8; 32], [u8; 32]>>,
-    ) {
-        for node in self.roots.iter_mut() {
+    pub fn append_leaf(&mut self, mut value: Hash, mut _sibling_map: Option<&mut SiblingMap>) {
+        for node in self.0.iter_mut() {
             if node.is_none() {
                 node.replace(value);
                 return;
@@ -43,58 +55,78 @@ impl MerkleMountainRange {
                     sibling_map.insert(value, sibling);
                     sibling_map.insert(sibling, value);
                 }
-                value = Self::branch_hash(&value, &sibling)
+                value = branch_hash(&value, &sibling)
             }
         }
-        self.roots.push(Some(value));
+        self.0.push(Some(value));
     }
 
     /// Returns the root of the (unbalanced) Merkle tree that covers all the present range
     /// roots. The optional `_sibling_map` parameter can be specified in order to
     /// record all visited siblings.
-    pub fn root(
-        &self,
-        mut _sibling_map: Option<&mut HashMap<[u8; 32], [u8; 32]>>,
-    ) -> Option<[u8; 32]> {
-        let mut result: Option<[u8; 32]> = None;
-        for root in self.roots.iter().flatten() {
+    pub fn root(&self, mut _sibling_map: Option<&mut SiblingMap>) -> Option<Hash> {
+        let mut result: Option<Hash> = None;
+        for root in self.0.iter().flatten() {
             if let Some(sibling) = result {
                 // We only need to log siblings outside the zkVM
                 #[cfg(not(target_os = "zkvm"))]
                 if let Some(sibling_map) = _sibling_map.as_mut() {
                     sibling_map.insert(*root, sibling);
                     sibling_map.insert(sibling, *root);
                 }
-                result.replace(Self::branch_hash(&sibling, root));
+                result.replace(branch_hash(&sibling, root));
             } else {
                 result.replace(*root);
             }
         }
         result
     }
+}
 
-    /// Returns the hash of two sibling nodes by appending them together while placing
-    /// the lexicographically smaller node first.
-    #[inline]
-    fn branch_hash(a: &[u8; 32], b: &[u8; 32]) -> [u8; 32] {
-        let mut hasher = Sha256::new();
-        if a < b {
-            hasher.update(a);
-            hasher.update(b);
-        } else {
-            hasher.update(b);
-            hasher.update(a);
-        };
-        hasher.finalize().into()
-    }
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct MerkleProof(Vec<Hash>);
 
+impl MerkleProof {
     /// Returns the inclusion proof of the input `value` using the provided `sibling_map`.
-    pub fn proof(sibling_map: &HashMap<[u8; 32], [u8; 32]>, mut value: [u8; 32]) -> Self {
-        let mut roots = vec![Some(value)];
+    pub fn new(sibling_map: &SiblingMap, mut value: Hash) -> Self {
+        let mut roots = vec![value];
         while let Some(sibling) = sibling_map.get(&value) {
-            roots.push(Some(*sibling));
-            value = Self::branch_hash(sibling, &value);
+            roots.push(*sibling);
+            value = branch_hash(sibling, &value);
+        }
+        Self(roots)
+    }
+
+    /// Verifies the inclusion proof against the provided `root` of a
+    /// [MerkleMountainRange] and a `value`.
+    pub fn verify(&self, root: &Hash, value: &Hash) -> bool {
+        let mut iter = self.0.iter();
+        match iter.next() {
+            None => false,
+            Some(first) => {
+                first == value && iter.fold(*first, |result, e| branch_hash(&result, e)) == *root
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_proof_verify() {
+        let values = vec![[0u8; 32], [1; 32], [2; 32], [3; 32], [4; 32]];
+        let value = values[0];
+        let mut mmr = MerkleMountainRange::default();
+        let mut sibling_map = SiblingMap::new();
+        for value in values {
+            mmr.append_leaf(value, Some(&mut sibling_map));
         }
-        Self { roots }
+        let root = mmr.root(Some(&mut sibling_map)).unwrap();
+        let proof = MerkleProof::new(&sibling_map, value);
+        assert!(proof.verify(&root, &value));
+        // test that the proof is not valid for a different value/root
+        assert!(!proof.verify(&root, &[0xff; 32]));
+        assert!(!proof.verify(&[0x00; 32], &value));
     }
 }