diff --git a/Cargo.lock b/Cargo.lock index fb91136e5..7c0480f5c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3211,8 +3211,7 @@ checksum = "af150ab688ff2122fcef229be89cb50dd66af9e01a4ff320cc137eecc9bacc38" [[package]] name = "jmt" -version = "0.6.0" -source = "git+https://github.com/penumbra-zone/jmt#46b4b0042f66506536097d689ac2201e3e430afd" +version = "0.7.0" dependencies = [ "anyhow", "borsh", diff --git a/Cargo.toml b/Cargo.toml index d5a0fbf9d..ebebd30ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -54,8 +54,7 @@ rust-version = "1.66" [workspace.dependencies] # Dependencies maintained by sovereign -# TODO: replace by release number once available on crates.io: tracking issue https://github.com/Sovereign-Labs/sovereign-sdk/issues/632 -jmt = { git = "https://github.com/penumbra-zone/jmt", commit = "46b4b00" } +jmt = "0.7.0" # External dependencies async-trait = "0.1.71" @@ -112,4 +111,5 @@ secp256k1 = { version = "0.27.0", default-features = false, features = ["global- [patch.crates-io] # See reth: https://github.com/paradigmxyz/reth/blob/main/Cargo.toml#L79 revm = { git = "https://github.com/bluealloy/revm/", branch = "release/v25" } -revm-primitives = { git = "https://github.com/bluealloy/revm/", branch = "release/v25" } \ No newline at end of file +revm-primitives = { git = "https://github.com/bluealloy/revm/", branch = "release/v25" } +jmt = { path = "../penumbra-jmt/jmt" } diff --git a/module-system/sov-state/src/lib.rs b/module-system/sov-state/src/lib.rs index 360f6167d..88de994a0 100644 --- a/module-system/sov-state/src/lib.rs +++ b/module-system/sov-state/src/lib.rs @@ -5,9 +5,6 @@ mod map; #[cfg(feature = "native")] mod prover_storage; -#[cfg(feature = "native")] -mod tree_db; - mod scratchpad; pub mod storage; @@ -37,7 +34,7 @@ pub use storage::Storage; use utils::AlignedVec; pub use value::StateValue; -pub use crate::witness::{ArrayWitness, TreeWitnessReader, Witness}; +pub use crate::witness::{ArrayWitness, Witness}; // A prefix prepended to each key before insertion and retrieval from the storage. // All the collection types in this crate are backed by the same storage instance, this means that insertions of the same key diff --git a/module-system/sov-state/src/prover_storage.rs b/module-system/sov-state/src/prover_storage.rs index 28fe20839..f50f4ca12 100644 --- a/module-system/sov-state/src/prover_storage.rs +++ b/module-system/sov-state/src/prover_storage.rs @@ -10,7 +10,6 @@ use sov_db::state_db::StateDB; use crate::config::Config; use crate::internal_cache::OrderedReadsAndWrites; use crate::storage::{NativeStorage, StorageKey, StorageProof, StorageValue}; -use crate::tree_db::TreeReadLogger; use crate::witness::Witness; use crate::{MerkleProofSpec, Storage}; @@ -85,19 +84,14 @@ impl Storage for ProverStorage { witness: &Self::Witness, ) -> Result<[u8; 32], anyhow::Error> { let latest_version = self.db.get_next_version() - 1; - witness.add_hint(latest_version); - let read_logger = TreeReadLogger::with_db_and_witness(self.db.clone(), witness); - let untracked_jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&self.db); + let jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&self.db); // Handle empty untracked_jmt - if untracked_jmt - .get_root_hash_option(latest_version)? - .is_none() - { + if jmt.get_root_hash_option(latest_version)?.is_none() { assert_eq!(latest_version, 0); let empty_batch = Vec::default().into_iter(); - let (_, tree_update) = untracked_jmt + let (_, tree_update) = jmt .put_value_set(empty_batch, latest_version) .expect("JMT update must succeed"); @@ -110,14 +104,13 @@ impl Storage for ProverStorage { for (key, read_value) in state_accesses.ordered_reads { let key_hash = KeyHash::with::(key.key.as_ref()); // TODO: Switch to the batch read API once it becomes available - let (result, proof) = untracked_jmt.get_with_proof(key_hash, latest_version)?; + let (result, proof) = jmt.get_with_proof(key_hash, latest_version)?; if result.as_ref() != read_value.as_ref().map(|f| f.value.as_ref()) { anyhow::bail!("Bug! Incorrect value read from jmt"); } witness.add_hint(proof); } - let tracked_jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&read_logger); // Compute the jmt update from the write batch let batch = state_accesses .ordered_writes @@ -135,13 +128,17 @@ impl Storage for ProverStorage { let next_version = self.db.get_next_version(); - let (new_root, tree_update) = tracked_jmt - .put_value_set(batch, next_version) + let (new_root, update_proof, tree_update) = jmt + .put_value_set_with_proof(batch, next_version) .expect("JMT update must succeed"); self.db .write_node_batch(&tree_update.node_batch) .expect("db write must succeed"); + + witness.add_hint(update_proof); + witness.add_hint(&new_root.0); + self.db.inc_next_version(); Ok(new_root.0) } diff --git a/module-system/sov-state/src/tree_db.rs b/module-system/sov-state/src/tree_db.rs deleted file mode 100644 index d41e14d48..000000000 --- a/module-system/sov-state/src/tree_db.rs +++ /dev/null @@ -1,50 +0,0 @@ -use borsh::BorshSerialize; -use jmt::storage::TreeReader; -use jmt::OwnedValue; -use sov_db::state_db::StateDB; - -use crate::witness::Witness; - -pub struct TreeReadLogger<'a, W> { - state_db: StateDB, - witness: &'a W, -} - -impl<'a, W: Witness> TreeReadLogger<'a, W> { - /// Creates a tree read logger wrapping the provided StateDB. - /// The logger is recording by default - pub fn with_db_and_witness(db: StateDB, witness: &'a W) -> Self { - Self { - state_db: db, - witness, - } - } -} - -impl<'a, W: Witness> TreeReader for TreeReadLogger<'a, W> { - fn get_node_option( - &self, - node_key: &jmt::storage::NodeKey, - ) -> anyhow::Result> { - let node_opt = self.state_db.get_node_option(node_key)?; - self.witness - .add_hint(node_opt.as_ref().map(|node| node.try_to_vec().unwrap())); - Ok(node_opt) - } - - fn get_value_option( - &self, - max_version: jmt::Version, - key_hash: jmt::KeyHash, - ) -> anyhow::Result> { - let value_opt = self.state_db.get_value_option(max_version, key_hash)?; - self.witness.add_hint(value_opt.clone()); - Ok(value_opt) - } - - fn get_rightmost_leaf( - &self, - ) -> anyhow::Result> { - unimplemented!() - } -} diff --git a/module-system/sov-state/src/witness.rs b/module-system/sov-state/src/witness.rs index eaa202cda..bce2e77a2 100644 --- a/module-system/sov-state/src/witness.rs +++ b/module-system/sov-state/src/witness.rs @@ -2,7 +2,6 @@ use std::sync::atomic::AtomicUsize; use std::sync::Mutex; use borsh::{BorshDeserialize, BorshSerialize}; -use jmt::storage::TreeReader; use serde::{Deserialize, Serialize}; // TODO: Refactor witness trait so it only require Serialize / Deserialize @@ -13,42 +12,6 @@ pub trait Witness: Default + Serialize { fn merge(&self, rhs: &Self); } -#[derive(Debug)] -pub struct TreeWitnessReader<'a, T: Witness>(&'a T); - -impl<'a, T: Witness> TreeWitnessReader<'a, T> { - pub fn new(witness: &'a T) -> Self { - Self(witness) - } -} - -impl<'a, T: Witness> TreeReader for TreeWitnessReader<'a, T> { - fn get_node_option( - &self, - _node_key: &jmt::storage::NodeKey, - ) -> anyhow::Result> { - let serialized_node_opt: Option> = self.0.get_hint(); - match serialized_node_opt { - Some(val) => Ok(Some(jmt::storage::Node::deserialize_reader(&mut &val[..])?)), - None => Ok(None), - } - } - - fn get_value_option( - &self, - _max_version: jmt::Version, - _key_hash: jmt::KeyHash, - ) -> anyhow::Result> { - Ok(self.0.get_hint()) - } - - fn get_rightmost_leaf( - &self, - ) -> anyhow::Result> { - unimplemented!() - } -} - #[derive(Default, Debug, Serialize, Deserialize)] pub struct ArrayWitness { next_idx: AtomicUsize, diff --git a/module-system/sov-state/src/zk_storage.rs b/module-system/sov-state/src/zk_storage.rs index 9e4ba55ba..a991f8924 100644 --- a/module-system/sov-state/src/zk_storage.rs +++ b/module-system/sov-state/src/zk_storage.rs @@ -1,13 +1,13 @@ use std::marker::PhantomData; use std::sync::Arc; -use jmt::{JellyfishMerkleTree, KeyHash, Version}; +use jmt::{KeyHash, RootHash}; #[cfg(all(target_os = "zkvm", feature = "bench"))] use zk_cycle_macros::cycle_tracker; use crate::internal_cache::OrderedReadsAndWrites; use crate::storage::{StorageKey, StorageProof, StorageValue}; -use crate::witness::{TreeWitnessReader, Witness}; +use crate::witness::Witness; use crate::{MerkleProofSpec, Storage}; #[cfg(all(target_os = "zkvm", feature = "bench"))] @@ -61,9 +61,6 @@ impl Storage for ZkStorage { state_accesses: OrderedReadsAndWrites, witness: &Self::Witness, ) -> Result<[u8; 32], anyhow::Error> { - let latest_version: Version = witness.get_hint(); - let reader = TreeWitnessReader::new(witness); - // For each value that's been read from the tree, verify the provided smt proof for (key, read_value) in state_accesses.ordered_reads { let key_hash = KeyHash::with::(key.key.as_ref()); @@ -79,6 +76,9 @@ impl Storage for ZkStorage { } } + let update_proof: jmt::proof::UpdateMerkleProof = witness.get_hint(); + let new_root = RootHash(witness.get_hint()); + // Compute the jmt update from the write batch let batch = state_accesses .ordered_writes @@ -89,16 +89,15 @@ impl Storage for ZkStorage { key_hash, value.map(|v| Arc::try_unwrap(v.value).unwrap_or_else(|arc| (*arc).clone())), ) - }); + }) + .collect::>(); - let next_version = latest_version + 1; // TODO: Make updates verifiable. Currently, writes don't verify that the provided siblings existed in the old tree // because the TreeReader is trusted - let jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&reader); + update_proof + .verify_update(RootHash(self.prev_state_root), new_root, batch) + .expect("Update proof was invalid! The prover was malicious"); - let (new_root, _tree_update) = jmt - .put_value_set(batch, next_version) - .expect("JMT update must succeed"); Ok(new_root.0) }