From d9a5f5ca1c15d2658891f3ee06a91d55563739b0 Mon Sep 17 00:00:00 2001 From: konsti Date: Sat, 21 Sep 2024 16:09:14 +0200 Subject: [PATCH] Add `only_authenticated` option to the client (#7545) --- crates/uv-auth/src/middleware.rs | 82 +++++++++++++++++++---------- crates/uv-client/src/base_client.rs | 38 ++++++++----- 2 files changed, 81 insertions(+), 39 deletions(-) diff --git a/crates/uv-auth/src/middleware.rs b/crates/uv-auth/src/middleware.rs index 1c938d5cba8e..cd6b177495f1 100644 --- a/crates/uv-auth/src/middleware.rs +++ b/crates/uv-auth/src/middleware.rs @@ -8,7 +8,7 @@ use crate::{ realm::Realm, CredentialsCache, KeyringProvider, CREDENTIALS_CACHE, }; -use anyhow::anyhow; +use anyhow::{anyhow, format_err}; use netrc::Netrc; use reqwest::{Request, Response}; use reqwest_middleware::{Error, Middleware, Next}; @@ -22,6 +22,9 @@ pub struct AuthMiddleware { netrc: Option, keyring: Option, cache: Option, + /// We know that the endpoint needs authentication, so we don't try to send an unauthenticated + /// request, avoiding cloning an uncloneable request. + only_authenticated: bool, } impl AuthMiddleware { @@ -30,6 +33,7 @@ impl AuthMiddleware { netrc: Netrc::new().ok(), keyring: None, cache: None, + only_authenticated: false, } } @@ -56,6 +60,14 @@ impl AuthMiddleware { self } + /// We know that the endpoint needs authentication, so we don't try to send an unauthenticated + /// request, avoiding cloning an uncloneable request. + #[must_use] + pub fn with_only_authenticated(mut self, only_authenticated: bool) -> Self { + self.only_authenticated = only_authenticated; + self + } + /// Get the configured authentication store. /// /// If not set, the global store is used. @@ -198,32 +210,42 @@ impl Middleware for AuthMiddleware { .as_ref() .is_some_and(|credentials| credentials.username().is_some()); - // Otherwise, attempt an anonymous request - trace!("Attempting unauthenticated request for {url}"); - - // - // Clone the request so we can retry it on authentication failure - let mut retry_request = request.try_clone().ok_or_else(|| { - Error::Middleware(anyhow!( - "Request object is not cloneable. Are you passing a streaming body?".to_string() - )) - })?; - - let response = next.clone().run(request, extensions).await?; - - // If we don't fail with authorization related codes, return the response - if !matches!( - response.status(), - StatusCode::FORBIDDEN | StatusCode::NOT_FOUND | StatusCode::UNAUTHORIZED - ) { - return Ok(response); - } + let (mut retry_request, response) = if self.only_authenticated { + // For endpoints where we require the user to provide credentials, we don't try the + // unauthenticated request first. + trace!("Checking for credentials for {url}"); + (request, None) + } else { + // Otherwise, attempt an anonymous request + trace!("Attempting unauthenticated request for {url}"); + + // + // Clone the request so we can retry it on authentication failure + let retry_request = request.try_clone().ok_or_else(|| { + Error::Middleware(anyhow!( + "Request object is not cloneable. Are you passing a streaming body?" + .to_string() + )) + })?; + + let response = next.clone().run(request, extensions).await?; + + // If we don't fail with authorization related codes, return the response + if !matches!( + response.status(), + StatusCode::FORBIDDEN | StatusCode::NOT_FOUND | StatusCode::UNAUTHORIZED + ) { + return Ok(response); + } - // Otherwise, search for credentials - trace!( - "Request for {url} failed with {}, checking for credentials", - response.status() - ); + // Otherwise, search for credentials + trace!( + "Request for {url} failed with {}, checking for credentials", + response.status() + ); + + (retry_request, Some(response)) + }; // Check in the cache first let credentials = self.cache().get_realm( @@ -265,7 +287,13 @@ impl Middleware for AuthMiddleware { } } - Ok(response) + if let Some(response) = response { + Ok(response) + } else { + Err(Error::Middleware(format_err!( + "Missing credentials for {url}" + ))) + } } } diff --git a/crates/uv-client/src/base_client.rs b/crates/uv-client/src/base_client.rs index 189d960ce7ea..c9a2b5d596a5 100644 --- a/crates/uv-client/src/base_client.rs +++ b/crates/uv-client/src/base_client.rs @@ -36,6 +36,7 @@ pub struct BaseClientBuilder<'a> { client: Option, markers: Option<&'a MarkerEnvironment>, platform: Option<&'a Platform>, + only_authenticated: bool, } impl Default for BaseClientBuilder<'_> { @@ -55,6 +56,7 @@ impl BaseClientBuilder<'_> { client: None, markers: None, platform: None, + only_authenticated: false, } } } @@ -108,6 +110,12 @@ impl<'a> BaseClientBuilder<'a> { self } + #[must_use] + pub fn only_authenticated(mut self, only_authenticated: bool) -> Self { + self.only_authenticated = only_authenticated; + self + } + pub fn is_offline(&self) -> bool { matches!(self.connectivity, Connectivity::Offline) } @@ -230,20 +238,26 @@ impl<'a> BaseClientBuilder<'a> { fn apply_middleware(&self, client: Client) -> ClientWithMiddleware { match self.connectivity { Connectivity::Online => { - let client = reqwest_middleware::ClientBuilder::new(client); - - // Initialize the retry strategy. - let retry_policy = - ExponentialBackoff::builder().build_with_max_retries(self.retries); - let retry_strategy = RetryTransientMiddleware::new_with_policy_and_strategy( - retry_policy, - UvRetryableStrategy, - ); - let client = client.with(retry_strategy); + let mut client = reqwest_middleware::ClientBuilder::new(client); + + // Avoid uncloneable errors with a streaming body during publish. + if self.retries > 0 { + // Initialize the retry strategy. + let retry_policy = + ExponentialBackoff::builder().build_with_max_retries(self.retries); + let retry_strategy = RetryTransientMiddleware::new_with_policy_and_strategy( + retry_policy, + UvRetryableStrategy, + ); + client = client.with(retry_strategy); + } // Initialize the authentication middleware to set headers. - let client = - client.with(AuthMiddleware::new().with_keyring(self.keyring.to_provider())); + client = client.with( + AuthMiddleware::new() + .with_keyring(self.keyring.to_provider()) + .with_only_authenticated(self.only_authenticated), + ); client.build() }