Skip to content

Commit 324c1b4

Browse files
jtfmummjtfmumm
committed
Exit early on auth failures to improve first-index security by default
1 parent fa75458 commit 324c1b4

File tree

14 files changed

+511
-72
lines changed

14 files changed

+511
-72
lines changed

‎Cargo.lock

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/uv-auth/src/index.rs

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,10 @@ pub struct Index {
5959
pub auth_policy: AuthPolicy,
6060
}
6161

62+
// TODO(john): Multiple methods in this struct need to iterate over
63+
// all the indexes in the set. There are probably not many URLs to
64+
// iterate through, but we could use a trie instead of a HashSet here
65+
// for more efficient search.
6266
#[derive(Debug, Default, Clone, Eq, PartialEq)]
6367
pub struct Indexes(FxHashSet<Index>);
6468

@@ -78,26 +82,20 @@ impl Indexes {
7882

7983
/// Get the index URL prefix for a URL if one exists.
8084
pub fn index_url_for(&self, url: &Url) -> Option<&Url> {
81-
// TODO(john): There are probably not many URLs to iterate through,
82-
// but we could use a trie instead of a HashSet here for more
83-
// efficient search.
84-
self.0
85-
.iter()
86-
.find(|index| is_url_prefix(&index.root_url, url))
87-
.map(|index| &index.root_url)
85+
self.find_prefix_index(url).map(|index| &index.root_url)
8886
}
8987

9088
/// Get the [`AuthPolicy`] for a URL.
91-
pub fn policy_for(&self, url: &Url) -> AuthPolicy {
92-
// TODO(john): There are probably not many URLs to iterate through,
93-
// but we could use a trie instead of a HashMap here for more
94-
// efficient search.
95-
for index in &self.0 {
96-
if is_url_prefix(&index.root_url, url) {
97-
return index.auth_policy;
98-
}
99-
}
100-
AuthPolicy::Auto
89+
pub fn auth_policy_for(&self, url: &Url) -> AuthPolicy {
90+
self.find_prefix_index(url)
91+
.map(|index| index.auth_policy)
92+
.unwrap_or(AuthPolicy::Auto)
93+
}
94+
95+
fn find_prefix_index(&self, url: &Url) -> Option<&Index> {
96+
self.0
97+
.iter()
98+
.find(|&index| is_url_prefix(&index.root_url, url))
10199
}
102100
}
103101

‎crates/uv-auth/src/middleware.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,7 @@ impl Middleware for AuthMiddleware {
182182
// to the headers so for display purposes we restore some information
183183
let url = tracing_url(&request, request_credentials.as_ref());
184184
let maybe_index_url = self.indexes.index_url_for(request.url());
185-
let auth_policy = self.indexes.policy_for(request.url());
185+
let auth_policy = self.indexes.auth_policy_for(request.url());
186186
trace!("Handling request for {url} with authentication policy {auth_policy}");
187187

188188
let credentials: Option<Arc<Credentials>> = if matches!(auth_policy, AuthPolicy::Never) {

‎crates/uv-client/src/registry_client.rs

Lines changed: 31 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use async_http_range_reader::AsyncHttpRangeReader;
99
use futures::{FutureExt, StreamExt, TryStreamExt};
1010
use http::HeaderMap;
1111
use itertools::Either;
12-
use reqwest::{Proxy, Response, StatusCode};
12+
use reqwest::{Proxy, Response};
1313
use reqwest_middleware::ClientWithMiddleware;
1414
use rustc_hash::FxHashMap;
1515
use tokio::sync::{Mutex, Semaphore};
@@ -23,7 +23,7 @@ use uv_configuration::{IndexStrategy, TrustedHost};
2323
use uv_distribution_filename::{DistFilename, SourceDistFilename, WheelFilename};
2424
use uv_distribution_types::{
2525
BuiltDist, File, FileLocation, IndexCapabilities, IndexFormat, IndexLocations,
26-
IndexMetadataRef, IndexUrl, IndexUrls, Name,
26+
IndexMetadataRef, IndexStatusCodeDecision, IndexStatusCodeStrategy, IndexUrl, IndexUrls, Name,
2727
};
2828
use uv_metadata::{read_metadata_async_seek, read_metadata_async_stream};
2929
use uv_normalize::PackageName;
@@ -331,8 +331,15 @@ impl RegistryClient {
331331
let _permit = download_concurrency.acquire().await;
332332
match index.format {
333333
IndexFormat::Simple => {
334+
let status_code_strategy =
335+
self.index_urls.status_code_strategy_for(index.url);
334336
if let Some(metadata) = self
335-
.simple_single_index(package_name, index.url, capabilities)
337+
.simple_single_index(
338+
package_name,
339+
index.url,
340+
capabilities,
341+
&status_code_strategy,
342+
)
336343
.await?
337344
{
338345
results.push((index.url, MetadataFormat::Simple(metadata)));
@@ -357,8 +364,15 @@ impl RegistryClient {
357364
let _permit = download_concurrency.acquire().await;
358365
match index.format {
359366
IndexFormat::Simple => {
367+
let status_code_strategy =
368+
self.index_urls.status_code_strategy_for(index.url);
360369
let metadata = self
361-
.simple_single_index(package_name, index.url, capabilities)
370+
.simple_single_index(
371+
package_name,
372+
index.url,
373+
capabilities,
374+
&status_code_strategy,
375+
)
362376
.await?;
363377
Ok((index.url, metadata.map(MetadataFormat::Simple)))
364378
}
@@ -445,6 +459,7 @@ impl RegistryClient {
445459
package_name: &PackageName,
446460
index: &IndexUrl,
447461
capabilities: &IndexCapabilities,
462+
status_code_strategy: &IndexStatusCodeStrategy,
448463
) -> Result<Option<OwnedArchive<SimpleMetadata>>, Error> {
449464
// Format the URL for PyPI.
450465
let mut url = index.url().clone();
@@ -490,18 +505,19 @@ impl RegistryClient {
490505
Ok(metadata) => Ok(Some(metadata)),
491506
Err(err) => match err.into_kind() {
492507
// The package could not be found in the remote index.
493-
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
494-
Some(StatusCode::NOT_FOUND) => Ok(None),
495-
Some(StatusCode::UNAUTHORIZED) => {
496-
capabilities.set_unauthorized(index.clone());
497-
Ok(None)
498-
}
499-
Some(StatusCode::FORBIDDEN) => {
500-
capabilities.set_forbidden(index.clone());
501-
Ok(None)
508+
ErrorKind::WrappedReqwestError(url, err) => {
509+
let decision = if let Some(status) = err.status() {
510+
status_code_strategy.handle_status_code(status, index, capabilities)
511+
} else {
512+
IndexStatusCodeDecision::Stop
513+
};
514+
match decision {
515+
IndexStatusCodeDecision::Continue => Ok(None),
516+
IndexStatusCodeDecision::Stop => {
517+
Err(ErrorKind::WrappedReqwestError(url, err).into())
518+
}
502519
}
503-
_ => Err(ErrorKind::WrappedReqwestError(url, err).into()),
504-
},
520+
}
505521

506522
// The package is unavailable due to a lack of connectivity.
507523
ErrorKind::Offline(_) => Ok(None),

‎crates/uv-distribution-types/Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,10 @@ uv-pypi-types = { workspace = true }
3030
uv-small-str = { workspace = true }
3131

3232
arcstr = { workspace = true }
33+
anyhow = { workspace = true }
3334
bitflags = { workspace = true }
3435
fs-err = { workspace = true }
36+
http = { workspace = true }
3537
itertools = { workspace = true }
3638
jiff = { workspace = true }
3739
owo-colors = { workspace = true }

‎crates/uv-distribution-types/src/index.rs

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,18 @@
11
use std::path::Path;
22
use std::str::FromStr;
33

4+
use http::StatusCode;
5+
use serde::{Deserialize, Deserializer, Serialize};
46
use thiserror::Error;
57
use url::Url;
68

79
use uv_auth::{AuthPolicy, Credentials};
810

911
use crate::index_name::{IndexName, IndexNameError};
1012
use crate::origin::Origin;
11-
use crate::{IndexUrl, IndexUrlError};
13+
use crate::{IndexStatusCodeStrategy, IndexUrl, IndexUrlError};
1214

13-
#[derive(
14-
Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, serde::Serialize, serde::Deserialize,
15-
)]
15+
#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
1616
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
1717
#[serde(rename_all = "kebab-case")]
1818
pub struct Index {
@@ -94,6 +94,17 @@ pub struct Index {
9494
/// ```
9595
#[serde(default)]
9696
pub authenticate: AuthPolicy,
97+
/// Status codes that uv should ignore when deciding whether
98+
/// to continue searching in the next index after a failure.
99+
///
100+
/// ```toml
101+
/// [[tool.uv.index]]
102+
/// name = "my-index"
103+
/// url = "https://<omitted>/simple"
104+
/// ignore-error-codes = [401, 403]
105+
/// ```
106+
#[serde(deserialize_with = "validate_error_codes", default)]
107+
pub ignore_error_codes: Vec<u16>,
97108
}
98109

99110
#[derive(
@@ -131,6 +142,7 @@ impl Index {
131142
format: IndexFormat::Simple,
132143
publish_url: None,
133144
authenticate: AuthPolicy::default(),
145+
ignore_error_codes: Vec::new(),
134146
}
135147
}
136148

@@ -145,6 +157,7 @@ impl Index {
145157
format: IndexFormat::Simple,
146158
publish_url: None,
147159
authenticate: AuthPolicy::default(),
160+
ignore_error_codes: Vec::new(),
148161
}
149162
}
150163

@@ -159,6 +172,7 @@ impl Index {
159172
format: IndexFormat::Flat,
160173
publish_url: None,
161174
authenticate: AuthPolicy::default(),
175+
ignore_error_codes: Vec::new(),
162176
}
163177
}
164178

@@ -214,6 +228,15 @@ impl Index {
214228
}
215229
Ok(self)
216230
}
231+
232+
/// Return the [`IndexStatusCodeStrategy`] for this index.
233+
pub fn status_code_strategy(&self) -> IndexStatusCodeStrategy {
234+
if !self.ignore_error_codes.is_empty() {
235+
IndexStatusCodeStrategy::from_ignored_error_codes(&self.ignore_error_codes)
236+
} else {
237+
IndexStatusCodeStrategy::from_index_url(self.url.url())
238+
}
239+
}
217240
}
218241

219242
impl From<IndexUrl> for Index {
@@ -227,6 +250,7 @@ impl From<IndexUrl> for Index {
227250
format: IndexFormat::Simple,
228251
publish_url: None,
229252
authenticate: AuthPolicy::default(),
253+
ignore_error_codes: Vec::new(),
230254
}
231255
}
232256
}
@@ -249,6 +273,7 @@ impl FromStr for Index {
249273
format: IndexFormat::Simple,
250274
publish_url: None,
251275
authenticate: AuthPolicy::default(),
276+
ignore_error_codes: Vec::new(),
252277
});
253278
}
254279
}
@@ -264,6 +289,7 @@ impl FromStr for Index {
264289
format: IndexFormat::Simple,
265290
publish_url: None,
266291
authenticate: AuthPolicy::default(),
292+
ignore_error_codes: Vec::new(),
267293
})
268294
}
269295
}
@@ -349,6 +375,25 @@ impl<'a> From<&'a IndexUrl> for IndexMetadataRef<'a> {
349375
}
350376
}
351377

378+
/// Validate the provided error codes.
379+
fn validate_error_codes<'de, D>(deserializer: D) -> Result<Vec<u16>, D::Error>
380+
where
381+
D: Deserializer<'de>,
382+
{
383+
let status_codes = Option::<Vec<u16>>::deserialize(deserializer)?;
384+
if let Some(codes) = status_codes {
385+
for code in &codes {
386+
if StatusCode::from_u16(*code).is_err() {
387+
return Err(serde::de::Error::custom(format!(
388+
"{code} is not a valid HTTP status code"
389+
)));
390+
}
391+
}
392+
return Ok(codes);
393+
}
394+
Ok(Vec::new())
395+
}
396+
352397
/// An error that can occur when parsing an [`Index`].
353398
#[derive(Error, Debug)]
354399
pub enum IndexSourceError {

‎crates/uv-distribution-types/src/index_url.rs

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,15 @@ use std::path::Path;
55
use std::str::FromStr;
66
use std::sync::{Arc, LazyLock, RwLock};
77

8+
use http::StatusCode;
89
use itertools::Either;
910
use rustc_hash::{FxHashMap, FxHashSet};
1011
use thiserror::Error;
1112
use url::{ParseError, Url};
1213

1314
use uv_pep508::{split_scheme, Scheme, VerbatimUrl, VerbatimUrlError};
1415

15-
use crate::{Index, Verbatim};
16+
use crate::{Index, IndexStatusCodeStrategy, Verbatim};
1617

1718
static PYPI_URL: LazyLock<Url> = LazyLock::new(|| Url::parse("https://pypi.org/simple").unwrap());
1819

@@ -535,6 +536,16 @@ impl<'a> IndexUrls {
535536
pub fn no_index(&self) -> bool {
536537
self.no_index
537538
}
539+
540+
/// Return the [`IndexStatusCodeStrategy`] for an [`IndexUrl`].
541+
pub fn status_code_strategy_for(&self, url: &IndexUrl) -> IndexStatusCodeStrategy {
542+
for index in &self.indexes {
543+
if index.url() == url {
544+
return index.status_code_strategy();
545+
}
546+
}
547+
IndexStatusCodeStrategy::Default
548+
}
538549
}
539550

540551
bitflags::bitflags! {
@@ -578,6 +589,18 @@ impl IndexCapabilities {
578589
.insert(Flags::NO_RANGE_REQUESTS);
579590
}
580591

592+
pub fn set_by_status_code(&self, status_code: StatusCode, index_url: &IndexUrl) {
593+
match status_code {
594+
StatusCode::UNAUTHORIZED => {
595+
self.set_unauthorized(index_url.clone());
596+
}
597+
StatusCode::FORBIDDEN => {
598+
self.set_forbidden(index_url.clone());
599+
}
600+
_ => {}
601+
}
602+
}
603+
581604
/// Returns `true` if the given [`IndexUrl`] returns a `401 Unauthorized` status code.
582605
pub fn unauthorized(&self, index_url: &IndexUrl) -> bool {
583606
self.0

‎crates/uv-distribution-types/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ pub use crate::requirement::*;
7575
pub use crate::resolution::*;
7676
pub use crate::resolved::*;
7777
pub use crate::specified_requirement::*;
78+
pub use crate::status_code_strategy::*;
7879
pub use crate::traits::*;
7980

8081
mod annotation;
@@ -101,6 +102,7 @@ mod requirement;
101102
mod resolution;
102103
mod resolved;
103104
mod specified_requirement;
105+
mod status_code_strategy;
104106
mod traits;
105107

106108
#[derive(Debug, Clone)]

0 commit comments

Comments
 (0)