Skip to content

Commit ab87372

Browse files
committed
wip
1 parent f6c578e commit ab87372

File tree

3 files changed

+442
-15
lines changed

3 files changed

+442
-15
lines changed

firewood/Cargo.toml

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ fastrace.workspace = true
2929
firewood-macros.workspace = true
3030
hex.workspace = true
3131
metrics.workspace = true
32+
smallvec = { workspace = true, features = ["write", "union"] }
3233
sha2.workspace = true
3334
test-case.workspace = true
3435
thiserror.workspace = true
@@ -44,8 +45,8 @@ default = []
4445
nightly = []
4546
io-uring = ["firewood-storage/io-uring"]
4647
logger = ["firewood-storage/logger"]
47-
branch_factor_256 = [ "firewood-storage/branch_factor_256" ]
48-
ethhash = [ "firewood-storage/ethhash" ]
48+
branch_factor_256 = ["firewood-storage/branch_factor_256"]
49+
ethhash = ["firewood-storage/ethhash"]
4950

5051
[dev-dependencies]
5152
# Workspace dependencies

firewood/src/merkle.rs

Lines changed: 282 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,18 @@
44
#[cfg(test)]
55
mod tests;
66

7-
use crate::proof::{Proof, ProofCollection, ProofError, ProofNode};
7+
use crate::proof::{Proof, ProofCollection, ProofError, ProofNode, verify_opt_value_digest};
88
use crate::range_proof::RangeProof;
99
use crate::stream::{MerkleKeyValueStream, PathIterator};
1010
use crate::v2::api::{self, FrozenProof, FrozenRangeProof, KeyType, ValueType};
1111
use firewood_storage::{
12-
BranchNode, Child, FileIoError, HashType, HashedNodeReader, ImmutableProposal, IntoHashType,
13-
LeafNode, MaybePersistedNode, MutableProposal, NibblesIterator, Node, NodeStore, Parentable,
14-
Path, ReadableStorage, SharedNode, TrieHash, TrieReader, ValueDigest,
12+
BranchNode, Child, FileIoError, HashType, Hashable, HashedNodeReader, ImmutableProposal,
13+
IntoHashType, LeafNode, MaybePersistedNode, MemStore, MutableProposal, NibblesIterator, Node,
14+
NodeStore, Parentable, Path, ReadableStorage, SharedNode, TrieHash, TrieReader, ValueDigest,
1515
};
1616
use futures::{StreamExt, TryStreamExt};
1717
use metrics::counter;
18+
use smallvec::SmallVec;
1819
use std::collections::HashSet;
1920
use std::fmt::{Debug, Write};
2021
use std::future::ready;
@@ -153,7 +154,6 @@ impl<T: TrieReader> Merkle<T> {
153154
self.nodestore.root_node()
154155
}
155156

156-
#[cfg(test)]
157157
pub(crate) const fn nodestore(&self) -> &T {
158158
&self.nodestore
159159
}
@@ -180,7 +180,6 @@ impl<T: TrieReader> Merkle<T> {
180180
if proof.is_empty() {
181181
// No nodes, even the root, are before `key`.
182182
// The root alone proves the non-existence of `key`.
183-
// TODO reduce duplicate code with ProofNode::from<PathIterItem>
184183
let child_hashes = if let Some(branch) = root.as_branch() {
185184
branch.children_hashes()
186185
} else {
@@ -277,12 +276,284 @@ impl<T: TrieReader> Merkle<T> {
277276
/// incremental range proof verification
278277
pub fn verify_range_proof(
279278
&self,
280-
_first_key: Option<impl KeyType>,
281-
_last_key: Option<impl KeyType>,
282-
_root_hash: &TrieHash,
283-
_proof: &RangeProof<impl KeyType, impl ValueType, impl ProofCollection>,
279+
first_key: Option<impl KeyType>,
280+
last_key: Option<impl KeyType>,
281+
root_hash: &TrieHash,
282+
proof: &RangeProof<impl KeyType, impl ValueType, impl ProofCollection>,
284283
) -> Result<(), api::Error> {
285-
todo!()
284+
let first_key = first_key.map(Path::from).unwrap_or_default();
285+
let last_key = last_key.map(Path::from).unwrap_or_default();
286+
287+
// 1. Validate proof structure (similar to validateChangeProof in Go)
288+
self.validate_range_proof_structure(&first_key, &last_key, proof)?;
289+
290+
// 2. Verify start proof nodes.
291+
self.verify_proof_nodes_values(
292+
proof.start_proof(),
293+
&first_key,
294+
&last_key,
295+
// validate_range_proof_structure will have verified that the keys are
296+
// in order, allowing us to use binary search on lookup
297+
proof.key_values(),
298+
)?;
299+
300+
// 3. Verify end proof nodes
301+
self.verify_proof_nodes_values(
302+
proof.end_proof(),
303+
&first_key,
304+
&last_key,
305+
// validate_range_proof_structure will have verified that the keys are
306+
// in order, allowing us to use binary search on lookup
307+
proof.key_values(),
308+
)?;
309+
310+
// 4. Reconstruct trie and verify root
311+
self.verify_reconstructed_trie_root(proof, root_hash)?;
312+
313+
Ok(())
314+
}
315+
316+
/// Verify that the range proof is structurally valid and that we can use it
317+
/// to verify the trie root once reconstructed.
318+
fn validate_range_proof_structure(
319+
&self,
320+
first_key: &Path,
321+
last_key: &Path,
322+
proof: &RangeProof<impl KeyType, impl ValueType, impl ProofCollection>,
323+
) -> Result<(), ProofError> {
324+
// 1. Basic validation
325+
if proof.is_empty() {
326+
return Err(ProofError::Empty);
327+
}
328+
329+
// 2. Range validation
330+
if !first_key.is_empty() && !last_key.is_empty() && first_key > last_key {
331+
return Err(ProofError::InvalidRange);
332+
}
333+
334+
// 3. Proof structure validation
335+
match (
336+
first_key.is_empty(),
337+
last_key.is_empty(),
338+
proof.key_values().is_empty(),
339+
) {
340+
(_, true, true) if !proof.end_proof().is_empty() => {
341+
return Err(ProofError::UnexpectedEndProof);
342+
}
343+
(true, _, _) if !proof.start_proof().is_empty() => {
344+
return Err(ProofError::UnexpectedStartProof);
345+
}
346+
(_, false, _) | (_, _, false) if proof.end_proof().is_empty() => {
347+
return Err(ProofError::ExpectedEndProof);
348+
}
349+
_ => {} // Valid combination
350+
}
351+
352+
let last_key = if proof.key_values().is_empty() {
353+
// no key-values, `last_key` remains the `last_key` provided by the caller
354+
last_key
355+
} else {
356+
// two re-usable buffers to expand the keys into nibbles. re-using
357+
// these buffers avoids multiple allocations and deallocations in the
358+
// loop below.
359+
//
360+
// Uses a smallvec so we can convert directly to a Path without
361+
// an extra copy step.
362+
let mut this_key_buf = SmallVec::<[u8; 64]>::new();
363+
let mut last_key_buf = SmallVec::<[u8; 64]>::new();
364+
365+
// 4. verify key-values are in strict order by key (no duplicates either)
366+
for (key, _) in proof.key_values() {
367+
this_key_buf.clear();
368+
this_key_buf.extend(NibblesIterator::new(key.as_ref()));
369+
debug_assert!(!this_key_buf.is_empty(), "key must not be empty");
370+
debug_assert!(this_key_buf.len() % 2 == 0, "key must be even length");
371+
372+
// verify that the first key is not larger than the first key
373+
// in the range. Only check the first key as all remaining keys
374+
// are implicitly larger if other checks hold.
375+
if last_key_buf.is_empty() && !first_key.is_empty() && *this_key_buf < **first_key {
376+
return Err(ProofError::StateFromOutsideOfRange);
377+
}
378+
379+
// For every key, check that it is less than or equal to the last
380+
// key in the range.
381+
if !last_key.is_empty() && *this_key_buf <= **last_key {
382+
return Err(ProofError::StateFromOutsideOfRange);
383+
}
384+
385+
if !last_key_buf.is_empty() && this_key_buf < last_key_buf {
386+
// we have a last key but it is greater than the current key
387+
// therefore, the list is not sorted or has duplicates
388+
return Err(ProofError::NonIncreasingValues);
389+
}
390+
391+
// swap the buffers so that `last_key_buf` contains the key we
392+
// processed in this iteration.
393+
std::mem::swap(&mut last_key_buf, &mut this_key_buf);
394+
}
395+
396+
// at this point, `last_key_buf` is filled with the nibbles of the last
397+
// and largest key in the key-values. We can re-use it for the
398+
// verification below. It overrides the `last_key` provided by the
399+
// caller in order to verify the end proof.
400+
401+
&Path(last_key_buf)
402+
};
403+
404+
// 5. Validate proof paths (structural only, not root verification)
405+
if !proof.start_proof().is_empty() {
406+
proof.start_proof().verify_proof_path_structure(first_key)?;
407+
}
408+
409+
if !proof.end_proof().is_empty() {
410+
proof.end_proof().verify_proof_path_structure(last_key)?;
411+
}
412+
413+
Ok(())
414+
}
415+
416+
fn verify_proof_nodes_values(
417+
&self,
418+
proof: &Proof<impl ProofCollection>,
419+
first_key: &Path,
420+
last_key: &Path,
421+
key_values_sorted_by_key: &[(impl KeyType, impl ValueType)],
422+
) -> Result<(), ProofError> {
423+
// cache the root node to avoid multiple lookups
424+
let root = self.root();
425+
let root = root.as_deref();
426+
427+
let mut node_key = Path::default();
428+
for node in proof.as_ref() {
429+
node_key.0.clear();
430+
node_key.0.extend(node.key());
431+
432+
// skip partial paths as they cannot have values
433+
#[cfg(not(feature = "branch_factor_256"))]
434+
if node_key.len() % 2 != 0 {
435+
continue;
436+
}
437+
438+
if !first_key.is_empty() && *node_key < **first_key
439+
|| !last_key.is_empty() && *node_key > **last_key
440+
{
441+
// node not in range, ignore it
442+
continue;
443+
}
444+
445+
verify_opt_value_digest(
446+
// must be inline with the call to verify_opt_value_digest
447+
// in order for the lifetime of the temp storage to be valid,
448+
// which ends before the semicolon
449+
self.get_node_value_for_proof(
450+
key_values_sorted_by_key,
451+
&node_key,
452+
root,
453+
// temp storage for the fetched node so we don't need to
454+
// copy the value off the node
455+
&mut None,
456+
)?,
457+
node.value_digest(),
458+
)?; // temp storage is dropped here
459+
}
460+
461+
Ok(())
462+
}
463+
464+
fn verify_reconstructed_trie_root(
465+
&self,
466+
proof: &RangeProof<impl KeyType, impl ValueType, impl ProofCollection>,
467+
root_hash: &TrieHash,
468+
) -> Result<(), api::Error> {
469+
// Create in-memory trie for reconstruction
470+
let memstore = MemStore::new(vec![]);
471+
let nodestore = NodeStore::new_empty_proposal(memstore.into());
472+
let mut merkle = Merkle { nodestore };
473+
474+
// Insert all key-value pairs from the range proof
475+
for (key, value) in proof.key_values() {
476+
merkle
477+
.insert(key.as_ref(), value.as_ref().into())
478+
.map_err(ProofError::IO)?;
479+
}
480+
481+
// Hash the trie and get root
482+
let merkle: Merkle<NodeStore<Arc<ImmutableProposal>, _>> = merkle.try_into()?;
483+
let computed_root = merkle.nodestore().root_hash().ok_or(ProofError::Empty)?;
484+
485+
// Compare with expected root
486+
if computed_root == *root_hash {
487+
Ok(())
488+
} else {
489+
Err(api::Error::IncorrectRootHash {
490+
provided: root_hash.clone(),
491+
current: computed_root,
492+
})
493+
}
494+
}
495+
496+
/// Get the value for the given key from the proof key-values. If not found,
497+
/// it will look for the value in the trie.
498+
///
499+
/// This lifetime means `'out` is dependent on `'this`, `'kvs`, and `'store`.
500+
///
501+
/// This allows us to return a reference to the value in the node without
502+
/// copying it to a new buffer.
503+
fn get_node_value_for_proof<'this: 'out, 'kvs: 'out, 'store: 'out, 'out>(
504+
&'this self,
505+
key_values_sorted_by_key: &'kvs [(impl KeyType, impl ValueType)],
506+
path: &Path,
507+
root: Option<&Node>,
508+
fetched_node: &'store mut Option<SharedNode>,
509+
) -> Result<Option<&'out [u8]>, FileIoError> {
510+
use std::cmp::Ordering::{self, Equal, Greater, Less};
511+
512+
fn cmp_path_with_key(path: &Path, key: &[u8]) -> Ordering {
513+
let mut path = path.iter();
514+
let mut key = NibblesIterator::new(key);
515+
loop {
516+
break match (path.next(), key.next()) {
517+
(Some(path), Some(key)) => match path.cmp(&key) {
518+
Equal => continue, // all other branches break
519+
ord => ord,
520+
},
521+
(Some(_), None) => Greater, // path is longer than key
522+
(None, Some(_)) => Less, // key is longer than path
523+
(None, None) => Equal, // both are empty
524+
};
525+
}
526+
}
527+
528+
// use binary search to find the value for the key in `key_values_sorted_by_key`
529+
// (if it exists)
530+
if let Ok(found) = key_values_sorted_by_key
531+
.binary_search_by(|(key, _)| cmp_path_with_key(path, key.as_ref()))
532+
.map(
533+
#[expect(
534+
clippy::indexing_slicing,
535+
reason = "binary_search guarantees the index is in bounds"
536+
)]
537+
|index| key_values_sorted_by_key[index].1.as_ref(),
538+
)
539+
{
540+
return Ok(Some(found));
541+
}
542+
543+
// otherwise, look for it in the trie
544+
let Some(root) = root else {
545+
// no root, so no value
546+
return Ok(None);
547+
};
548+
549+
let Some(node) = get_helper(&self.nodestore, root, path.as_ref())? else {
550+
// node was not found, so no value
551+
return Ok(None);
552+
};
553+
554+
let node = &*fetched_node.insert(node);
555+
556+
Ok(node.value())
286557
}
287558

288559
pub(crate) fn path_iter<'a>(

0 commit comments

Comments
 (0)