Skip to content

Commit

Permalink
fix(trie): removing a blinded leaf should result in an error
Browse files Browse the repository at this point in the history
  • Loading branch information
@shekhirin
shekhirin committed Oct 18, 2024
1 parent 5859f93 commit 21f303f
Showing 1 changed file with 52 additions and 16 deletions.
68 changes: 52 additions & 16 deletions crates/trie/sparse/src/trie.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,11 +267,11 @@ impl RevealedSparseTrie {
/// Remove leaf node from the trie.
pub fn remove_leaf(&mut self, path: Nibbles) -> SparseTrieResult<()> {
self.prefix_set.insert(path.clone());
let existing = self.values.remove(&path);
if existing.is_none() {
// trie structure unchanged, return immediately
return Ok(())
}
self.values.remove(&path);

// If the path wasn't present in `values`, we still need to walk the trie and ensure that
// there is no node at the path. When a leaf node is a blinded `Hash`, it will have an entry
// in `nodes`, but not in the `values`.

let mut removed_nodes = self.take_nodes_for_path(&path)?;
debug!(target: "trie::sparse", ?path, ?removed_nodes, "Removed nodes for path");
Expand Down Expand Up @@ -726,6 +726,7 @@ mod tests {

use super::*;
use alloy_primitives::U256;
use assert_matches::assert_matches;
use itertools::Itertools;
use proptest::prelude::*;
use rand::seq::IteratorRandom;
Expand Down Expand Up @@ -959,7 +960,7 @@ mod tests {
pretty_assertions::assert_eq!(
sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
BTreeMap::from_iter([
(Nibbles::new(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1101.into())),
(
Nibbles::from_nibbles([0x5, 0x0]),
Expand All @@ -971,11 +972,11 @@ mod tests {
),
(
Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
SparseNode::new_leaf(Nibbles::new())
SparseNode::new_leaf(Nibbles::default())
),
(
Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
SparseNode::new_leaf(Nibbles::new())
SparseNode::new_leaf(Nibbles::default())
),
(
Nibbles::from_nibbles([0x5, 0x2]),
Expand Down Expand Up @@ -1014,7 +1015,7 @@ mod tests {
pretty_assertions::assert_eq!(
sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
BTreeMap::from_iter([
(Nibbles::new(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
(
Nibbles::from_nibbles([0x5, 0x0]),
Expand All @@ -1026,11 +1027,11 @@ mod tests {
),
(
Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
SparseNode::new_leaf(Nibbles::new())
SparseNode::new_leaf(Nibbles::default())
),
(
Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
SparseNode::new_leaf(Nibbles::new())
SparseNode::new_leaf(Nibbles::default())
),
(Nibbles::from_nibbles([0x5, 0x3]), SparseNode::new_branch(0b1010.into())),
(
Expand Down Expand Up @@ -1062,7 +1063,7 @@ mod tests {
pretty_assertions::assert_eq!(
sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
BTreeMap::from_iter([
(Nibbles::new(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
(
Nibbles::from_nibbles([0x5, 0x0]),
Expand Down Expand Up @@ -1096,7 +1097,7 @@ mod tests {
pretty_assertions::assert_eq!(
sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
BTreeMap::from_iter([
(Nibbles::new(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
(
Nibbles::from_nibbles([0x5, 0x0]),
Expand Down Expand Up @@ -1127,7 +1128,7 @@ mod tests {
pretty_assertions::assert_eq!(
sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
BTreeMap::from_iter([
(Nibbles::new(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
(Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
(
Nibbles::from_nibbles([0x5, 0x0]),
Expand All @@ -1146,7 +1147,7 @@ mod tests {
pretty_assertions::assert_eq!(
sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
BTreeMap::from_iter([(
Nibbles::new(),
Nibbles::default(),
SparseNode::new_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]))
),])
);
Expand All @@ -1156,7 +1157,42 @@ mod tests {
// Empty
pretty_assertions::assert_eq!(
sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
BTreeMap::from_iter([(Nibbles::new(), SparseNode::Empty),])
BTreeMap::from_iter([(Nibbles::default(), SparseNode::Empty),])
);
}

#[test]
fn sparse_trie_remove_leaf_blinded() {
let mut sparse = RevealedSparseTrie::default();

let leaf = LeafNode::new(
Nibbles::default(),
alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec(),
);

// Reveal a branch node and one of its children
//
// Branch (Mask = 11)
// ├── 0 -> Hash (Path = 0)
// └── 1 -> Leaf (Path = 1)
sparse
.reveal_node(
Nibbles::default(),
TrieNode::Branch(BranchNode::new(
vec![
RlpNode::word_rlp(&B256::repeat_byte(1)),
RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(),
],
TrieMask::new(0b11),
)),
)
.unwrap();
sparse.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf)).unwrap();

// Removing a blinded leaf should result in an error
assert_matches!(
sparse.remove_leaf(Nibbles::from_nibbles([0x0])),
Err(SparseTrieError::BlindedNode { path, hash }) if path == Nibbles::from_nibbles([0x0]) && hash == B256::repeat_byte(1)
);
}

Expand Down

0 comments on commit 21f303f

Please sign in to comment.