fix(trie): reveal blinded node along with masks in sparse trie (#13827)

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<TrieMask>,
+    /// Branch node hash mask, if any.
+    pub hash_mask: Option<TrieMask>,
+}
+
 /// Trie node provider for retrieving blinded nodes.
 pub trait BlindedProvider {
     /// Retrieve blinded node by path.
-    fn blinded_node(&mut self, path: &Nibbles) -> Result<Option<Bytes>, SparseTrieError>;
+    fn blinded_node(&mut self, path: &Nibbles) -> Result<Option<RevealedNode>, 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<Option<Bytes>, SparseTrieError> {
+    fn blinded_node(&mut self, _path: &Nibbles) -> Result<Option<RevealedNode>, SparseTrieError> {
         Ok(None)
     }
 }

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<P: BlindedProvider> RevealedSparseTrie<P> {
                         if self.updates.is_some() {
                             // Check if the extension node child is a hash that needs to be revealed
                             if self.nodes.get(&current).unwrap().is_hash() {
-                                if let Some(node) = self.provider.blinded_node(&current)? {
+                                if let Some(RevealedNode { node, tree_mask, hash_mask }) =
+                                    self.provider.blinded_node(&current)?
+                                {
                                     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<P: BlindedProvider> RevealedSparseTrie<P> {
                 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<P: BlindedProvider> RevealedSparseTrie<P> {
 
                         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,
+                                )?;
                             }
                         }
 

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<Option<Bytes>, SparseTrieError> {
+    fn blinded_node(&mut self, path: &Nibbles) -> Result<Option<RevealedNode>, 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<Option<Bytes>, SparseTrieError> {
+    fn blinded_node(&mut self, path: &Nibbles) -> Result<Option<RevealedNode>, 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 }))
     }
 }

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<P> WitnessBlindedProvider<P> {
 }
 
 impl<P: BlindedProvider> BlindedProvider for WitnessBlindedProvider<P> {
-    fn blinded_node(&mut self, path: &Nibbles) -> Result<Option<Bytes>, SparseTrieError> {
+    fn blinded_node(&mut self, path: &Nibbles) -> Result<Option<RevealedNode>, 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)