NOMT storage support for open proof

crates/module-system/sov-modules-api/tests/integration/proof_tests.rs

@@ -106,7 +106,6 @@ mod map {
     use super::{make_user_map_proof, S};
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_roundtrip() {
         let (root, proof, map) = make_user_map_proof(1, 2);
         let (key, val) = map.verify_proof::<S>(root, proof).unwrap();
@@ -115,31 +114,27 @@ mod map {
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_wrong_namespace() {
         let (root, mut proof, map) = make_user_map_proof(1, 2);
         proof.namespace = ProvableNamespace::Kernel;
         assert!(map.verify_proof::<S>(root, proof).is_err());
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_wrong_key() {
         let (root, mut proof, map) = make_user_map_proof(1, 2);
         proof.key = SlotKey::new(&Prefix::new(1, 0), &1, map.codec());
         assert!(map.verify_proof::<S>(root, proof).is_err());
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_missing_value() {
         let (root, mut proof, map) = make_user_map_proof(1, 2);
         proof.value = None;
         assert!(map.verify_proof::<S>(root, proof).is_err());
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_wrong_value() {
         let (root, mut proof, map) = make_user_map_proof(1, 2);
         proof.value = Some(SlotValue::new(&3, map.codec()));
@@ -154,39 +149,34 @@ mod value {
     use super::{make_user_value_proof, S};
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_roundtrip() {
         let (root, proof, map) = make_user_value_proof(1);
         let val = map.verify_proof::<S>(root, proof).unwrap();
         assert_eq!(val, Some(1));
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_wrong_namespace() {
         let (root, mut proof, map) = make_user_value_proof(1);
         proof.namespace = ProvableNamespace::Kernel;
         assert!(map.verify_proof::<S>(root, proof).is_err());
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_wrong_key() {
         let (root, mut proof, map) = make_user_value_proof(1);
         proof.key = SlotKey::new(&Prefix::new(255, 0), &1, map.codec()); // Use the wrong prefix
         assert!(map.verify_proof::<S>(root, proof).is_err());
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_missing_value() {
         let (root, mut proof, map) = make_user_value_proof(1);
         proof.value = None;
         assert!(map.verify_proof::<S>(root, proof).is_err());
     }
 
     #[test]
-    #[ignore = "NOMT does not support open_proof yet"]
     fn test_state_proof_wrong_value() {
         let (root, mut proof, map) = make_user_value_proof(1);
         proof.value = Some(SlotValue::new(&3, map.codec()));
@@ -195,7 +185,7 @@ mod value {
 }
 
 #[test]
-#[ignore = "NOMT does not support open_proof yet"]
+#[ignore = "NOMT does not support archival proof generation - proofs are always generated against the current state"]
 fn test_archival_proof_gen() {
     let mut kernel = MockKernel::<S>::default();
     let mut storage_manager = SimpleNomtStorageManager::new();

crates/module-system/sov-state/src/nomt/mod.rs

@@ -1,5 +1,24 @@
 //! Storage implementation based on [Nearly Optimal Merkle Tree(NOMT)](https://github.com/thrumdev/nomt/) implementation.
 
+use borsh::{BorshDeserialize, BorshSerialize};
+use nomt_core::proof::MultiProof;
+use serde::Serialize;
+
 #[cfg(feature = "native")]
 pub mod prover_storage;
 pub mod zk_storage;
+
+/// Wrapper around [`nomt_core::proof::MultiProof`] that implements `PartialEq`/`Eq`
+/// by comparing borsh-serialized representations.
+///
+/// This is necessary because `MultiProof` doesn't derive these traits.
+#[derive(Debug, Clone, Serialize, serde::Deserialize, BorshSerialize, BorshDeserialize)]
+pub struct NomtMultiProof(pub MultiProof);
+
+impl PartialEq for NomtMultiProof {
+    fn eq(&self, other: &Self) -> bool {
+        borsh::to_vec(&self.0).ok() == borsh::to_vec(&other.0).ok()
+    }
+}
+
+impl Eq for NomtMultiProof {}

crates/module-system/sov-state/src/nomt/prover_storage.rs

@@ -9,6 +9,7 @@ use anyhow::Context;
 use nomt::hasher::BinaryHasher;
 use nomt::proof::MultiProof;
 use nomt::FinishedSession;
+use nomt_core::trie::{KeyPath, LeafData, Node, ValueHash};
 use sov_db::accessory_db::AccessoryDb;
 use sov_db::historical_state::HistoricalStateReader;
 use sov_db::state_db_nomt::{HistoricalValueError, NomtSessionBuilder, SessionsContainer};
@@ -18,12 +19,14 @@ use sov_db::storage_manager::{
 use sov_rollup_interface::common::SlotNumber;
 use sov_rollup_interface::reexports::digest::Digest;
 
+use crate::nomt::NomtMultiProof;
 use crate::pinned_cache::PinnedCache;
 use crate::storage::ReadType;
 use crate::{
     Accessory, CompileTimeNamespace, MerkleProofSpec, Namespace, NativeStorage, NodeLeaf,
     NodeLeafAndMaybeValue, OrderedReadsAndWrites, ProvableCompileTimeNamespace, ProvableNamespace,
-    SlotKey, SlotValue, StateAccesses, StateUpdate, Storage, StorageProof, StorageRoot, Witness,
+    SlotKey, SlotValue, StateAccesses, StateRoot, StateUpdate, Storage, StorageProof, StorageRoot,
+    Witness,
 };
 
 type NomtSession<H> = nomt::Session<BinaryHasher<H>>;
@@ -401,7 +404,7 @@ where
 {
     type Hasher = S::Hasher;
     type Witness = S::Witness;
-    type Proof = ();
+    type Proof = NomtMultiProof;
     type Root = StorageRoot<S>;
     // These 2 are effectively the same thing, `StateUpdate` is not materialized, `ChangeSet` is materialized.
     type StateUpdate = NomtStateUpdate<S>;
@@ -617,10 +620,51 @@ where
     }
 
     fn open_proof(
-        _state_root: Self::Root,
-        _proof: StorageProof<Self::Proof>,
+        state_root: Self::Root,
+        proof: StorageProof<Self::Proof>,
     ) -> anyhow::Result<(SlotKey, Option<SlotValue>)> {
-        unimplemented!("The NomtProverStorage does not support `open_proof` yet.")
+        let StorageProof {
+            key,
+            value,
+            proof: multi_proof,
+            namespace,
+        } = proof;
+        let root_node: Node = state_root.namespace_root(namespace);
+
+        let verified = nomt_core::proof::verify_multi_proof::<BinaryHasher<S::Hasher>>(
+            &multi_proof.0,
+            root_node,
+        )
+        .map_err(|e| anyhow::anyhow!("Failed to verify proof: {:?}", e))?;
+
+        let key_path: KeyPath = S::Hasher::digest(key.as_ref()).into();
+
+        match &value {
+            None => {
+                if !verified
+                    .confirm_nonexistence(&key_path)
+                    .map_err(|e| anyhow::anyhow!("Key out of scope: {:?}", e))?
+                {
+                    anyhow::bail!("Failed to verify non-existence of key");
+                }
+            }
+            Some(slot_value) => {
+                let authenticated_write = slot_value.combine_val_hash_and_size::<S::Hasher>();
+                let value_hash: ValueHash = S::Hasher::digest(&authenticated_write).into();
+                let leaf = LeafData {
+                    key_path,
+                    value_hash,
+                };
+                if !verified
+                    .confirm_value(&leaf)
+                    .map_err(|e| anyhow::anyhow!("Key out of scope: {:?}", e))?
+                {
+                    anyhow::bail!("Failed to verify value for key");
+                }
+            }
+        }
+
+        Ok((key, value))
     }
 }
 
@@ -675,11 +719,23 @@ where
             ProvableNamespace::Kernel => self.read_value::<crate::Kernel>(&key, slot_number)?,
         };
 
+        let session = match namespace {
+            ProvableNamespace::User => self
+                .state_session_builder
+                .begin_user_session_without_witness()?,
+            ProvableNamespace::Kernel => self
+                .state_session_builder
+                .begin_kernel_session_without_witness()?,
+        };
+
+        let key_path: KeyPath = S::Hasher::digest(key.as_ref()).into();
+        let path_proof = session.prove(key_path)?;
+        let multi_proof = MultiProof::from_path_proofs(vec![path_proof]);
+
         Ok(StorageProof {
             key,
             value,
-            // TODO: Proof is empty now, will be fixed in follow
-            proof: (),
+            proof: NomtMultiProof(multi_proof),
             namespace,
         })
     }

crates/module-system/sov-state/src/nomt/zk_storage.rs

@@ -8,6 +8,7 @@ use nomt_core::trie::{KeyPath, LeafData, Node, ValueHash};
 use sov_rollup_interface::common::SlotNumber;
 use sov_rollup_interface::reexports::digest::Digest;
 
+use crate::nomt::NomtMultiProof;
 use crate::pinned_cache::PinnedCache;
 use crate::storage::ReadType;
 use crate::{
@@ -106,7 +107,7 @@ impl<S: MerkleProofSpec> NomtVerifierStorage<S> {
 impl<S: MerkleProofSpec> Storage for NomtVerifierStorage<S> {
     type Hasher = S::Hasher;
     type Witness = S::Witness;
-    type Proof = ();
+    type Proof = NomtMultiProof;
     type Root = StorageRoot<S>;
     type StateUpdate = ();
     type ChangeSet = ();
@@ -160,10 +161,51 @@ impl<S: MerkleProofSpec> Storage for NomtVerifierStorage<S> {
     fn materialize_changes(self, _state_update: Self::StateUpdate) -> Self::ChangeSet {}
 
     fn open_proof(
-        _state_root: Self::Root,
-        _proof: StorageProof<Self::Proof>,
+        state_root: Self::Root,
+        proof: StorageProof<Self::Proof>,
     ) -> anyhow::Result<(SlotKey, Option<SlotValue>)> {
-        unimplemented!("The NomtZkStorage does not support `open_proof` yet.");
+        let StorageProof {
+            key,
+            value,
+            proof: multi_proof,
+            namespace,
+        } = proof;
+        let root_node: Node = state_root.namespace_root(namespace);
+
+        let verified = nomt_core::proof::verify_multi_proof::<BinaryHasher<S::Hasher>>(
+            &multi_proof.0,
+            root_node,
+        )
+        .map_err(|e| anyhow::anyhow!("Failed to verify proof: {:?}", e))?;
+
+        let key_path: KeyPath = S::Hasher::digest(key.as_ref()).into();
+
+        match &value {
+            None => {
+                if !verified
+                    .confirm_nonexistence(&key_path)
+                    .map_err(|e| anyhow::anyhow!("Key out of scope: {:?}", e))?
+                {
+                    anyhow::bail!("Failed to verify non-existence of key");
+                }
+            }
+            Some(slot_value) => {
+                let authenticated_write = slot_value.combine_val_hash_and_size::<S::Hasher>();
+                let value_hash: ValueHash = S::Hasher::digest(&authenticated_write).into();
+                let leaf = LeafData {
+                    key_path,
+                    value_hash,
+                };
+                if !verified
+                    .confirm_value(&leaf)
+                    .map_err(|e| anyhow::anyhow!("Key out of scope: {:?}", e))?
+                {
+                    anyhow::bail!("Failed to verify value for key");
+                }
+            }
+        }
+
+        Ok((key, value))
     }
 }