From b01ccc271d04a9f702b8d1c172fbc8e4b4fb7be4 Mon Sep 17 00:00:00 2001
From: greged93 <82421016+greged93@users.noreply.github.com>
Date: Tue, 6 Aug 2024 20:04:00 +0200
Subject: [PATCH] feat: add header to stage command (#10127)

---
 crates/cli/commands/src/stage/run.rs | 65 +++++++++++++++++++++++++---
 1 file changed, 58 insertions(+), 7 deletions(-)

diff --git a/crates/cli/commands/src/stage/run.rs b/crates/cli/commands/src/stage/run.rs
index a7349f6e95fb..7b097be42182 100644
--- a/crates/cli/commands/src/stage/run.rs
+++ b/crates/cli/commands/src/stage/run.rs
@@ -2,20 +2,24 @@
 //!
 //! Stage debugging tool
 
-use std::{any::Any, net::SocketAddr, sync::Arc, time::Instant};
-
+use crate::common::{AccessRights, Environment, EnvironmentArgs};
 use clap::Parser;
 use reth_beacon_consensus::EthBeaconConsensus;
 use reth_chainspec::ChainSpec;
 use reth_cli_runner::CliContext;
 use reth_cli_util::get_secret_key;
 use reth_config::config::{HashingConfig, SenderRecoveryConfig, TransactionLookupConfig};
-use reth_downloaders::bodies::bodies::BodiesDownloaderBuilder;
+use reth_downloaders::{
+    bodies::bodies::BodiesDownloaderBuilder,
+    headers::reverse_headers::ReverseHeadersDownloaderBuilder,
+};
 use reth_evm::execute::BlockExecutorProvider;
 use reth_exex::ExExManagerHandle;
 use reth_network::BlockDownloaderProvider;
+use reth_network_p2p::HeadersClient;
 use reth_node_core::{
     args::{NetworkArgs, StageEnum},
+    primitives::BlockHashOrNumber,
     version::{
         BUILD_PROFILE_NAME, CARGO_PKG_VERSION, VERGEN_BUILD_TIMESTAMP, VERGEN_CARGO_FEATURES,
         VERGEN_CARGO_TARGET_TRIPLE, VERGEN_GIT_SHA,
@@ -32,16 +36,17 @@ use reth_provider::{
 };
 use reth_stages::{
     stages::{
-        AccountHashingStage, BodyStage, ExecutionStage, IndexAccountHistoryStage,
+        AccountHashingStage, BodyStage, ExecutionStage, HeaderStage, IndexAccountHistoryStage,
         IndexStorageHistoryStage, MerkleStage, SenderRecoveryStage, StorageHashingStage,
         TransactionLookupStage,
     },
-    ExecInput, ExecOutput, ExecutionStageThresholds, Stage, StageExt, UnwindInput, UnwindOutput,
+    ExecInput, ExecOutput, ExecutionStageThresholds, Stage, StageError, StageExt, UnwindInput,
+    UnwindOutput,
 };
+use std::{any::Any, net::SocketAddr, sync::Arc, time::Instant};
+use tokio::sync::watch;
 use tracing::*;
 
-use crate::common::{AccessRights, Environment, EnvironmentArgs};
-
 /// `reth stage` command
 #[derive(Debug, Parser)]
 pub struct Command {
@@ -138,6 +143,52 @@ impl Command {
 
         let (mut exec_stage, mut unwind_stage): (Box<dyn Stage<_>>, Option<Box<dyn Stage<_>>>) =
             match self.stage {
+                StageEnum::Headers => {
+                    let consensus =
+                        Arc::new(EthBeaconConsensus::new(provider_factory.chain_spec()));
+
+                    let network_secret_path = self
+                        .network
+                        .p2p_secret_key
+                        .clone()
+                        .unwrap_or_else(|| data_dir.p2p_secret());
+                    let p2p_secret_key = get_secret_key(&network_secret_path)?;
+
+                    let default_peers_path = data_dir.known_peers();
+
+                    let network = self
+                        .network
+                        .network_config(
+                            &config,
+                            provider_factory.chain_spec(),
+                            p2p_secret_key,
+                            default_peers_path,
+                        )
+                        .build(provider_factory.clone())
+                        .start_network()
+                        .await?;
+                    let fetch_client = Arc::new(network.fetch_client().await?);
+
+                    // Use `to` as the tip for the stage
+                    let tip = fetch_client
+                        .get_header(BlockHashOrNumber::Number(self.to))
+                        .await?
+                        .into_data()
+                        .ok_or(StageError::MissingSyncGap)?;
+                    let (_, rx) = watch::channel(tip.hash_slow());
+
+                    (
+                        Box::new(HeaderStage::new(
+                            provider_factory.clone(),
+                            ReverseHeadersDownloaderBuilder::new(config.stages.headers)
+                                .build(fetch_client, consensus.clone()),
+                            rx,
+                            consensus,
+                            etl_config,
+                        )),
+                        None,
+                    )
+                }
                 StageEnum::Bodies => {
                     let consensus =
                         Arc::new(EthBeaconConsensus::new(provider_factory.chain_spec()));