Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
167 changes: 166 additions & 1 deletion durable-storage/src/database.rs
Original file line number Diff line number Diff line change
Expand Up @@ -584,7 +584,7 @@ impl<KV: BackgroundKeyValueStore, M: DatabaseMode> Database<KV, M> {
}
}

mod traced_database;
pub(crate) mod traced_database;

#[cfg(test)]
pub(crate) mod tests {
Expand Down Expand Up @@ -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<KV: BackgroundKeyValueStore>(
handle: &Handle,
Expand Down Expand Up @@ -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<Verify>` 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::<KV>(handle, repo);

let keys: Vec<Key> = 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<Key, usize> =
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<Operation> = 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<KV, Prove<'static>> = normal_db
.try_start_proof()
.expect("starting proof should succeed");

let prove_reads: Vec<Vec<u8>> = 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<KV, Verify> =
TracedDatabase::from(to_verify::<KV>(&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<Vec<u8>> = 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")]
Expand Down
2 changes: 2 additions & 0 deletions durable-storage/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
71 changes: 9 additions & 62 deletions durable-storage/src/merkle_layer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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<KV: KeyValueStore> MerkleLayer<KV, Normal> {
fn tree(&self) -> &Tree<LazyNodeId> {
Expand Down Expand Up @@ -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<u8>),
/// 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<u8>),
/// 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<KV: KeyValueStore, M: MerkleLayerMode + BytesMode>(
ml: &mut MerkleLayer<KV, M>,
op: &Operation,
) -> Option<Vec<u8>> {
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.
///
Expand Down Expand Up @@ -878,7 +825,7 @@ mod tests {
let mut prove_ml = normal_ml.start_proof();
let prove_reads: Vec<Vec<u8>> = 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();

Expand All @@ -895,7 +842,7 @@ mod tests {
// ---- Normal: replay same operations for reference hash and read values ----
let normal_reads: Vec<Vec<u8>> = 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();

Expand Down Expand Up @@ -926,7 +873,7 @@ mod tests {
let (verify_final_hash, verify_reads) = catch_not_found(move || {
let reads: Vec<Vec<u8>> = 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)
})
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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();

Expand Down
Loading
Loading