[feature/jmt (3/4)] Proof generation and verification

fuel-merkle/src/jellyfish.rs

@@ -8,6 +8,9 @@ mod error;
 
 pub use error::MerkleTreeError;
 
+/// Inclusion and exclusion proofs for the Jellyfish Merkle Tree.
+pub mod proof;
+
 // Re-export dependencies from the jmt crate necessary for defining implementations
 // of the Mappable trait required by the JellyfishMerkleTree integration.
 pub use jmt;

fuel-merkle/src/jellyfish/merkle_tree.rs

@@ -26,6 +26,12 @@ use jmt::{
     Sha256Jmt,
 };
 
+use crate::jellyfish::proof::{
+    ExclusionProof,
+    InclusionProof,
+    MerkleProof,
+};
+
 use crate::jellyfish::error::MerkleTreeError;
 
 // Obtained by creating an empty tree.
@@ -206,6 +212,33 @@ where
             Err(MerkleTreeError::RootHashMismatch(*root, root_from_storage))
         }
     }
+
+    pub fn generate_proof(
+        &self,
+        key: &MerkleTreeKey,
+    ) -> Result<MerkleProof, MerkleTreeError<StorageError>> {
+        let jmt = self.as_jmt();
+        let key_hash = jmt::KeyHash(**key);
+        let version = self
+            .get_latest_root_version()
+            .unwrap_or_default()
+            .unwrap_or_default();
+        let (value_vec, proof) = jmt
+            .get_with_proof(key_hash, version)
+            .map_err(MerkleTreeError::JmtError)?;
+        let proof = match value_vec {
+            Some(value) => MerkleProof::Inclusion(InclusionProof {
+                proof,
+                key: key_hash,
+                value,
+            }),
+            None => MerkleProof::Exclusion(ExclusionProof {
+                proof,
+                key: key_hash,
+            }),
+        };
+        Ok(proof)
+    }
 }
 
 impl<

fuel-merkle/src/jellyfish/proof.rs

@@ -0,0 +1,56 @@
+use jmt::proof::SparseMerkleProof;
+
+use crate::common::Bytes32;
+
+// Give crate access to fields for testing tampering with proofs
+pub struct InclusionProof {
+    pub(crate) proof: SparseMerkleProof<sha2::Sha256>,
+    pub(crate) key: jmt::KeyHash,
+    pub(crate) value: jmt::OwnedValue,
+}
+
+pub struct ExclusionProof {
+    pub(crate) proof: SparseMerkleProof<sha2::Sha256>,
+    pub(crate) key: jmt::KeyHash,
+}
+
+pub enum MerkleProof {
+    Inclusion(InclusionProof),
+    Exclusion(ExclusionProof),
+}
+
+impl MerkleProof {
+    pub fn verify(&self, root_hash: Bytes32) -> bool {
+        match self {
+            MerkleProof::Inclusion(inclusion_proof) => {
+                let root_hash = jmt::RootHash(root_hash);
+                let key = inclusion_proof.key;
+                let value = &inclusion_proof.value;
+                let proof = &inclusion_proof.proof;
+
+                proof.verify_existence(root_hash, key, value).is_ok()
+            }
+            MerkleProof::Exclusion(exclusion_proof) => {
+                let root_hash = jmt::RootHash(root_hash);
+                let key = exclusion_proof.key;
+                let proof = &exclusion_proof.proof;
+
+                proof.verify_nonexistence(root_hash, key).is_ok()
+            }
+        }
+    }
+
+    pub fn is_inclusion_proof(&self) -> bool {
+        match self {
+            MerkleProof::Inclusion(_) => true,
+            MerkleProof::Exclusion(_) => false,
+        }
+    }
+
+    pub fn is_exclusion_proof(&self) -> bool {
+        match self {
+            MerkleProof::Inclusion(_) => false,
+            MerkleProof::Exclusion(_) => true,
+        }
+    }
+}