From 92785ee8acf7de148ebe769e1db0c459d6bdb6dc Mon Sep 17 00:00:00 2001
From: Anukul Sangwan <anukulsangwan@icloud.com>
Date: Thu, 12 Dec 2024 18:05:38 +0200
Subject: [PATCH] set auth header for fork

---
 crates/cheatcodes/src/evm/fork.rs | 3 +++
 crates/common/src/provider/mod.rs | 6 ++++++
 crates/evm/core/src/fork/multi.rs | 1 +
 crates/evm/core/src/opts.rs       | 4 ++++
 4 files changed, 14 insertions(+)

diff --git a/crates/cheatcodes/src/evm/fork.rs b/crates/cheatcodes/src/evm/fork.rs
index 9970a9641f0f..e78aa3be9f69 100644
--- a/crates/cheatcodes/src/evm/fork.rs
+++ b/crates/cheatcodes/src/evm/fork.rs
@@ -332,6 +332,9 @@ fn create_fork_request(
     evm_opts.fork_block_number = block;
     evm_opts.fork_retries = rpc_endpoint.config.retries;
     evm_opts.fork_retry_backoff = rpc_endpoint.config.retry_backoff;
+    if let Some(Ok(auth)) = rpc_endpoint.auth {
+        evm_opts.fork_headers = Some(vec![format!("Authorization: {auth}")]);
+    }
     let fork = CreateFork {
         enable_caching: !ccx.state.config.no_storage_caching &&
             ccx.state.config.rpc_storage_caching.enable_for_endpoint(&url),
diff --git a/crates/common/src/provider/mod.rs b/crates/common/src/provider/mod.rs
index 75efdb869f9f..cd34f94e9520 100644
--- a/crates/common/src/provider/mod.rs
+++ b/crates/common/src/provider/mod.rs
@@ -241,6 +241,12 @@ impl ProviderBuilder {
         self
     }
 
+    /// Sets http headers. If `None`, defaults to the already-set value.
+    pub fn maybe_headers(mut self, headers: Option<Vec<String>>) -> Self {
+        self.headers = headers.unwrap_or(self.headers);
+        self
+    }
+
     /// Constructs the `RetryProvider` taking all configs into account.
     pub fn build(self) -> Result<RetryProvider> {
         let Self {
diff --git a/crates/evm/core/src/fork/multi.rs b/crates/evm/core/src/fork/multi.rs
index 112e1eeda7d4..490ae7379ae5 100644
--- a/crates/evm/core/src/fork/multi.rs
+++ b/crates/evm/core/src/fork/multi.rs
@@ -514,6 +514,7 @@ async fn create_fork(mut fork: CreateFork) -> eyre::Result<(ForkId, CreatedFork,
         ProviderBuilder::new(fork.url.as_str())
             .maybe_max_retry(fork.evm_opts.fork_retries)
             .maybe_initial_backoff(fork.evm_opts.fork_retry_backoff)
+            .maybe_headers(fork.evm_opts.fork_headers.clone())
             .compute_units_per_second(fork.evm_opts.get_compute_units_per_second())
             .build()?,
     );
diff --git a/crates/evm/core/src/opts.rs b/crates/evm/core/src/opts.rs
index 76be5937bfde..4c13369b6ae6 100644
--- a/crates/evm/core/src/opts.rs
+++ b/crates/evm/core/src/opts.rs
@@ -29,6 +29,9 @@ pub struct EvmOpts {
     /// Initial retry backoff.
     pub fork_retry_backoff: Option<u64>,
 
+    /// Headers to use with `fork_url`
+    pub fork_headers: Option<Vec<String>>,
+
     /// The available compute units per second.
     ///
     /// See also <https://docs.alchemy.com/reference/compute-units#what-are-cups-compute-units-per-second>
@@ -80,6 +83,7 @@ impl Default for EvmOpts {
             fork_block_number: None,
             fork_retries: None,
             fork_retry_backoff: None,
+            fork_headers: None,
             compute_units_per_second: None,
             no_rpc_rate_limit: false,
             no_storage_caching: false,