From 76c5aef911d909161be47113cd88b7ea36fafffd Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Thu, 31 Oct 2024 11:53:59 +0100
Subject: [PATCH] fix(trie): move to sibling on invalid tree mask (#12193)

Co-authored-by: Federico Gimenez <fgimenez@users.noreply.github.com>
---
 crates/trie/trie/src/metrics.rs | 22 +++++++++++++++++++++-
 crates/trie/trie/src/walker.rs  | 32 ++++++++++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/crates/trie/trie/src/metrics.rs b/crates/trie/trie/src/metrics.rs
index 7582f37418..006dc7e365 100644
--- a/crates/trie/trie/src/metrics.rs
+++ b/crates/trie/trie/src/metrics.rs
@@ -1,5 +1,5 @@
 use crate::stats::TrieStats;
-use metrics::Histogram;
+use metrics::{Counter, Histogram};
 use reth_metrics::Metrics;
 
 /// Wrapper for state root metrics.
@@ -63,3 +63,23 @@ impl TrieType {
         }
     }
 }
+
+/// Metrics for trie walker
+#[derive(Clone, Metrics)]
+#[metrics(scope = "trie.walker")]
+pub struct WalkerMetrics {
+    /// The number of subnodes out of order due to wrong tree mask.
+    out_of_order_subnode: Counter,
+}
+
+impl WalkerMetrics {
+    /// Create new metrics for the given trie type.
+    pub fn new(ty: TrieType) -> Self {
+        Self::new_with_labels(&[("type", ty.as_str())])
+    }
+
+    /// Increment `out_of_order_subnode`.
+    pub fn inc_out_of_order_subnode(&self, amount: u64) {
+        self.out_of_order_subnode.increment(amount);
+    }
+}
diff --git a/crates/trie/trie/src/walker.rs b/crates/trie/trie/src/walker.rs
index e75a96d0f1..aaff293b37 100644
--- a/crates/trie/trie/src/walker.rs
+++ b/crates/trie/trie/src/walker.rs
@@ -7,6 +7,9 @@ use alloy_primitives::B256;
 use reth_storage_errors::db::DatabaseError;
 use std::collections::HashSet;
 
+#[cfg(feature = "metrics")]
+use crate::metrics::WalkerMetrics;
+
 /// `TrieWalker` is a structure that enables traversal of a Merkle trie.
 /// It allows moving through the trie in a depth-first manner, skipping certain branches
 /// if they have not changed.
@@ -24,13 +27,23 @@ pub struct TrieWalker<C> {
     pub changes: PrefixSet,
     /// The retained trie node keys that need to be removed.
     removed_keys: Option<HashSet<Nibbles>>,
+    #[cfg(feature = "metrics")]
+    /// Walker metrics.
+    metrics: WalkerMetrics,
 }
 
 impl<C> TrieWalker<C> {
     /// Constructs a new `TrieWalker` from existing stack and a cursor.
     pub fn from_stack(cursor: C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
-        let mut this =
-            Self { cursor, changes, stack, can_skip_current_node: false, removed_keys: None };
+        let mut this = Self {
+            cursor,
+            changes,
+            stack,
+            can_skip_current_node: false,
+            removed_keys: None,
+            #[cfg(feature = "metrics")]
+            metrics: WalkerMetrics::default(),
+        };
         this.update_skip_node();
         this
     }
@@ -113,6 +126,8 @@ impl<C: TrieCursor> TrieWalker<C> {
             stack: vec![CursorSubNode::default()],
             can_skip_current_node: false,
             removed_keys: None,
+            #[cfg(feature = "metrics")]
+            metrics: WalkerMetrics::default(),
         };
 
         // Set up the root node of the trie in the stack, if it exists.
@@ -179,6 +194,19 @@ impl<C: TrieCursor> TrieWalker<C> {
             self.stack[0].set_nibble(key[0] as i8);
         }
 
+        // The current tree mask might have been set incorrectly.
+        // Sanity check that the newly retrieved trie node key is the child of the last item
+        // on the stack. If not, advance to the next sibling instead of adding the node to the
+        // stack.
+        if let Some(subnode) = self.stack.last() {
+            if !key.starts_with(subnode.full_key()) {
+                #[cfg(feature = "metrics")]
+                self.metrics.inc_out_of_order_subnode(1);
+                self.move_to_next_sibling(false)?;
+                return Ok(())
+            }
+        }
+
         // Create a new CursorSubNode and push it to the stack.
         let subnode = CursorSubNode::new(key, Some(node));
         let nibble = subnode.nibble();