From fe1847561c52f28103ab8ce8984649fc9f27b5fb Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Fri, 23 Feb 2024 09:21:10 -0600
Subject: [PATCH] Retain authentication when making range requests (#1902)

Needs https://github.com/prefix-dev/async_http_range_reader/pull/9
Closes https://github.com/astral-sh/uv/issues/1709
---
 Cargo.lock                              |  4 ++--
 Cargo.toml                              |  2 +-
 crates/uv-client/src/registry_client.rs | 29 ++++++++++++++++---------
 3 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3fc212fde..fe3169189 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -203,9 +203,9 @@ dependencies = [
 
 [[package]]
 name = "async_http_range_reader"
-version = "0.6.1"
+version = "0.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5143aaae4ec035a5d7cfda666eab896fe5428a2a8ab09ca651a2dce3a8f06912"
+checksum = "cf8eeab30c68da4dc2c51f3afc4327ab06fe0f3f028ca423f7ca398c7ed8c5e7"
 dependencies = [
  "bisection",
  "futures",
diff --git a/Cargo.toml b/Cargo.toml
index 008873e38..0d250e41a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,7 +21,7 @@ anstream = { version = "0.6.5" }
 anyhow = { version = "1.0.79" }
 async-compression = { version = "0.4.6" }
 async-trait = { version = "0.1.77" }
-async_http_range_reader = { version = "0.6.1" }
+async_http_range_reader = { version = "0.7.0" }
 async_zip = { git = "https://github.com/charliermarsh/rs-async-zip", rev = "d76801da0943de985254fc6255c0e476b57c5836", features = ["deflate"] }
 base64 = { version = "0.21.7" }
 cachedir = { version = "0.3.1" }
diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs
index 3416c3027..ed7eb4da2 100644
--- a/crates/uv-client/src/registry_client.rs
+++ b/crates/uv-client/src/registry_client.rs
@@ -7,6 +7,7 @@ use std::str::FromStr;
 use async_http_range_reader::AsyncHttpRangeReader;
 use futures::{FutureExt, TryStreamExt};
 
+use http::HeaderMap;
 use reqwest::{Client, ClientBuilder, Response, StatusCode};
 use reqwest_retry::policies::ExponentialBackoff;
 use reqwest_retry::RetryTransientMiddleware;
@@ -440,14 +441,28 @@ impl RegistryClient {
             Connectivity::Offline => CacheControl::AllowStale,
         };
 
+        let client = self.client_raw.clone();
+        let req = self
+            .client
+            .uncached()
+            .head(url.clone())
+            .build()
+            .map_err(ErrorKind::RequestError)?;
+
+        // Copy authorization headers from the HEAD request to subsequent requests
+        let mut headers = HeaderMap::default();
+        if let Some(authorization) = req.headers().get("authorization") {
+            headers.append("authorization", authorization.clone());
+        }
+
         // This response callback is special, we actually make a number of subsequent requests to
         // fetch the file from the remote zip.
-        let client = self.client_raw.clone();
         let read_metadata_range_request = |response: Response| {
             async {
-                let mut reader = AsyncHttpRangeReader::from_head_response(client, response)
-                    .await
-                    .map_err(ErrorKind::AsyncHttpRangeReader)?;
+                let mut reader =
+                    AsyncHttpRangeReader::from_head_response(client, response, headers)
+                        .await
+                        .map_err(ErrorKind::AsyncHttpRangeReader)?;
                 trace!("Getting metadata for {filename} by range request");
                 let text = wheel_metadata_from_remote_zip(filename, &mut reader).await?;
                 let metadata = Metadata21::parse(text.as_bytes()).map_err(|err| {
@@ -463,12 +478,6 @@ impl RegistryClient {
             .instrument(info_span!("read_metadata_range_request", wheel = %filename))
         };
 
-        let req = self
-            .client
-            .uncached()
-            .head(url.clone())
-            .build()
-            .map_err(ErrorKind::RequestError)?;
         let result = self
             .client
             .get_serde(