feat: add header to stage command (paradigmxyz#10127)

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()));