From e0d5735672af43e10927608bd0ce3c006043940d Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Sun, 2 Jul 2023 13:55:27 +0200
Subject: [PATCH] feat: add another distance check (#3501)

---
 crates/blockchain-tree/src/blockchain_tree.rs |  2 +-
 crates/blockchain-tree/src/shareable.rs       |  8 ++++++++
 crates/consensus/beacon/src/engine/mod.rs     |  9 +++++++++
 crates/interfaces/src/blockchain_tree/mod.rs  | 17 ++++++++++++++++-
 crates/storage/provider/src/providers/mod.rs  |  8 ++++++++
 5 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/crates/blockchain-tree/src/blockchain_tree.rs b/crates/blockchain-tree/src/blockchain_tree.rs
index 70cd1cb221..ce10bbd907 100644
--- a/crates/blockchain-tree/src/blockchain_tree.rs
+++ b/crates/blockchain-tree/src/blockchain_tree.rs
@@ -592,7 +592,7 @@ impl<DB: Database, C: Consensus, EF: ExecutorFactory> BlockchainTree<DB, C, EF>
     }
 
     /// Checks the block buffer for the given block.
-    pub fn get_buffered_block(&mut self, hash: &BlockHash) -> Option<&SealedBlockWithSenders> {
+    pub fn get_buffered_block(&self, hash: &BlockHash) -> Option<&SealedBlockWithSenders> {
         self.buffered_blocks.block_by_hash(hash)
     }
 
diff --git a/crates/blockchain-tree/src/shareable.rs b/crates/blockchain-tree/src/shareable.rs
index d23dff9954..ef0ea0323d 100644
--- a/crates/blockchain-tree/src/shareable.rs
+++ b/crates/blockchain-tree/src/shareable.rs
@@ -109,6 +109,14 @@ impl<DB: Database, C: Consensus, EF: ExecutorFactory> BlockchainTreeViewer
         self.tree.read().block_by_hash(block_hash).cloned()
     }
 
+    fn buffered_block_by_hash(&self, block_hash: BlockHash) -> Option<SealedBlock> {
+        self.tree.read().get_buffered_block(&block_hash).map(|b| b.block.clone())
+    }
+
+    fn buffered_header_by_hash(&self, block_hash: BlockHash) -> Option<SealedHeader> {
+        self.tree.read().get_buffered_block(&block_hash).map(|b| b.header.clone())
+    }
+
     fn canonical_blocks(&self) -> BTreeMap<BlockNumber, BlockHash> {
         trace!(target: "blockchain_tree", "Returning canonical blocks in tree");
         self.tree.read().block_indices().canonical_chain().inner().clone()
diff --git a/crates/consensus/beacon/src/engine/mod.rs b/crates/consensus/beacon/src/engine/mod.rs
index d815d91b31..57028c9410 100644
--- a/crates/consensus/beacon/src/engine/mod.rs
+++ b/crates/consensus/beacon/src/engine/mod.rs
@@ -1082,6 +1082,15 @@ where
                                     canonical_tip_num,
                                     downloaded_num_hash.number,
                                 );
+                            } else if let Some(buffered_finalized) =
+                                self.blockchain.buffered_header_by_hash(state.finalized_block_hash)
+                            {
+                                // if we have buffered the finalized block, we should check how far
+                                // we're off
+                                requires_pipeline = self.exceeds_pipeline_run_threshold(
+                                    canonical_tip_num,
+                                    buffered_finalized.number,
+                                );
                             }
                         }
 
diff --git a/crates/interfaces/src/blockchain_tree/mod.rs b/crates/interfaces/src/blockchain_tree/mod.rs
index 8b1eda0a3a..63f628c49d 100644
--- a/crates/interfaces/src/blockchain_tree/mod.rs
+++ b/crates/interfaces/src/blockchain_tree/mod.rs
@@ -172,9 +172,24 @@ pub trait BlockchainTreeViewer: Send + Sync {
 
     /// Returns the block with matching hash from the tree, if it exists.
     ///
-    /// Caution: This will not return blocks from the canonical chain.
+    /// Caution: This will not return blocks from the canonical chain or buffered blocks that are
+    /// disconnected from the canonical chain.
     fn block_by_hash(&self, hash: BlockHash) -> Option<SealedBlock>;
 
+    /// Returns the _buffered_ (disconnected) block with matching hash from the internal buffer if
+    /// it exists.
+    ///
+    /// Caution: Unlike [Self::block_by_hash] this will only return blocks that are currently
+    /// disconnected from the canonical chain.
+    fn buffered_block_by_hash(&self, block_hash: BlockHash) -> Option<SealedBlock>;
+
+    /// Returns the _buffered_ (disconnected) header with matching hash from the internal buffer if
+    /// it exists.
+    ///
+    /// Caution: Unlike [Self::block_by_hash] this will only return headers that are currently
+    /// disconnected from the canonical chain.
+    fn buffered_header_by_hash(&self, block_hash: BlockHash) -> Option<SealedHeader>;
+
     /// Returns true if the tree contains the block with matching hash.
     fn contains(&self, hash: BlockHash) -> bool {
         self.block_by_hash(hash).is_some()
diff --git a/crates/storage/provider/src/providers/mod.rs b/crates/storage/provider/src/providers/mod.rs
index 295d9fcc7f..df1c04b663 100644
--- a/crates/storage/provider/src/providers/mod.rs
+++ b/crates/storage/provider/src/providers/mod.rs
@@ -608,6 +608,14 @@ where
         self.tree.block_by_hash(block_hash)
     }
 
+    fn buffered_block_by_hash(&self, block_hash: BlockHash) -> Option<SealedBlock> {
+        self.tree.buffered_block_by_hash(block_hash)
+    }
+
+    fn buffered_header_by_hash(&self, block_hash: BlockHash) -> Option<SealedHeader> {
+        self.tree.buffered_header_by_hash(block_hash)
+    }
+
     fn canonical_blocks(&self) -> BTreeMap<BlockNumber, BlockHash> {
         self.tree.canonical_blocks()
     }