Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix backward iterator seek #215

Merged
merged 6 commits into from
May 14, 2024
Merged
Show file tree
Hide file tree
Changes from 5 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
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
}
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);
}
}
Loading