From 43bd94ac4e9638f8b6fc455739f26a4d696e83dd Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com> Date: Fri, 17 Jan 2025 11:00:57 +0000 Subject: [PATCH] fix(trie): reveal blinded node along with masks in sparse trie (#13827) --- crates/trie/sparse/src/blinded.rs | 17 +++++++-- crates/trie/sparse/src/trie.rs | 50 +++++++++++++++++++-------- crates/trie/trie/src/proof/blinded.rs | 28 +++++++++------ crates/trie/trie/src/witness.rs | 6 ++-- 4 files changed, 71 insertions(+), 30 deletions(-) diff --git a/crates/trie/sparse/src/blinded.rs b/crates/trie/sparse/src/blinded.rs index 28a41ba11..46f7e1655 100644 --- a/crates/trie/sparse/src/blinded.rs +++ b/crates/trie/sparse/src/blinded.rs @@ -2,7 +2,7 @@ use alloy_primitives::{Bytes, B256}; use reth_execution_errors::SparseTrieError; -use reth_trie_common::Nibbles; +use reth_trie_common::{Nibbles, TrieMask}; /// Factory for instantiating blinded node providers. pub trait BlindedProviderFactory { @@ -18,10 +18,21 @@ pub trait BlindedProviderFactory { fn storage_node_provider(&self, account: B256) -> Self::StorageNodeProvider; } +/// Revealed blinded trie node. +#[derive(Debug)] +pub struct RevealedNode { + /// Raw trie node. + pub node: Bytes, + /// Branch node tree mask, if any. + pub tree_mask: Option, + /// Branch node hash mask, if any. + pub hash_mask: Option, +} + /// Trie node provider for retrieving blinded nodes. pub trait BlindedProvider { /// Retrieve blinded node by path. - fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError>; + fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError>; } /// Default blinded node provider factory that creates [`DefaultBlindedProvider`]. @@ -46,7 +57,7 @@ impl BlindedProviderFactory for DefaultBlindedProviderFactory { pub struct DefaultBlindedProvider; impl BlindedProvider for DefaultBlindedProvider { - fn blinded_node(&mut self, _path: &Nibbles) -> Result, SparseTrieError> { + fn blinded_node(&mut self, _path: &Nibbles) -> Result, SparseTrieError> { Ok(None) } } diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 3441daa4b..9c35f7dee 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -1,4 +1,4 @@ -use crate::blinded::{BlindedProvider, DefaultBlindedProvider}; +use crate::blinded::{BlindedProvider, DefaultBlindedProvider, RevealedNode}; use alloy_primitives::{ hex, keccak256, map::{Entry, HashMap, HashSet}, @@ -945,14 +945,24 @@ impl RevealedSparseTrie

{ if self.updates.is_some() { // Check if the extension node child is a hash that needs to be revealed if self.nodes.get(¤t).unwrap().is_hash() { - if let Some(node) = self.provider.blinded_node(¤t)? { + if let Some(RevealedNode { node, tree_mask, hash_mask }) = + self.provider.blinded_node(¤t)? + { let decoded = TrieNode::decode(&mut &node[..])?; - trace!(target: "trie::sparse", ?current, ?decoded, "Revealing extension node child"); - // We'll never have to update the revealed child node, only - // remove or do nothing, so - // we can safely ignore the hash mask here and - // pass `None`. - self.reveal_node(current.clone(), decoded, None, None)?; + trace!( + target: "trie::sparse", + ?current, + ?decoded, + ?tree_mask, + ?hash_mask, + "Revealing extension node child", + ); + self.reveal_node( + current.clone(), + decoded, + tree_mask, + hash_mask, + )?; } } } @@ -1000,6 +1010,7 @@ impl RevealedSparseTrie

{ return Err(SparseTrieErrorKind::BlindedNode { path: path.clone(), hash }.into()) } + trace!(target: "trie::sparse", ?path, "Leaf node is not present in the trie"); // Leaf is not present in the trie. return Ok(()) } @@ -1098,13 +1109,24 @@ impl RevealedSparseTrie

{ if self.nodes.get(&child_path).unwrap().is_hash() { trace!(target: "trie::sparse", ?child_path, "Retrieving remaining blinded branch child"); - if let Some(node) = self.provider.blinded_node(&child_path)? { + if let Some(RevealedNode { node, tree_mask, hash_mask }) = + self.provider.blinded_node(&child_path)? + { let decoded = TrieNode::decode(&mut &node[..])?; - trace!(target: "trie::sparse", ?child_path, ?decoded, "Revealing remaining blinded branch child"); - // We'll never have to update the revealed branch node, only remove - // or do nothing, so we can safely ignore the hash mask here and - // pass `None`. - self.reveal_node(child_path.clone(), decoded, None, None)?; + trace!( + target: "trie::sparse", + ?child_path, + ?decoded, + ?tree_mask, + ?hash_mask, + "Revealing remaining blinded branch child" + ); + self.reveal_node( + child_path.clone(), + decoded, + tree_mask, + hash_mask, + )?; } } diff --git a/crates/trie/trie/src/proof/blinded.rs b/crates/trie/trie/src/proof/blinded.rs index 57d0de97f..9b838c2e9 100644 --- a/crates/trie/trie/src/proof/blinded.rs +++ b/crates/trie/trie/src/proof/blinded.rs @@ -2,11 +2,13 @@ use super::{Proof, StorageProof}; use crate::{hashed_cursor::HashedCursorFactory, trie_cursor::TrieCursorFactory}; use alloy_primitives::{ map::{HashMap, HashSet}, - Bytes, B256, + B256, }; use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind}; use reth_trie_common::{prefix_set::TriePrefixSetsMut, Nibbles}; -use reth_trie_sparse::blinded::{pad_path_to_key, BlindedProvider, BlindedProviderFactory}; +use reth_trie_sparse::blinded::{ + pad_path_to_key, BlindedProvider, BlindedProviderFactory, RevealedNode, +}; use std::sync::Arc; use tracing::trace; @@ -85,17 +87,20 @@ where T: TrieCursorFactory + Clone + Send + Sync, H: HashedCursorFactory + Clone + Send + Sync, { - fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError> { + fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError> { let targets = HashMap::from_iter([(pad_path_to_key(path), HashSet::default())]); - let proof = + let mut proof = Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone()) .with_prefix_sets_mut(self.prefix_sets.as_ref().clone()) + .with_branch_node_masks(true) .multiproof(targets) .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; let node = proof.account_subtree.into_inner().remove(path); - + let tree_mask = proof.branch_node_tree_masks.remove(path); + let hash_mask = proof.branch_node_hash_masks.remove(path); trace!(target: "trie::proof::blinded", ?path, ?node, "Blinded node for account trie"); - Ok(node) + + Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask })) } } @@ -129,21 +134,24 @@ where T: TrieCursorFactory + Clone + Send + Sync, H: HashedCursorFactory + Clone + Send + Sync, { - fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError> { + fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError> { let targets = HashSet::from_iter([pad_path_to_key(path)]); let storage_prefix_set = self.prefix_sets.storage_prefix_sets.get(&self.account).cloned().unwrap_or_default(); - let proof = StorageProof::new_hashed( + let mut proof = StorageProof::new_hashed( self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone(), self.account, ) .with_prefix_set_mut(storage_prefix_set) + .with_branch_node_masks(true) .storage_multiproof(targets) .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; let node = proof.subtree.into_inner().remove(path); - + let tree_mask = proof.branch_node_tree_masks.remove(path); + let hash_mask = proof.branch_node_hash_masks.remove(path); trace!(target: "trie::proof::blinded", account = ?self.account, ?path, ?node, "Blinded node for storage trie"); - Ok(node) + + Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask })) } } diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index 5d43c0ea1..b044b476c 100644 --- a/crates/trie/trie/src/witness.rs +++ b/crates/trie/trie/src/witness.rs @@ -17,7 +17,7 @@ use reth_execution_errors::{ }; use reth_trie_common::{MultiProofTargets, Nibbles}; use reth_trie_sparse::{ - blinded::{BlindedProvider, BlindedProviderFactory}, + blinded::{BlindedProvider, BlindedProviderFactory, RevealedNode}, SparseStateTrie, }; use std::sync::{mpsc, Arc}; @@ -244,11 +244,11 @@ impl

WitnessBlindedProvider

{ } impl BlindedProvider for WitnessBlindedProvider

{ - fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError> { + fn blinded_node(&mut self, path: &Nibbles) -> Result, SparseTrieError> { let maybe_node = self.provider.blinded_node(path)?; if let Some(node) = &maybe_node { self.tx - .send(node.clone()) + .send(node.node.clone()) .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; } Ok(maybe_node)