diff --git a/crates/module-system/sov-state/src/nomt/prover_storage.rs b/crates/module-system/sov-state/src/nomt/prover_storage.rs index 09fa20da5f..d94093b65b 100644 --- a/crates/module-system/sov-state/src/nomt/prover_storage.rs +++ b/crates/module-system/sov-state/src/nomt/prover_storage.rs @@ -537,21 +537,7 @@ fn compute_state_update_namespace( let mut finished = session.finish(accesses)?; if write_witness { let nomt_witness = finished.take_witness().expect("Witness cannot be missing"); - let nomt::Witness { - path_proofs, - operations: nomt::WitnessedOperations { .. }, - } = nomt_witness; - // Note, we discard `p.path`, but maybe there's a way to use to have more efficient verification? - let mut path_proofs_inner = path_proofs.into_iter().map(|p| p.inner).collect::>(); - - // Sort them as required by - // Note that the path proofs produced within a crate::witness::Witness are not guaranteed to be ordered, - // so the input should be sorted lexicographically by the terminal path prior to calling this function. - // https://github.com/thrumdev/nomt/issues/904 - path_proofs_inner.sort_by(|a, b| a.terminal.path().cmp(b.terminal.path())); - - let multi_proof = MultiProof::from_path_proofs(path_proofs_inner); - witness.add_hint(&multi_proof); + witness.add_hint(&nomt_witness); } Ok(finished) } diff --git a/crates/module-system/sov-state/src/nomt/zk_storage.rs b/crates/module-system/sov-state/src/nomt/zk_storage.rs index 26bbf867a1..4dfa97553d 100644 --- a/crates/module-system/sov-state/src/nomt/zk_storage.rs +++ b/crates/module-system/sov-state/src/nomt/zk_storage.rs @@ -1,9 +1,11 @@ //! ZK Verifier part of the NOMT based Storage implementation +use std::collections::BTreeMap; use std::marker::PhantomData; use nomt_core::hasher::BinaryHasher; -use nomt_core::proof::MultiProof; +use nomt_core::proof::{MultiProof, PathUpdate}; use nomt_core::trie::{KeyPath, LeafData, Node, ValueHash}; +use nomt_core::witness::Witness as NomtWitness; #[cfg(all(feature = "test-utils", feature = "native"))] use sov_rollup_interface::common::SlotNumber; use sov_rollup_interface::reexports::digest::Digest; @@ -41,69 +43,173 @@ impl NomtVerifierStorage { ordered_reads: state_reads, ordered_writes: state_writes, } = state_accesses; + let mut expected_reads = BTreeMap::new(); + for (key, value) in state_reads { + let key_hash: KeyPath = S::Hasher::digest(key.as_ref()).into(); + let value_hash = value.map(|node_leaf| { + S::Hasher::digest(node_leaf.combine_val_hash_and_size()).into() + }); + if expected_reads.insert(key_hash, value_hash).is_some() { + anyhow::bail!("Duplicate key read in state accesses: {:?}", key); + } + } + + let mut expected_writes = BTreeMap::new(); + for (key, value) in state_writes { + let key_hash: KeyPath = S::Hasher::digest(key.as_ref()).into(); + let value_hash = value.map(|slot_value| { + // Authenticated write is hash of a combination of size and original value hash. + S::Hasher::digest(slot_value.combine_val_hash_and_size::()).into() + }); + if expected_writes.insert(key_hash, value_hash).is_some() { + anyhow::bail!("Duplicate key write in state accesses: {:?}", key); + } + } - let multi_proof: MultiProof = array_witness.get_hint(); + let nomt_witness: NomtWitness = array_witness.get_hint(); + let mut path_proofs_inner = nomt_witness + .path_proofs + .iter() + .map(|path| path.inner.clone()) + .collect::>(); + path_proofs_inner.sort_by(|a, b| a.terminal.path().cmp(b.terminal.path())); + let multi_proof = MultiProof::from_path_proofs(path_proofs_inner); let verified_multi_proof = nomt_core::proof::verify_multi_proof::>( &multi_proof, prev_root, ) .map_err(|e| anyhow::anyhow!("Failed to verify multi proof: {:?}", e))?; - for (key, value) in state_reads { - let key_hash: KeyPath = S::Hasher::digest(key.as_ref()).into(); + for (key, value) in &expected_reads { match value { None => { if !verified_multi_proof - .confirm_nonexistence(&key_hash) + .confirm_nonexistence(key) .map_err(|e| anyhow::anyhow!("Failed to confirm non-existence: {:?}", e))? { - anyhow::bail!("Failed to verify non-existence of key: {:?}", key); + anyhow::bail!("Failed to verify non-existence of key"); } } - Some(node_leaf) => { - let authenticated_write = node_leaf.combine_val_hash_and_size(); - let value_hash = S::Hasher::digest(&authenticated_write).into(); + Some(value_hash) => { let leaf = LeafData { - key_path: key_hash, - value_hash, + key_path: *key, + value_hash: *value_hash, }; if !verified_multi_proof .confirm_value(&leaf) .map_err(|e| anyhow::anyhow!("Failed to confirm value: {:?}", e))? { - anyhow::bail!("Failed to verify inclusion of key: {:?}", key); + anyhow::bail!("Failed to verify inclusion of key"); } } } } - let mut updates = state_writes + let mut witnessed_writes = BTreeMap::new(); + for write in nomt_witness.operations.writes { + let Some(expected_value) = expected_writes.remove(&write.key) else { + anyhow::bail!("Unexpected or duplicate NOMT write witness entry"); + }; + if write.value != expected_value { + anyhow::bail!( + "Mismatched NOMT write witness value for key {:?}: witness={:?}, expected={:?}", + write.key, + write.value, + expected_value + ); + } + if witnessed_writes + .insert(write.key, expected_value) + .is_some() + { + anyhow::bail!("Duplicate NOMT write witness entry"); + }; + } + + if !expected_writes.is_empty() { + anyhow::bail!("Missing NOMT write witness entries"); + } + + let mut verified_paths = nomt_witness + .path_proofs .into_iter() - .map(|(key, value)| { - ( - S::Hasher::digest(key.as_ref()).into(), - value.map(|slot_value| { - // Authenticated write is hash of a combination of size and orignal value hash. - S::Hasher::digest(slot_value.combine_val_hash_and_size::()) - .into() - }), - ) + .map(|path| { + let verified = path + .inner + .verify::>(path.path.path(), prev_root) + .map_err(|e| anyhow::anyhow!("Failed to verify path proof: {:?}", e))?; + Ok(( + verified, + path.path.depth() as usize, + truncate_key_path(path.path.raw_path(), path.path.depth() as usize), + Vec::<(KeyPath, Option)>::new(), + )) }) - .collect::)>>(); + .collect::>>()?; + verified_paths.sort_by(|a, b| a.0.path().cmp(b.0.path())); + + let mut path_index_by_prefix = BTreeMap::new(); + let mut path_depths_desc = Vec::new(); + for (index, (_, depth, prefix, _)) in verified_paths.iter().enumerate() { + if path_index_by_prefix + .insert((*depth, *prefix), index) + .is_some() + { + anyhow::bail!("Duplicate NOMT path proof prefix"); + } + path_depths_desc.push(*depth); + } + path_depths_desc.sort_unstable(); + path_depths_desc.dedup(); + path_depths_desc.reverse(); + + for (key, value) in witnessed_writes { + let matching_index = path_depths_desc + .iter() + .find_map(|depth| { + path_index_by_prefix + .get(&(*depth, truncate_key_path(key, *depth))) + .copied() + }) + .ok_or_else(|| anyhow::anyhow!("No NOMT path proof covers write key {:?}", key))?; + + verified_paths[matching_index].3.push((key, value)); + } - // Sort them by key hash, as required by [`nomt_core::proof::verify_multi_proof_update`] - updates.sort_by(|a, b| a.0.cmp(&b.0)); + let mut updates = Vec::new(); + for (verified, _, _, writes) in verified_paths { + if !writes.is_empty() { + updates.push(PathUpdate { + inner: verified, + ops: writes, + }); + } + } - nomt_core::proof::verify_multi_proof_update::>( - &verified_multi_proof, - updates, - ) - .map_err(|e| anyhow::anyhow!("Failed to verify update: {:?}", e)) - // Note: we don't check exhaustion of the proof - // because it does not impact the correctness of the guest, only performance. + nomt_core::proof::verify_update::>(prev_root, &updates) + .map_err(|e| anyhow::anyhow!("Failed to verify update: {:?}", e)) } } +fn truncate_key_path(mut key: KeyPath, depth: usize) -> KeyPath { + debug_assert!(depth <= 256); + + let full_bytes = depth / 8; + let partial_bits = depth % 8; + + if full_bytes < key.len() { + if partial_bits == 0 { + key[full_bytes..].fill(0); + } else { + let keep_mask = u8::MAX << (8 - partial_bits); + key[full_bytes] &= keep_mask; + key[full_bytes + 1..].fill(0); + } + } + + key +} + impl Storage for NomtVerifierStorage { type Hasher = S::Hasher; type Witness = S::Witness; diff --git a/examples/demo-rollup/tests/prover/sp1_cpu_prover.rs b/examples/demo-rollup/tests/prover/sp1_cpu_prover.rs index 73621b36b5..d04f4312ab 100644 --- a/examples/demo-rollup/tests/prover/sp1_cpu_prover.rs +++ b/examples/demo-rollup/tests/prover/sp1_cpu_prover.rs @@ -11,6 +11,7 @@ use sov_rollup_interface::zk::{ }; use sov_sp1_adapter::host::{MockSp1Prover, SP1Host}; use sov_sp1_adapter::{SP1MethodId, SP1Verifier}; +use tokio::time::Instant; type ProofInput = StateTransitionWitnessWithAddress< ::Address, @@ -55,7 +56,7 @@ async fn generate_proofs(host: &TestHost) -> Vec::Address::try_from([0u8; 28].as_ref()).unwrap(); let mut proofs = Vec::new(); - for witness in witnesses { + for (i, witness) in witnesses.into_iter().enumerate() { let da_block_header = witness.da_block_header.clone(); let data: ProofInput = StateTransitionWitnessWithAddress { @@ -63,7 +64,9 @@ async fn generate_proofs(host: &TestHost) -> Vec Vec::Address::try_from([0u8; 28].as_ref()).unwrap(); + + for witness in witnesses { + let expected_initial_state_root = witness.initial_state_root.clone(); + let expected_final_state_root = witness.final_state_root.clone(); + let expected_slot_hash = witness.da_block_header.hash(); + + let public_data = host + .run_mock_public_data(ProofInput { + stf_witness: witness, + prover_address, + }) + .await; + + assert_eq!(public_data.initial_state_root, expected_initial_state_root); + assert_eq!(public_data.final_state_root, expected_final_state_root); + assert_eq!(public_data.slot_hash, expected_slot_hash); + } +} + // The SP1 prover manages its own Tokio runtime, which conflicts with the `tokio::test` runtime. // To avoid this, all blocking work must be executed inside `tokio::task::spawn_blocking`. struct TestHost { @@ -129,6 +156,25 @@ impl TestHost { .unwrap() } } + + async fn run_mock_public_data( + &self, + data: ProofInput, + ) -> StateTransitionPublicData<::Address, MockDaSpec, ProofStateRoot> { + tokio::task::spawn_blocking(move || { + let mock_host = MockSp1Prover::new(*sp1::SP1_GUEST_MOCK_ELF) + .expect("MockSp1Prover should be created successfully"); + + let proof = mock_host + .add_hint_and_run(&data) + .expect("Mock prover should run successfully"); + + bincode::deserialize(proof.public_values.as_slice()) + .expect("Mock proof public values should deserialize") + }) + .await + .unwrap() + } } #[allow(dead_code)]