From 1c9567e5e45fb2e64b385ad537b585bf09512299 Mon Sep 17 00:00:00 2001 From: Neelay Sant <52956732+NSant215@users.noreply.github.com> Date: Fri, 15 May 2026 15:08:38 +0100 Subject: [PATCH] test(dursto): add integration test for database for roundtrip of proofs --- durable-storage/src/database.rs | 167 +++++++++++++++++++++++- durable-storage/src/lib.rs | 2 + durable-storage/src/merkle_layer.rs | 71 ++-------- durable-storage/src/proof_test_utils.rs | 128 ++++++++++++++++++ 4 files changed, 305 insertions(+), 63 deletions(-) create mode 100644 durable-storage/src/proof_test_utils.rs diff --git a/durable-storage/src/database.rs b/durable-storage/src/database.rs index a9c67145ed..ab08d73597 100644 --- a/durable-storage/src/database.rs +++ b/durable-storage/src/database.rs @@ -584,7 +584,7 @@ impl Database { } } -mod traced_database; +pub(crate) mod traced_database; #[cfg(test)] pub(crate) mod tests { @@ -613,9 +613,12 @@ pub(crate) mod tests { use crate::merkle_layer::MerkleLayer; use crate::merkle_worker::BackgroundKeyValueStore; use crate::merkle_worker::BackgroundPersistentKeyValueStore; + use crate::proof_test_utils::apply_to_database; + use crate::proof_test_utils::setup_data_len; use crate::storage::KeyValueStore; use crate::storage::TestKeyValueStoreSetup; use crate::storage::kv_test; + use crate::test_helpers::Operation; fn new_database( handle: &Handle, @@ -2197,6 +2200,168 @@ pub(crate) mod tests { database.into_trace() }); + // Multi-step proof: starting from a populated state, run a sequence of independent proof + // batches. Each batch + // - calls `try_start_proof` on the current Normal-mode `Database`, + // - applies operations through the Prove-mode `Database` API, + // - derives a `Database` from the generated proof, and replays the same operations. + // - The Normal-mode source-of-truth is then advanced with the same operations so the next + // batch's initial state reflects the accumulated mutations. + // + // Clone operations (e.g. `try_clone_with`) are intentionally excluded — only set, write, + // delete, and read are exercised. + kv_test!(test_database_multi_step_prove_verify, KV: BackgroundKeyValueStore, + setup_runtime |handle, repo| = { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .build() + .expect("Creating a Tokio runtime should succeed"); + let handle = runtime.handle().clone(); + let (_keepalive, repo) = KV::setup_repo(); + (runtime, handle, _keepalive, repo) + }, + [ + setup_keys in prop::collection::vec(any::<[u8; 2]>(), 1..15), + batch_seeds in prop::collection::vec(0usize..1000, 2..6), + ], { + let mut normal_db = new_database::(handle, repo); + + let keys: Vec = setup_keys + .iter() + .map(|b| Key::new(b).expect("key should be valid")) + .collect(); + + // Track each key's current data length so generated ops stay valid as state evolves + // across steps. Mirrors the strategy used in `prove_verify_round_trips_mixed`. + let mut lengths: std::collections::HashMap = + std::collections::HashMap::new(); + for (bytes, key) in setup_keys.iter().zip(keys.iter()) { + let data_len = setup_data_len(bytes); + let data = Bytes::from(vec![bytes[1]; data_len]); + normal_db + .set(key.clone(), data) + .expect("setup set should succeed"); + lengths.insert(key.clone(), data_len); + } + + for &seed in batch_seeds.iter() { + // Build a small batch of ops for this batch. Mix of set/delete/write/read so we + // exercise structural and value-data accesses across the proof. + // + // Generated reads stay within `[0, len]` because `Database::read_bytes` errors on + // out-of-range offsets. + let ops_count = (seed % 8) + 1; + let mut step_ops: Vec = Vec::with_capacity(ops_count); + for i in 0..ops_count { + let key = keys[(seed + i) % keys.len()].clone(); + let data = Bytes::from(vec![(seed + i) as u8; 1 + ((seed + i) % 5)]); + match (seed + i) % 4 { + 0 => { + lengths.insert(key.clone(), data.len()); + step_ops.push(Operation::Set(key, data)); + } + 1 => { + lengths.remove(&key); + step_ops.push(Operation::Delete(key)); + } + 2 => match lengths.get(&key).copied() { + Some(len) => { + let offset = (seed + i) % (len + 1); + lengths.insert( + key.clone(), + std::cmp::max(len, offset + data.len()), + ); + step_ops.push(Operation::Write(key, offset, data)); + } + None => { + // Writing to a non-existent key at non-zero offset would fail; fall + // back to a Set instead so the op stays valid. + lengths.insert(key.clone(), data.len()); + step_ops.push(Operation::Set(key, data)); + } + }, + _ => match lengths.get(&key).copied() { + Some(len) => { + let offset = (seed + i) % (len + 1); + let max_bytes = ((seed + i) % 5) + 1; + step_ops.push(Operation::Read(key, offset, max_bytes)); + } + None => { + lengths.insert(key.clone(), data.len()); + step_ops.push(Operation::Set(key, data)); + } + }, + } + } + + let initial_hash = normal_db + .hash() + .expect("hashing normal db should succeed"); + + let mut prove_db: TracedDatabase> = normal_db + .try_start_proof() + .expect("starting proof should succeed"); + + let prove_reads: Vec> = step_ops + .iter() + .filter_map(|op| apply_to_database(&mut prove_db, op)) + .collect(); + let prove_final_hash = prove_db + .hash() + .expect("hashing prove db should succeed"); + + let (prove_inner, _) = prove_db.into_parts(); + let mut verify_db: TracedDatabase = + TracedDatabase::from(to_verify::(&prove_inner)); + + let verify_initial_hash = verify_db + .hash() + .expect("hashing verify db should succeed"); + prop_assert_eq!( + verify_initial_hash, initial_hash, + "verify initial hash must match normal initial hash" + ); + + let step_ops_for_verify = step_ops.clone(); + let (verify_final_hash, verify_reads) = catch_not_found(move || { + let reads: Vec> = step_ops_for_verify + .iter() + .filter_map(|op| apply_to_database(&mut verify_db, op)) + .collect(); + let h = verify_db + .hash() + .expect("hashing verify db should succeed"); + (h, reads) + }) + .expect("verify replay must not trigger not_found"); + + prop_assert_eq!( + &prove_reads, &verify_reads, + "verify reads must match prove reads" + ); + prop_assert_eq!( + prove_final_hash, verify_final_hash, + "verify final hash must match prove final hash" + ); + + // advance Normal mode with the same operations so the next batch's initial + // state reflects the accumulated mutations. + for op in &step_ops { + apply_to_database(&mut normal_db, op); + } + let normal_final_hash = normal_db + .hash() + .expect("hashing normal db should succeed"); + + prop_assert_eq!( + normal_final_hash, prove_final_hash, + "normal final hash must match prove final hash" + ); + } + + normal_db.into_trace() + }); + #[cfg(feature = "rocksdb")] kv_test!( #[should_panic(expected = "trace mismatch")] diff --git a/durable-storage/src/lib.rs b/durable-storage/src/lib.rs index 2b66c5b42d..9e9c650156 100644 --- a/durable-storage/src/lib.rs +++ b/durable-storage/src/lib.rs @@ -36,6 +36,8 @@ pub mod key; mod merkle_layer; mod merkle_worker; pub mod persistence_layer; +#[cfg(test)] +mod proof_test_utils; pub mod registry; pub mod repo; pub mod storage; diff --git a/durable-storage/src/merkle_layer.rs b/durable-storage/src/merkle_layer.rs index 11651d9c36..6b9518f882 100644 --- a/durable-storage/src/merkle_layer.rs +++ b/durable-storage/src/merkle_layer.rs @@ -701,7 +701,6 @@ mod tests { use std::sync::Arc; use octez_riscv_data::components::bytes::Bytes; - use octez_riscv_data::components::bytes::BytesMode; use octez_riscv_data::merkle_proof::FromProof; use octez_riscv_data::merkle_proof::ProofError; use octez_riscv_data::merkle_proof::proof_tree::MerkleProof; @@ -714,7 +713,6 @@ mod tests { use proptest::prop_assert_eq; use super::MerkleLayer; - use super::MerkleLayerMode; use super::ProveImpl; use super::new_merkle_layer; use super::new_verify_layer; @@ -731,10 +729,13 @@ mod tests { use crate::avl::tree::Tree; use crate::key::Key; use crate::merkle_layer::VerifyImpl; + use crate::proof_test_utils::apply_to_merkle_layer; + use crate::proof_test_utils::setup_data_len; use crate::storage::KeyValueStore; use crate::storage::PersistentKeyValueStore; use crate::storage::TestKeyValueStoreSetup; use crate::storage::kv_test; + use crate::test_helpers::Operation; impl MerkleLayer { fn tree(&self) -> &Tree { @@ -793,60 +794,6 @@ mod tests { } } - /// Operations executed during a prove step and replayed during verification. - #[derive(Debug, Clone)] - enum Operation { - /// Set (insert or overwrite) a key. - Set(Key, Vec), - /// Delete a key (may be absent — that is a no-op on both sides). - Delete(Key), - /// In-place write at the given offset. - /// Only applied to keys that are known to exist. - /// The offset must be `<= existing_data_len` for the write to succeed. - Write(Key, usize, Vec), - /// Read up to `count` bytes starting at `offset`. A missing key reads zero bytes; reads - /// past the end clamp to the data length. Recording the read range is what makes the - /// data resolvable in Verify mode. - Read(Key, usize, usize), - } - - /// Apply `op` against `ml`. Returns `Some(bytes)` for [`Operation::Read`] so callers can - /// compare read results across modes, and `None` for state-mutating ops. - fn apply_operation( - ml: &mut MerkleLayer, - op: &Operation, - ) -> Option> { - match op { - Operation::Set(key, data) => { - ml.set(key, data).expect("set should succeed"); - None - } - Operation::Delete(key) => { - ml.delete(key).expect("delete should succeed"); - None - } - Operation::Write(key, offset, data) => { - ml.write(key, *offset, data).expect("write should succeed"); - None - } - Operation::Read(key, offset, count) => { - let mut buf = vec![0u8; *count]; - let n = ml - .get(key) - .expect("get should succeed") - .map_or(0, |data| data.read(*offset, &mut buf)); - buf.truncate(n); - Some(buf) - } - } - } - - /// Derive a data length in the range [1, 8] from the key bytes. Used to generate - /// variable-length data for testing. - fn setup_data_len(key_bytes: &[u8; 2]) -> usize { - (key_bytes[0] as usize % 8) + 1 - } - /// Core assertion: prove-mode proof and hash are consistent with Normal mode and the proof /// can be replayed under Verify mode to reach the same final hash. /// @@ -878,7 +825,7 @@ mod tests { let mut prove_ml = normal_ml.start_proof(); let prove_reads: Vec> = operations .iter() - .filter_map(|op| apply_operation(&mut prove_ml, op)) + .filter_map(|op| apply_to_merkle_layer(&mut prove_ml, op)) .collect(); let prove_final_hash = prove_ml.hash(); @@ -895,7 +842,7 @@ mod tests { // ---- Normal: replay same operations for reference hash and read values ---- let normal_reads: Vec> = operations .iter() - .filter_map(|op| apply_operation(&mut normal_ml, op)) + .filter_map(|op| apply_to_merkle_layer(&mut normal_ml, op)) .collect(); let normal_final_hash = normal_ml.hash(); @@ -926,7 +873,7 @@ mod tests { let (verify_final_hash, verify_reads) = catch_not_found(move || { let reads: Vec> = operations_for_verify .iter() - .filter_map(|op| apply_operation(&mut verify_ml, op)) + .filter_map(|op| apply_to_merkle_layer(&mut verify_ml, op)) .collect(); (verify_ml.hash(), reads) }) @@ -2311,7 +2258,7 @@ mod tests { let ops_count = seed % 20; for i in 0..ops_count { let key = keys[i % keys.len()].clone(); - let data = vec![seed as u8; 5]; + let data = bytes::Bytes::from(vec![seed as u8; 5]); match (i + seed) % 4 { 0 => { lengths.insert(key.clone(), data.len()); @@ -2364,7 +2311,7 @@ mod tests { let key = Key::new(bytes).expect("key should be valid"); // The offset must be valid for the existing data length. let offset = offsets[i % offsets.len()] % (setup_data_len(bytes) + 1); - Operation::Write(key, offset, vec![0xAA; 1 + (i % 8)]) + Operation::Write(key, offset, bytes::Bytes::from(vec![0xAA; 1 + (i % 8)])) }) .collect(); @@ -2402,7 +2349,7 @@ mod tests { .enumerate() .map(|(i, bytes)| { let key = Key::new(bytes).expect("key should be valid"); - Operation::Set(key, vec![0xBB; 1 + (i % 8)]) + Operation::Set(key, bytes::Bytes::from(vec![0xBB; 1 + (i % 8)])) }) .collect(); diff --git a/durable-storage/src/proof_test_utils.rs b/durable-storage/src/proof_test_utils.rs new file mode 100644 index 0000000000..52083d6c33 --- /dev/null +++ b/durable-storage/src/proof_test_utils.rs @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: 2026 Nomadic Labs +// +// SPDX-License-Identifier: MIT + +//! Shared test helpers for prove/verify proof round-trip tests. +//! +//! Hosts `apply_*` helpers that consume the [`Operation`] enum defined in [`test_helpers`]. +//! Lives behind `#[cfg(test)]` (rather than the `unstable-test-utils` feature like +//! [`test_helpers`]) because it depends on the test-only [`TracedDatabase`] wrapper. +//! +//! [`test_helpers`]: crate::test_helpers +//! [`Operation`]: crate::test_helpers::Operation +//! [`TracedDatabase`]: crate::database::traced_database::TracedDatabase + +use octez_riscv_data::components::bytes::BytesMode; + +use crate::database::DatabaseMode; +use crate::database::traced_database::TracedDatabase; +use crate::merkle_layer::MerkleLayer; +use crate::merkle_layer::MerkleLayerMode; +use crate::merkle_worker::BackgroundKeyValueStore; +use crate::storage::KeyValueStore; +use crate::test_helpers::Operation; + +/// Derive a setup-time data length in the range `[1, 8]` from key bytes. Used to seed setup +/// values with varying length so proof tests cover non-uniform initial state. +pub(crate) fn setup_data_len(key_bytes: &[u8; 2]) -> usize { + (key_bytes[0] as usize % 8) + 1 +} + +/// Apply `op` against a [`MerkleLayer`] in any mode that implements [`MerkleLayerMode`] and +/// [`BytesMode`]. Returns `Some(bytes)` for [`Operation::Read`] and `None` for mutating ops. +/// +/// Only the prove/verify-relevant subset of [`Operation`] is supported — variants that target +/// registry- or commit-level state are unreachable in a single proof step. Read clamps to the +/// existing data length and returns zero bytes for a missing key. +pub(crate) fn apply_to_merkle_layer( + ml: &mut MerkleLayer, + op: &Operation, +) -> Option> +where + KV: KeyValueStore, + M: MerkleLayerMode + BytesMode, +{ + match op { + Operation::Set(key, data) => { + ml.set(key, data).expect("set should succeed"); + None + } + Operation::Delete(key) => { + ml.delete(key).expect("delete should succeed"); + None + } + Operation::Write(key, offset, data) => { + ml.write(key, *offset, data).expect("write should succeed"); + None + } + Operation::Read(key, offset, count) => { + let mut buf = vec![0u8; *count]; + let n = ml + .get(key) + .expect("get should succeed") + .map_or(0, |data| data.read(*offset, &mut buf)); + buf.truncate(n); + Some(buf) + } + Operation::Exists(_) + | Operation::ValueLength(_) + | Operation::Hash + | Operation::Commit + | Operation::Checkout + | Operation::GrowRegistry + | Operation::ShrinkRegistry + | Operation::CopyDatabase + | Operation::MoveDatabase + | Operation::ClearDatabase => { + unimplemented!("{op:?} is not supported when applying to a MerkleLayer") + } + } +} + +/// Apply `op` against a [`TracedDatabase`] in any mode that implements [`DatabaseMode`]. +/// Returns `Some(bytes)` for [`Operation::Read`] and `None` for mutating ops. +/// +/// Only the prove/verify-relevant subset of [`Operation`] is supported. +/// `Database::read_bytes` requires the key to exist and the offset to be `<= value_length`; +/// callers must constrain `Read` ops accordingly. +pub(crate) fn apply_to_database( + db: &mut TracedDatabase, + op: &Operation, +) -> Option> +where + KV: BackgroundKeyValueStore, + M: DatabaseMode, +{ + match op { + Operation::Set(key, data) => { + db.set(key.clone(), data.clone()) + .expect("set should succeed"); + None + } + Operation::Delete(key) => { + db.delete(key.clone()).expect("delete should succeed"); + None + } + Operation::Write(key, offset, data) => { + db.write(key.clone(), *offset, data.clone()) + .expect("write should succeed"); + None + } + Operation::Read(key, offset, max_bytes) => Some( + db.read_bytes(key, *offset, *max_bytes) + .expect("read_bytes should succeed"), + ), + Operation::Exists(_) + | Operation::ValueLength(_) + | Operation::Hash + | Operation::Commit + | Operation::Checkout + | Operation::GrowRegistry + | Operation::ShrinkRegistry + | Operation::CopyDatabase + | Operation::MoveDatabase + | Operation::ClearDatabase => { + unimplemented!("{op:?} is not supported when applying to a TracedDatabase") + } + } +}