From c0e03518a0512c37b2d67ca7fd9244fa131287f5 Mon Sep 17 00:00:00 2001 From: BGluth Date: Fri, 9 Aug 2024 15:16:57 -0400 Subject: [PATCH] Feat/237 mpt trie ext to branch collapse error (#455) * Added `Nibbles::empty()` - While `Nibbles::default()` already does this, I think having an explicit method that makes an empty `Nibbles` is a bit cleaner, even if it's a bit redundant. * Now returns an error if we collapse an `E --> H` - Unsafe, as if the extension was pointing to a (hashed) leaf, then we must collapse the extension into the leaf. - However, because the leaf is hashed, we can not tell that the hash node is from a leaf and also can not extract the leaf's key. - It's the user's responsibility to ensure that this scenario never occurs, so we need to return an error when it does. * Requested PR changes for #455 --- mpt_trie/src/nibbles.rs | 21 ++++++++++- mpt_trie/src/trie_ops.rs | 75 ++++++++++++++++++++++++++++++++++------ 2 files changed, 85 insertions(+), 11 deletions(-) diff --git a/mpt_trie/src/nibbles.rs b/mpt_trie/src/nibbles.rs index d72e027cf..982692c27 100644 --- a/mpt_trie/src/nibbles.rs +++ b/mpt_trie/src/nibbles.rs @@ -254,7 +254,7 @@ impl From for NibblesIntern { } } -#[derive(Copy, Clone, Deserialize, Default, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +#[derive(Copy, Clone, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] /// A sequence of nibbles which is used as the key type into /// [`PartialTrie`][`crate::partial_trie::PartialTrie`]. /// @@ -323,6 +323,14 @@ impl Debug for Nibbles { } } +/// While we could just derive `Default` and it would be correct, it's a bit +/// cleaner to instead call [`Nibbles::new`] explicitly. +impl Default for Nibbles { + fn default() -> Self { + Self::new() + } +} + impl FromStr for Nibbles { type Err = StrToNibblesError; @@ -355,6 +363,17 @@ impl UpperHex for Nibbles { } impl Nibbles { + /// Create `Nibbles` that is empty. + /// + /// Note that mean that the key size is `0` and does not mean that the key + /// contains the `0` [`Nibble`]. + pub fn new() -> Self { + Self { + count: 0, + packed: NibblesIntern::default(), + } + } + /// Creates `Nibbles` from big endian bytes. /// /// Returns an error if the byte slice is empty or is longer than `32` diff --git a/mpt_trie/src/trie_ops.rs b/mpt_trie/src/trie_ops.rs index d142ea49a..1d6c0ab4d 100644 --- a/mpt_trie/src/trie_ops.rs +++ b/mpt_trie/src/trie_ops.rs @@ -31,11 +31,24 @@ pub enum TrieOpError { #[error("Attempted to delete a value that ended up inside a hash node! (hash: {0})")] HashNodeDeleteError(H256), - /// An error that occurs when encontered an unexisting type of node during - /// an extension node collapse. - #[error("Extension managed to get an unexisting child node type! (child: {0})")] + /// An error that occurs when we encounter an non-existing type of node + /// during an extension node collapse. + #[error("Extension managed to get an non-existing child node type! (child: {0})")] HashNodeExtError(TrieNodeType), + /// An error that occurs when we attempted to collapse an extension node + /// into a hash node. + /// + /// If this occurs, then there is a chance that we can not collapse + /// correctly and will produce the incorrect trie (and also the incorrect + /// trie hash). If the hash node is a hash of a leaf, then we need to + /// collapse the extension key into the leaf. However, this information + /// is lost if the node is hashed, and we can not tell if the hash node + /// was made from a leaf. As such, it's the responsibility of whoever is + /// constructing & mutating the trie that this will never occur. + #[error("Attempted to collapse an extension node into a hash node! This is unsafe! (See https://github.com/0xPolygonZero/zk_evm/issues/237 for more info) (Extension key: {0:x}, child hash node: {1:x})")] + ExtensionCollapsedIntoHashError(Nibbles, H256), + /// Failed to insert a hash node into the trie. #[error("Attempted to place a hash node on an existing node! (hash: {0})")] ExistingHashNodeError(H256), @@ -360,6 +373,11 @@ impl Node { } } + /// Deletes a key if it exists in the trie. + /// + /// If the key exists, then the existing node value that was deleted is + /// returned. Otherwise, if the key is not present, then `None` is returned + /// instead. pub(crate) fn trie_delete(&mut self, k: K) -> TrieOpResult>> where K: Into, @@ -368,8 +386,10 @@ impl Node { trace!("Deleting a leaf node with key {} if it exists", k); delete_intern(&self.clone(), k)?.map_or(Ok(None), |(updated_root, deleted_val)| { - // Final check at the root if we have an extension node - let wrapped_node = try_collapse_if_extension(updated_root)?; + // Final check at the root if we have an extension node. While this check also + // exists as we recursively traverse down the trie, it can not perform this + // check on the root node. + let wrapped_node = try_collapse_if_extension(updated_root, &Nibbles::default())?; let node_ref: &Node = &wrapped_node; *self = node_ref.clone(); @@ -521,12 +541,15 @@ fn delete_intern( { false => { // Branch stays. + let mut updated_children = children.clone(); updated_children[nibble as usize] = - try_collapse_if_extension(updated_child)?; + try_collapse_if_extension(updated_child, &curr_k)?; branch(updated_children, value.clone()) } true => { + // We need to collapse the branch into an extension/leaf node. + let (child_nibble, non_empty_node) = get_other_non_empty_child_and_nibble_in_two_elem_branch( children, nibble, @@ -559,7 +582,7 @@ fn delete_intern( delete_intern(child, curr_k).and_then(|res| { res.map_or(Ok(None), |(updated_child, value_deleted)| { let updated_node = - collapse_ext_node_if_needed(ext_nibbles, &updated_child)?; + collapse_ext_node_if_needed(ext_nibbles, &updated_child, &curr_k)?; Ok(Some((updated_node, value_deleted))) }) }) @@ -576,16 +599,44 @@ fn delete_intern( } } -fn try_collapse_if_extension(node: WrappedNode) -> TrieOpResult> { +fn try_collapse_if_extension( + node: WrappedNode, + curr_key: &Nibbles, +) -> TrieOpResult> { match node.as_ref() { - Node::Extension { nibbles, child } => collapse_ext_node_if_needed(nibbles, child), + Node::Extension { nibbles, child } => collapse_ext_node_if_needed(nibbles, child, curr_key), _ => Ok(node), } } +/// Attempt to collapse an extension node if we are required. +/// +/// The scenarios where we are required to do so are where the extension node is +/// pointing to a: +/// - Extension (the parent extension absorbs the child's key). +/// - Leaf (the leaf absorbs the extension's key). +/// +/// While an extension does not collapse when its child is a branch, we need to +/// still check that a specific edge case does not exist for the branch node. +/// Specifically, if all of the following holds true for the branch where: +/// - It has exactly two children. +/// - One child that is a hash node. +/// - One child that is a leaf node. +/// - The leaf child ends up getting deleted. +/// +/// Then we need to return an error, because if the hash node was created from a +/// leaf node, we are required to collapse the extension key into the leaf node. +/// However, since it's a hash node, we: +/// - Have no idea what the original node was. +/// - Can not access the underlying key if we needed to collapse it. +/// +/// Because of this, we need to rely on the user to not allow `mpt_trie` to +/// arrive at this state, as we can not ensure that we will be able to produce +/// the correct trie. fn collapse_ext_node_if_needed( ext_nibbles: &Nibbles, child: &WrappedNode, + curr_key: &Nibbles, ) -> TrieOpResult> { trace!( "Collapsing extension node ({:x}) with child {}...", @@ -606,7 +657,11 @@ fn collapse_ext_node_if_needed( nibbles: leaf_nibbles, value, } => Ok(leaf(ext_nibbles.merge_nibbles(leaf_nibbles), value.clone())), - Node::Hash(_) => Ok(extension(*ext_nibbles, child.clone())), + Node::Hash(h) => Err(TrieOpError::ExtensionCollapsedIntoHashError( + curr_key.merge_nibbles(ext_nibbles), + *h, + )), + // Can never do this safely, so return an error. _ => Err(TrieOpError::HashNodeExtError(TrieNodeType::from(child))), } }