Support redirects in `uv publish` (#17130)

<!--
Thank you for contributing to uv! To help us out with reviewing, please
consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

Follow redirects for `uv publish`. Related issue:
https://github.com/astral-sh/uv/issues/17126.

## Test Plan

<!-- How was it tested? -->
Added a unit test to test the custom redirect logic.

---------

Co-authored-by: konstin <konstin@mailbox.org>

Cargo.lock

@@ -6464,6 +6464,7 @@ name = "uv-publish"
 version = "0.0.7"
 dependencies = [
  "ambient-id",
+ "anstream",
  "astral-reqwest-middleware",
  "astral-reqwest-retry",
  "astral-tokio-tar",
@@ -6497,6 +6498,7 @@ dependencies = [
  "uv-redacted",
  "uv-static",
  "uv-warnings",
+ "wiremock",
 ]
 
 [[package]]

crates/uv-client/src/base_client.rs

@@ -50,7 +50,7 @@ pub const DEFAULT_RETRIES: u32 = 3;
 /// Maximum number of redirects to follow before giving up.
 ///
 /// This is the default used by [`reqwest`].
-const DEFAULT_MAX_REDIRECTS: u32 = 10;
+pub const DEFAULT_MAX_REDIRECTS: u32 = 10;
 
 /// Selectively skip parts or the entire auth middleware.
 #[derive(Debug, Clone, Copy, Default)]
@@ -104,6 +104,8 @@ pub enum RedirectPolicy {
     BypassMiddleware,
     /// Handle redirects manually, re-triggering our custom middleware for each request.
     RetriggerMiddleware,
+    /// No redirect for non-cloneable (e.g., streaming) requests with custom redirect logic.
+    NoRedirect,
 }
 
 impl RedirectPolicy {
@@ -111,6 +113,7 @@ impl RedirectPolicy {
         match self {
             Self::BypassMiddleware => reqwest::redirect::Policy::default(),
             Self::RetriggerMiddleware => reqwest::redirect::Policy::none(),
+            Self::NoRedirect => reqwest::redirect::Policy::none(),
         }
     }
 }
@@ -729,6 +732,7 @@ impl RedirectClientWithMiddleware {
         match self.redirect_policy {
             RedirectPolicy::BypassMiddleware => self.client.execute(req).await,
             RedirectPolicy::RetriggerMiddleware => self.execute_with_redirect_handling(req).await,
+            RedirectPolicy::NoRedirect => self.client.execute(req).await,
         }
     }
 

crates/uv-client/src/lib.rs

@@ -1,7 +1,7 @@
 pub use base_client::{
-    AuthIntegration, BaseClient, BaseClientBuilder, DEFAULT_RETRIES, ExtraMiddleware,
+    AuthIntegration, BaseClient, BaseClientBuilder, DEFAULT_MAX_REDIRECTS, DEFAULT_RETRIES,
-    RedirectClientWithMiddleware, RequestBuilder, RetryParsingError, UvRetryableStrategy,
+    ExtraMiddleware, RedirectClientWithMiddleware, RedirectPolicy, RequestBuilder,
-    is_transient_network_error,
+    RetryParsingError, UvRetryableStrategy, is_transient_network_error,
 };
 pub use cached_client::{CacheControl, CachedClient, CachedClientError, DataWithCachePolicy};
 pub use error::{Error, ErrorKind, WrappedReqwestError};

crates/uv-publish/Cargo.toml

@@ -46,8 +46,10 @@ tokio = { workspace = true }
 tokio-util = { workspace = true, features = ["io"] }
 tracing = { workspace = true }
 url = { workspace = true }
+wiremock = { workspace = true }
 
 [dev-dependencies]
+anstream = { workspace = true }
 insta = { workspace = true }
 fastrand = { workspace = true }
 

crates/uv-publish/src/lib.rs

@@ -10,7 +10,7 @@ use fs_err::tokio::File;
 use futures::TryStreamExt;
 use glob::{GlobError, PatternError, glob};
 use itertools::Itertools;
-use reqwest::header::AUTHORIZATION;
+use reqwest::header::{AUTHORIZATION, LOCATION, ToStrError};
 use reqwest::multipart::Part;
 use reqwest::{Body, Response, StatusCode};
 use reqwest_retry::RetryPolicy;
@@ -24,11 +24,11 @@ use tokio_util::io::ReaderStream;
 use tracing::{Level, debug, enabled, trace, warn};
 use url::Url;
 
-use uv_auth::{Credentials, PyxTokenStore};
+use uv_auth::{Credentials, PyxTokenStore, Realm};
 use uv_cache::{Cache, Refresh};
 use uv_client::{
-    BaseClient, MetadataFormat, OwnedArchive, RegistryClientBuilder, RequestBuilder,
+    BaseClient, DEFAULT_MAX_REDIRECTS, MetadataFormat, OwnedArchive, RegistryClientBuilder,
-    RetryParsingError, is_transient_network_error,
+    RequestBuilder, RetryParsingError, is_transient_network_error,
 };
 use uv_configuration::{KeyringProviderType, TrustedPublishing};
 use uv_distribution_filename::{DistFilename, SourceDistExtension, SourceDistFilename};
@@ -37,7 +37,7 @@ use uv_extract::hash::{HashReader, Hasher};
 use uv_fs::{ProgressReader, Simplified};
 use uv_metadata::read_metadata_async_seek;
 use uv_pypi_types::{HashAlgorithm, HashDigest, Metadata23, MetadataError};
-use uv_redacted::DisplaySafeUrl;
+use uv_redacted::{DisplaySafeUrl, DisplaySafeUrlError};
 use uv_warnings::warn_user;
 
 use crate::trusted_publishing::{TrustedPublishingError, TrustedPublishingToken};
@@ -132,11 +132,16 @@ pub enum PublishSendError {
     /// The registry returned a "403 Forbidden".
     #[error("Permission denied (status code {0}): {1}")]
     PermissionDenied(StatusCode, String),
-    /// See inline comment.
+    #[error("Too many redirects, only {0} redirects are allowed")]
-    #[error(
+    TooManyRedirects(u32),
-        "The request was redirected, but redirects are not allowed when publishing, please use the canonical URL: `{0}`"
+    #[error("Redirected URL is not in the same realm. Redirected to: {0}")]
-    )]
+    RedirectRealmMismatch(String),
-    RedirectError(Url),
+    #[error("Request was redirected, but no location header was provided")]
+    RedirectNoLocation,
+    #[error("Request was redirected, but location header is not a UTF-8 string")]
+    RedirectLocationInvalidStr(#[source] ToStrError),
+    #[error("Request was redirected, but location header is not a URL")]
+    RedirectInvalidLocation(#[source] DisplaySafeUrlError),
 }
 
 pub trait Reporter: Send + Sync + 'static {
@@ -470,11 +475,15 @@ pub async fn upload(
     reporter: Arc<impl Reporter>,
 ) -> Result<bool, PublishError> {
     let mut n_past_retries = 0;
+    let mut n_past_redirections = 0;
+    let max_redirects = DEFAULT_MAX_REDIRECTS;
     let start_time = SystemTime::now();
+    let mut current_registry = registry.clone();
+
     loop {
         let (request, idx) = build_upload_request(
             group,
-            registry,
+            &current_registry,
             client,
             credentials,
             form_metadata,
@@ -486,6 +495,60 @@ pub async fn upload(
         let result = request.send().await;
         let response = match result {
             Ok(response) => {
+                // When the user accidentally uses https://test.pypi.org/legacy (no slash) as publish URL, we
+                // get a redirect to https://test.pypi.org/legacy/ (the canonical index URL).
+                // In the above case we get 308, where reqwest or `RedirectClientWithMiddleware` would try
+                // cloning the streaming body, which is not possible.
+                // For https://test.pypi.org/simple (no slash), we get 301, which means we should make a GET request:
+                // https://fetch.spec.whatwg.org/#http-redirect-fetch).
+                // Reqwest doesn't support redirect policies conditional on the HTTP
+                // method (https://github.com/seanmonstar/reqwest/issues/1777#issuecomment-2303386160), so we're
+                // implementing our custom redirection logic.
+                if response.status().is_redirection() {
+                    if n_past_redirections >= max_redirects {
+                        return Err(PublishError::PublishSend(
+                            group.file.clone(),
+                            current_registry.clone().into(),
+                            PublishSendError::TooManyRedirects(n_past_redirections).into(),
+                        ));
+                    }
+                    let location = response
+                        .headers()
+                        .get(LOCATION)
+                        .ok_or_else(|| {
+                            PublishError::PublishSend(
+                                group.file.clone(),
+                                current_registry.clone().into(),
+                                PublishSendError::RedirectNoLocation.into(),
+                            )
+                        })?
+                        .to_str()
+                        .map_err(|err| {
+                            PublishError::PublishSend(
+                                group.file.clone(),
+                                current_registry.clone().into(),
+                                PublishSendError::RedirectLocationInvalidStr(err).into(),
+                            )
+                        })?;
+                    current_registry = DisplaySafeUrl::parse(location).map_err(|err| {
+                        PublishError::PublishSend(
+                            group.file.clone(),
+                            current_registry.clone().into(),
+                            PublishSendError::RedirectInvalidLocation(err).into(),
+                        )
+                    })?;
+                    if Realm::from(&current_registry) != Realm::from(registry) {
+                        return Err(PublishError::PublishSend(
+                            group.file.clone(),
+                            current_registry.clone().into(),
+                            PublishSendError::RedirectRealmMismatch(current_registry.to_string())
+                                .into(),
+                        ));
+                    }
+                    debug!("Redirecting the request to: {}", current_registry);
+                    n_past_redirections += 1;
+                    continue;
+                }
                 reporter.on_upload_complete(idx);
                 response
             }
@@ -498,7 +561,7 @@ pub async fn upload(
                             .unwrap_or_else(|_| Duration::default());
                         warn_user!(
                             "Transient failure while handling response for {}; retrying after {}s...",
-                            registry,
+                            current_registry,
                             duration.as_secs()
                         );
                         tokio::time::sleep(duration).await;
@@ -509,13 +572,13 @@ pub async fn upload(
 
                 return Err(PublishError::PublishSend(
                     group.file.clone(),
-                    registry.clone().into(),
+                    current_registry.clone().into(),
                     PublishSendError::ReqwestMiddleware(err).into(),
                 ));
             }
         };
 
-        return match handle_response(registry, response).await {
+        return match handle_response(&current_registry, response).await {
             Ok(()) => {
                 // Upload successful; for PyPI this can also mean a hash match in a raced upload
                 // (but it doesn't tell us), for other registries it should mean a fresh upload.
@@ -543,7 +606,7 @@ pub async fn upload(
                 }
                 Err(PublishError::PublishSend(
                     group.file.clone(),
-                    registry.clone().into(),
+                    current_registry.clone().into(),
                     err.into(),
                 ))
             }
@@ -1075,17 +1138,6 @@ async fn handle_response(registry: &Url, response: Response) -> Result<(), Publi
     debug!("Response code for {registry}: {status_code}");
     trace!("Response headers for {registry}: {response:?}");
 
-    // When the user accidentally uses https://test.pypi.org/simple (no slash) as publish URL, we
-    // get a redirect to https://test.pypi.org/simple/ (the canonical index URL), while changing the
-    // method to GET (see https://en.wikipedia.org/wiki/Post/Redirect/Get and
-    // https://fetch.spec.whatwg.org/#http-redirect-fetch). The user gets a 200 OK while we actually
-    // didn't upload anything! Reqwest doesn't support redirect policies conditional on the HTTP
-    // method (https://github.com/seanmonstar/reqwest/issues/1777#issuecomment-2303386160), so we're
-    // checking after the fact.
-    if response.url() != registry {
-        return Err(PublishSendError::RedirectError(response.url().clone()));
-    }
-
     if status_code.is_success() {
         if enabled!(Level::TRACE) {
             match response.text().await {
@@ -1142,13 +1194,21 @@ mod tests {
 
     use insta::{allow_duplicates, assert_debug_snapshot, assert_snapshot};
     use itertools::Itertools;
-
+    use thiserror::__private17::AsDynError;
     use uv_auth::Credentials;
-    use uv_client::BaseClientBuilder;
+    use uv_client::{AuthIntegration, BaseClientBuilder, RedirectPolicy};
     use uv_distribution_filename::DistFilename;
     use uv_redacted::DisplaySafeUrl;
 
-    use crate::{FormMetadata, Reporter, UploadDistribution, build_upload_request, group_files};
+    use crate::{
+        FormMetadata, PublishError, Reporter, UploadDistribution, build_upload_request,
+        group_files, upload,
+    };
+    use tokio::sync::Semaphore;
+    use uv_warnings::owo_colors::AnsiColors;
+    use uv_warnings::write_error_chain;
+    use wiremock::matchers::{method, path};
+    use wiremock::{Mock, MockServer, ResponseTemplate};
 
     struct DummyReporter;
 
@@ -1161,6 +1221,44 @@ mod tests {
         fn on_upload_complete(&self, _id: usize) {}
     }
 
+    async fn mock_server_upload(mock_server: &MockServer) -> Result<bool, PublishError> {
+        let raw_filename = "tqdm-4.66.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl";
+        let file = PathBuf::from("../../test/links/").join(raw_filename);
+        let filename = DistFilename::try_from_normalized_filename(raw_filename).unwrap();
+
+        let group = UploadDistribution {
+            file,
+            raw_filename: raw_filename.to_string(),
+            filename,
+            attestations: vec![],
+        };
+
+        let form_metadata = FormMetadata::read_from_file(&group.file, &group.filename)
+            .await
+            .unwrap();
+
+        let client = BaseClientBuilder::default()
+            .redirect(RedirectPolicy::NoRedirect)
+            .retries(0)
+            .auth_integration(AuthIntegration::NoAuthMiddleware)
+            .build();
+
+        let download_concurrency = Arc::new(Semaphore::new(1));
+        let registry = DisplaySafeUrl::parse(&format!("{}/final", &mock_server.uri())).unwrap();
+        upload(
+            &group,
+            &form_metadata,
+            &registry,
+            &client,
+            client.retry_policy(),
+            &Credentials::basic(Some("ferris".to_string()), Some("F3RR!S".to_string())),
+            None,
+            &download_concurrency,
+            Arc::new(DummyReporter),
+        )
+        .await
+    }
+
     #[test]
     fn test_group_files() {
         // Fisher-Yates shuffle.
@@ -1722,4 +1820,88 @@ mod tests {
             "#);
         });
     }
+
+    #[tokio::test]
+    async fn upload_redirect_308() {
+        let mock_server = MockServer::start().await;
+        Mock::given(method("POST"))
+            .and(path("/final"))
+            .respond_with(
+                ResponseTemplate::new(308)
+                    .insert_header("Location", format!("{}/final/", &mock_server.uri())),
+            )
+            .mount(&mock_server)
+            .await;
+        Mock::given(method("POST"))
+            .and(path("/final/"))
+            .respond_with(ResponseTemplate::new(200))
+            .mount(&mock_server)
+            .await;
+
+        assert!(mock_server_upload(&mock_server).await.unwrap());
+    }
+
+    #[tokio::test]
+    async fn upload_infinite_redirects() {
+        let mock_server = MockServer::start().await;
+        Mock::given(method("POST"))
+            .and(path("/final"))
+            .respond_with(
+                ResponseTemplate::new(308)
+                    .insert_header("Location", format!("{}/final/", &mock_server.uri())),
+            )
+            .mount(&mock_server)
+            .await;
+        Mock::given(method("POST"))
+            .and(path("/final/"))
+            .respond_with(
+                ResponseTemplate::new(308)
+                    .insert_header("Location", format!("{}/final", &mock_server.uri())),
+            )
+            .mount(&mock_server)
+            .await;
+
+        let err = mock_server_upload(&mock_server).await.unwrap_err();
+
+        let mut capture = String::new();
+        write_error_chain(err.as_dyn_error(), &mut capture, "error", AnsiColors::Red).unwrap();
+
+        let capture = capture.replace(&mock_server.uri(), "[SERVER]");
+        let capture = anstream::adapter::strip_str(&capture);
+        assert_snapshot!(
+            &capture,
+            @r"
+        error: Failed to publish `../../test/links/tqdm-4.66.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl` to [SERVER]/final
+          Caused by: Too many redirects, only 10 redirects are allowed
+        "
+        );
+    }
+
+    #[tokio::test]
+    async fn upload_redirect_different_realm() {
+        let mock_server = MockServer::start().await;
+        Mock::given(method("POST"))
+            .and(path("/final"))
+            .respond_with(
+                ResponseTemplate::new(308)
+                    .insert_header("Location", "https://different.auth.tld/final/"),
+            )
+            .mount(&mock_server)
+            .await;
+
+        let err = mock_server_upload(&mock_server).await.unwrap_err();
+
+        let mut capture = String::new();
+        write_error_chain(err.as_dyn_error(), &mut capture, "error", AnsiColors::Red).unwrap();
+
+        let capture = capture.replace(&mock_server.uri(), "[SERVER]");
+        let capture = anstream::adapter::strip_str(&capture);
+        assert_snapshot!(
+            &capture,
+            @r"
+        error: Failed to publish `../../test/links/tqdm-4.66.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl` to https://different.auth.tld/final/
+          Caused by: Redirected URL is not in the same realm. Redirected to: https://different.auth.tld/final/
+        "
+        );
+    }
 }

crates/uv/src/commands/publish.rs

@@ -8,7 +8,9 @@ use tokio::sync::Semaphore;
 use tracing::{debug, info, trace};
 use uv_auth::{Credentials, DEFAULT_TOLERANCE_SECS, PyxTokenStore};
 use uv_cache::Cache;
-use uv_client::{AuthIntegration, BaseClient, BaseClientBuilder, RegistryClientBuilder};
+use uv_client::{
+    AuthIntegration, BaseClient, BaseClientBuilder, RedirectPolicy, RegistryClientBuilder,
+};
 use uv_configuration::{KeyringProviderType, TrustedPublishing};
 use uv_distribution_types::{IndexCapabilities, IndexLocations, IndexUrl};
 use uv_publish::{
@@ -123,6 +125,9 @@ pub(crate) async fn publish(
         .keyring(keyring_provider)
         // Don't try cloning the request to make an unauthenticated request first.
         .auth_integration(AuthIntegration::OnlyAuthenticated)
+        // Disable automatic redirect, as the streaming publish request is not cloneable.
+        // Rely on custom redirect logic instead.
+        .redirect(RedirectPolicy::NoRedirect)
         .timeout(environment.upload_http_timeout)
         .build();
     let oidc_client = client_builder