Skip to content

Commit

Permalink
Fix backward iterator seek (#215)
Browse files Browse the repository at this point in the history
Fix backward iterator seek
  • Loading branch information
cheme authored May 14, 2024
1 parent 4c5bf65 commit 47f0dfe
Show file tree
Hide file tree
Showing 9 changed files with 207 additions and 56 deletions.
132 changes: 132 additions & 0 deletions test-support/reference-trie/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1235,3 +1235,135 @@ mod tests {
}
}
}

// This is a bit redundant with other iterator fuzzer
fn test_iterator<L, DB>(entries: Vec<(Vec<u8>, Vec<u8>)>, keys: Vec<Vec<u8>>, prefix: bool)
where
L: TrieLayout,
DB: hash_db::HashDB<L::Hash, DBValue> + hash_db::HashDBRef<L::Hash, DBValue> + Default,
{
let mut ref_tree = std::collections::BTreeMap::new();

// Populate DB with full trie from entries.
let (db, root) = {
let mut db = DB::default();
let mut root = Default::default();
{
let mut trie = TrieDBMutBuilder::<L>::new(&mut db, &mut root).build();
for (key, value) in entries.into_iter() {
trie.insert(&key, &value).unwrap();
ref_tree.insert(key, value);
}
}
(db, root)
};
// Lookup items in trie while recording traversed nodes.
let trie = TrieDBBuilder::<L>::new(&db, &root).build();
// standard iter
let mut iter = trie_db::triedb::TrieDBDoubleEndedIterator::new(&trie).unwrap();
let mut iter_ref = ref_tree.iter();
let mut iter_ref2 = ref_tree.iter();

loop {
let n = iter.next();
let nb = iter.next_back();
let n_ref = iter_ref.next();
let nb_ref = iter_ref2.next_back();
assert_eq!(n.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), n_ref);
assert_eq!(nb.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), nb_ref);
if n_ref.is_none() && nb_ref.is_none() {
break;
}
}
for key in keys {
use trie_db::TrieIterator;
let mut iter = if prefix {
trie_db::triedb::TrieDBDoubleEndedIterator::new_prefixed(&trie, &key).unwrap()
} else {
trie_db::triedb::TrieDBDoubleEndedIterator::new(&trie).unwrap()
};
if !prefix {
iter.seek(&key).unwrap();
}
let mut iter_ref =
ref_tree
.iter()
.filter(|k| if prefix { k.0.starts_with(&key) } else { k.0 >= &key });
let mut iter_ref2 =
ref_tree
.iter()
.filter(|k| if prefix { k.0.starts_with(&key) } else { k.0 <= &key });
loop {
let n = iter.next();
let nb = iter.next_back();
let n_ref = iter_ref.next();
let nb_ref = iter_ref2.next_back();
assert_eq!(n.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), n_ref);
assert_eq!(nb.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), nb_ref);
if n_ref.is_none() && nb_ref.is_none() {
break;
}
}
}
}

pub fn fuzz_double_iter<T, DB>(input: &[u8], prefix: bool)
where
T: TrieLayout,
DB: hash_db::HashDB<T::Hash, DBValue> + hash_db::HashDBRef<T::Hash, DBValue> + Default,
{
let mut data = fuzz_to_data(input);
// - the first 2/3 are added to the trie.
// - the last 1/3 is not added to the trie and use for random seek and prefix.
let mut keys = data[(data.len() / 3)..].iter().map(|(key, _)| key.clone()).collect::<Vec<_>>();
data.truncate(data.len() * 2 / 3);

let data = data_sorted_unique(data);
keys.sort();
keys.dedup();

test_iterator::<T, DB>(data, keys, prefix);
}

pub fn fuzz_to_data(input: &[u8]) -> Vec<(Vec<u8>, Vec<u8>)> {
let mut result = Vec::new();
// enc = (minkeylen, maxkeylen (min max up to 32), datas)
// fix data len 2 bytes
let mut minkeylen = if let Some(v) = input.get(0) {
let mut v = *v & 31u8;
v = v + 1;
v
} else {
return result
};
let mut maxkeylen = if let Some(v) = input.get(1) {
let mut v = *v & 31u8;
v = v + 1;
v
} else {
return result
};

if maxkeylen < minkeylen {
let v = minkeylen;
minkeylen = maxkeylen;
maxkeylen = v;
}
let mut ix = 2;
loop {
let keylen = if let Some(v) = input.get(ix) {
let mut v = *v & 31u8;
v = v + 1;
v = std::cmp::max(minkeylen, v);
v = std::cmp::min(maxkeylen, v);
v as usize
} else {
break
};
let key = if input.len() > ix + keylen { input[ix..ix + keylen].to_vec() } else { break };
ix += keylen;
let val = if input.len() > ix + 2 { input[ix..ix + 2].to_vec() } else { break };
result.push((key, val));
}
result
}
3 changes: 3 additions & 0 deletions trie-db/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@ The format is based on [Keep a Changelog].

[Keep a Changelog]: http://keepachangelog.com/en/1.0.0/

## [0.29.1] - 2024-05-24
- Fix backward iterator seek [#215](https://github.com/paritytech/trie/pull/215)

## [0.29.0] - 2024-03-04
- Implements `DoubleEndedIterator` for trie iterator [#208](https://github.com/paritytech/trie/pull/208)

Expand Down
2 changes: 1 addition & 1 deletion trie-db/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "trie-db"
version = "0.29.0"
version = "0.29.1"
authors = ["Parity Technologies <admin@parity.io>"]
description = "Merkle-Patricia Trie generic over key hasher and node encoding"
repository = "https://github.com/paritytech/trie"
Expand Down
6 changes: 5 additions & 1 deletion trie-db/fuzz/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ cargo-fuzz = true
[dependencies]
hash-db = { path = "../../hash-db", version = "0.16.0" }
memory-db = { path = "../../memory-db", version = "0.32.0" }
reference-trie = { path = "../../test-support/reference-trie", version = "0.29.0" }
reference-trie = { path = "../../test-support/reference-trie", version = "0.29.1" }
arbitrary = { version = "1.3.0", features = ["derive"] }
array-bytes = "6.0.0"

Expand Down Expand Up @@ -68,3 +68,7 @@ path = "fuzz_targets/trie_proof_invalid.rs"
[[bin]]
name = "prefix_seek_iter"
path = "fuzz_targets/prefix_seek_iter.rs"

[[bin]]
name = "double_iter"
path = "fuzz_targets/double_iter.rs"
14 changes: 14 additions & 0 deletions trie-db/fuzz/fuzz_targets/double_iter.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
#![no_main]

use libfuzzer_sys::fuzz_target;
use memory_db::{MemoryDB, PrefixedKey};
use trie_db::{DBValue, TrieLayout};
use trie_db_fuzz::fuzz_double_iter;

type T = reference_trie::NoExtensionLayout;
type DB = MemoryDB<<T as TrieLayout>::Hash, PrefixedKey<<T as TrieLayout>::Hash>, DBValue>;

fuzz_target!(|data: &[u8]| {
fuzz_double_iter::<T, DB>(data, false);
fuzz_double_iter::<T, DB>(data, true);
});
47 changes: 3 additions & 44 deletions trie-db/fuzz/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,58 +15,17 @@
use arbitrary::Arbitrary;
use hash_db::Hasher;
use memory_db::{HashKey, MemoryDB, PrefixedKey};
pub use reference_trie::fuzz_double_iter;
use reference_trie::{
calc_root, compare_insert_remove, reference_trie_root_iter_build as reference_trie_root,
calc_root, compare_insert_remove, fuzz_to_data,
reference_trie_root_iter_build as reference_trie_root,
};
use std::convert::TryInto;
use trie_db::{
proof::{generate_proof, verify_proof},
DBValue, Trie, TrieDBBuilder, TrieDBIterator, TrieDBMutBuilder, TrieLayout, TrieMut,
};

fn fuzz_to_data(input: &[u8]) -> Vec<(Vec<u8>, Vec<u8>)> {
let mut result = Vec::new();
// enc = (minkeylen, maxkeylen (min max up to 32), datas)
// fix data len 2 bytes
let mut minkeylen = if let Some(v) = input.get(0) {
let mut v = *v & 31u8;
v = v + 1;
v
} else {
return result
};
let mut maxkeylen = if let Some(v) = input.get(1) {
let mut v = *v & 31u8;
v = v + 1;
v
} else {
return result
};

if maxkeylen < minkeylen {
let v = minkeylen;
minkeylen = maxkeylen;
maxkeylen = v;
}
let mut ix = 2;
loop {
let keylen = if let Some(v) = input.get(ix) {
let mut v = *v & 31u8;
v = v + 1;
v = std::cmp::max(minkeylen, v);
v = std::cmp::min(maxkeylen, v);
v as usize
} else {
break
};
let key = if input.len() > ix + keylen { input[ix..ix + keylen].to_vec() } else { break };
ix += keylen;
let val = if input.len() > ix + 2 { input[ix..ix + 2].to_vec() } else { break };
result.push((key, val));
}
result
}

fn fuzz_removal(data: Vec<(Vec<u8>, Vec<u8>)>) -> Vec<(bool, Vec<u8>, Vec<u8>)> {
let mut res = Vec::new();
let mut torem = None;
Expand Down
9 changes: 3 additions & 6 deletions trie-db/src/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ impl<L: TrieLayout> TrieDBRawIterator<L> {
NodePlan::Leaf { partial: partial_plan, .. } => {
let slice = partial_plan.build(node_data);
if (fwd && slice < partial) || (!fwd && slice > partial) {
crumb.status = Status::Exiting;
crumb.status = Status::AftExiting;
return Ok(false);
}
return Ok(slice.starts_with(&partial));
Expand All @@ -182,8 +182,7 @@ impl<L: TrieLayout> TrieDBRawIterator<L> {
let slice = partial_plan.build(node_data);
if !partial.starts_with(&slice) {
if (fwd && slice < partial) || (!fwd && slice > partial) {
crumb.status = Status::Exiting;
self.key_nibbles.append_partial(slice.right());
crumb.status = Status::AftExiting;
return Ok(false);
}
return Ok(slice.starts_with(&partial));
Expand Down Expand Up @@ -230,9 +229,7 @@ impl<L: TrieLayout> TrieDBRawIterator<L> {
let slice = partial_plan.build(node_data);
if !partial.starts_with(&slice) {
if (fwd && slice < partial) || (!fwd && slice > partial) {
crumb.status = Status::Exiting;
self.key_nibbles.append_partial(slice.right());
self.key_nibbles.push((nibble_ops::NIBBLE_LENGTH - 1) as u8);
crumb.status = Status::AftExiting;
return Ok(false);
}
return Ok(slice.starts_with(&partial));
Expand Down
12 changes: 12 additions & 0 deletions trie-db/src/triedb.rs
Original file line number Diff line number Diff line change
Expand Up @@ -496,6 +496,18 @@ impl<'a, 'cache, L: TrieLayout> TrieDBDoubleEndedIterator<'a, 'cache, L> {
back_raw_iter: TrieDBRawIterator::new(db)?,
})
}

/// Create a new iterator, but limited to a given prefix.
pub fn new_prefixed(
db: &'a TrieDB<'a, 'cache, L>,
prefix: &[u8],
) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self {
db,
raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
back_raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
})
}
}

impl<L: TrieLayout> TrieDoubleEndedIterator<L> for TrieDBDoubleEndedIterator<'_, '_, L> {}
Expand Down
38 changes: 34 additions & 4 deletions trie-db/test/src/double_ended_iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -201,10 +201,7 @@ fn seek_back_works_internal<T: TrieLayout>() {
let mut iter = TrieDBNodeDoubleEndedIterator::new(&trie).unwrap();

<dyn TrieDoubleEndedIterator<T, Item = _>>::seek(&mut iter, &hex!("")[..]).unwrap();
match iter.next_back() {
Some(Ok((prefix, _, _))) => assert_eq!(prefix, nibble_vec(hex!(""), 0)),
_ => panic!("unexpected item"),
}
assert!(iter.next_back().is_none());

<dyn TrieDoubleEndedIterator<T, Item = _>>::seek(&mut iter, &hex!("03")[..]).unwrap();
match iter.next_back() {
Expand Down Expand Up @@ -406,3 +403,36 @@ fn prefix_over_empty_works_internal<T: TrieLayout>() {
iter.prefix(&hex!("00")[..]).unwrap();
assert!(iter.next_back().is_none());
}

test_layouts!(next_back_weird_behaviour_1, next_back_weird_behaviour_internal_1);
fn next_back_weird_behaviour_internal_1<T: TrieLayout>() {
use trie_db::TrieIterator;

let pairs = vec![(vec![11], b"bbbb".to_vec())];

let (memdb, root) = build_trie_db::<T>(&pairs);
let trie = TrieDBBuilder::<T>::new(&memdb, &root).build();

let mut iter = trie_db::triedb::TrieDBDoubleEndedIterator::new(&trie).unwrap();

iter.seek(&[10]).unwrap();
assert!(iter.next_back().is_none());
}

test_layouts!(fuzz_set, fuzz_set_internal);
fn fuzz_set_internal<T: TrieLayout>() {
type DB<T> = memory_db::MemoryDB<
<T as TrieLayout>::Hash,
memory_db::PrefixedKey<<T as TrieLayout>::Hash>,
trie_db::DBValue,
>;

let fuzz_inputs = [
vec![32, 65, 255, 254, 255, 213, 0, 0, 0, 254, 255, 0, 0, 235, 0, 0, 35],
vec![0, 5, 0, 0, 43, 0, 5, 0],
];
for i in fuzz_inputs {
reference_trie::fuzz_double_iter::<T, DB<T>>(&i, false);
reference_trie::fuzz_double_iter::<T, DB<T>>(&i, true);
}
}

0 comments on commit 47f0dfe

Please sign in to comment.