1use std::borrow::Cow;
2use std::fmt::{Display, Formatter};
3use std::ops::Deref;
4use std::path::Path;
5use std::str::FromStr;
6use std::sync::{Arc, LazyLock, RwLock};
7
8use http::StatusCode;
9use itertools::Either;
10use rustc_hash::{FxHashMap, FxHashSet};
11use thiserror::Error;
12use url::{ParseError, Url};
13use uv_auth::RealmRef;
14use uv_cache_key::CanonicalUrl;
15use uv_pep508::{Scheme, VerbatimUrl, VerbatimUrlError, split_scheme};
16use uv_pypi_types::HashAlgorithm;
17use uv_redacted::DisplaySafeUrl;
18use uv_warnings::warn_user;
19
20use crate::{ExcludeNewerOverride, Index, IndexStatusCodeStrategy, Verbatim};
21
22pub static PYPI_URL: LazyLock<DisplaySafeUrl> =
23 LazyLock::new(|| DisplaySafeUrl::parse("https://pypi.org/simple").unwrap());
24
25static DEFAULT_INDEX: LazyLock<Index> = LazyLock::new(|| {
26 Index::from_index_url(IndexUrl::Pypi(Arc::new(VerbatimUrl::from_url(
27 PYPI_URL.clone(),
28 ))))
29});
30
31#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
33pub enum IndexUrl {
34 Pypi(Arc<VerbatimUrl>),
35 Url(Arc<VerbatimUrl>),
36 Path(Arc<VerbatimUrl>),
37}
38
39impl IndexUrl {
40 pub fn parse(path: &str, root_dir: Option<&Path>) -> Result<Self, IndexUrlError> {
45 let url = VerbatimUrl::from_url_or_path(path, root_dir)?;
46 Ok(Self::from(url))
47 }
48
49 pub fn root(&self) -> Option<DisplaySafeUrl> {
54 let mut segments = self.url().path_segments()?;
55 let last = match segments.next_back()? {
56 "" => segments.next_back()?,
58 segment => segment,
59 };
60
61 if !(last.eq_ignore_ascii_case("simple") || last.eq_ignore_ascii_case("+simple")) {
63 return None;
64 }
65
66 let mut url = self.url().clone();
67 url.path_segments_mut().ok()?.pop_if_empty().pop();
68 Some(url)
69 }
70}
71
72#[cfg(feature = "schemars")]
73impl schemars::JsonSchema for IndexUrl {
74 fn schema_name() -> Cow<'static, str> {
75 Cow::Borrowed("IndexUrl")
76 }
77
78 fn json_schema(_generator: &mut schemars::generate::SchemaGenerator) -> schemars::Schema {
79 schemars::json_schema!({
80 "type": "string",
81 "description": "The URL of an index to use for fetching packages (e.g., `https://pypi.org/simple`), or a local path."
82 })
83 }
84}
85
86impl IndexUrl {
87 #[inline]
88 fn inner(&self) -> &VerbatimUrl {
89 match self {
90 Self::Pypi(url) | Self::Url(url) | Self::Path(url) => url,
91 }
92 }
93
94 pub fn url(&self) -> &DisplaySafeUrl {
96 self.inner().raw()
97 }
98
99 pub fn into_url(self) -> DisplaySafeUrl {
101 match self {
102 Self::Pypi(url) | Self::Url(url) | Self::Path(url) => {
103 Arc::unwrap_or_clone(url).into_url()
104 }
105 }
106 }
107
108 pub fn without_credentials(&self) -> Cow<'_, DisplaySafeUrl> {
110 let url = self.url();
111 if url.username().is_empty() && url.password().is_none() {
112 Cow::Borrowed(url)
113 } else {
114 let mut url = url.clone();
115 let _ = url.set_username("");
116 let _ = url.set_password(None);
117 Cow::Owned(url)
118 }
119 }
120
121 pub fn warn_on_disambiguated_relative_path(&self) {
126 let Self::Path(verbatim_url) = &self else {
127 return;
128 };
129
130 if let Some(path) = verbatim_url.given()
131 && !is_disambiguated_path(path)
132 {
133 if cfg!(windows) {
134 warn_user!(
135 "Relative paths passed to `--index` or `--default-index` should be disambiguated from index names (use `.\\{path}` or `./{path}`). Support for ambiguous values will be removed in the future"
136 );
137 } else {
138 warn_user!(
139 "Relative paths passed to `--index` or `--default-index` should be disambiguated from index names (use `./{path}`). Support for ambiguous values will be removed in the future"
140 );
141 }
142 }
143 }
144}
145
146impl Display for IndexUrl {
147 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
148 Display::fmt(self.inner(), f)
149 }
150}
151
152impl Verbatim for IndexUrl {
153 fn verbatim(&self) -> Cow<'_, str> {
154 self.inner().verbatim()
155 }
156}
157
158fn is_disambiguated_path(path: &str) -> bool {
164 if cfg!(windows) {
165 if path.starts_with(".\\") || path.starts_with("..\\") || path.starts_with('/') {
166 return true;
167 }
168 }
169 if path.starts_with("./") || path.starts_with("../") || Path::new(path).is_absolute() {
170 return true;
171 }
172 if let Some((scheme, _)) = split_scheme(path) {
174 return Scheme::parse(scheme).is_some();
175 }
176 false
178}
179
180#[derive(Error, Debug)]
182pub enum IndexUrlError {
183 #[error(transparent)]
184 Io(#[from] std::io::Error),
185 #[error(transparent)]
186 Url(#[from] ParseError),
187 #[error(transparent)]
188 VerbatimUrl(#[from] VerbatimUrlError),
189}
190
191impl FromStr for IndexUrl {
192 type Err = IndexUrlError;
193
194 fn from_str(s: &str) -> Result<Self, Self::Err> {
195 Self::parse(s, None)
196 }
197}
198
199impl serde::ser::Serialize for IndexUrl {
200 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
201 where
202 S: serde::ser::Serializer,
203 {
204 self.inner().without_credentials().serialize(serializer)
205 }
206}
207
208impl<'de> serde::de::Deserialize<'de> for IndexUrl {
209 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
210 where
211 D: serde::de::Deserializer<'de>,
212 {
213 struct Visitor;
214
215 impl serde::de::Visitor<'_> for Visitor {
216 type Value = IndexUrl;
217
218 fn expecting(&self, f: &mut Formatter) -> std::fmt::Result {
219 f.write_str("a string")
220 }
221
222 fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
223 IndexUrl::from_str(v).map_err(serde::de::Error::custom)
224 }
225 }
226
227 deserializer.deserialize_str(Visitor)
228 }
229}
230
231impl From<VerbatimUrl> for IndexUrl {
232 fn from(url: VerbatimUrl) -> Self {
233 if url.scheme() == "file" {
234 Self::Path(Arc::new(url))
235 } else if *url.raw() == *PYPI_URL {
236 Self::Pypi(Arc::new(url))
237 } else {
238 Self::Url(Arc::new(url))
239 }
240 }
241}
242
243impl From<IndexUrl> for DisplaySafeUrl {
244 fn from(index: IndexUrl) -> Self {
245 index.into_url()
246 }
247}
248
249impl Deref for IndexUrl {
250 type Target = Url;
251
252 fn deref(&self) -> &Self::Target {
253 self.inner()
254 }
255}
256
257#[derive(Default, Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
262#[serde(rename_all = "kebab-case", deny_unknown_fields)]
263pub struct IndexLocations {
264 indexes: Vec<Index>,
265 flat_index: Vec<Index>,
266 no_index: bool,
267}
268
269impl IndexLocations {
270 pub fn new(indexes: Vec<Index>, flat_index: Vec<Index>, no_index: bool) -> Self {
272 Self {
273 indexes,
274 flat_index,
275 no_index,
276 }
277 }
278
279 #[must_use]
286 pub fn combine(self, indexes: Vec<Index>, flat_index: Vec<Index>, no_index: bool) -> Self {
287 Self {
288 indexes: self.indexes.into_iter().chain(indexes).collect(),
289 flat_index: self.flat_index.into_iter().chain(flat_index).collect(),
290 no_index: self.no_index || no_index,
291 }
292 }
293
294 pub fn is_none(&self) -> bool {
297 *self == Self::default()
298 }
299}
300
301fn is_same_index(a: &IndexUrl, b: &IndexUrl) -> bool {
303 RealmRef::from(&**b.url()) == RealmRef::from(&**a.url())
304 && CanonicalUrl::new(a.url().clone()) == CanonicalUrl::new(b.url().clone())
305}
306
307impl<'a> IndexLocations {
308 pub fn default_index(&'a self) -> Option<&'a Index> {
314 if self.no_index {
315 None
316 } else {
317 let mut seen = FxHashSet::default();
318 self.indexes
319 .iter()
320 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
321 .find(|index| index.default)
322 .or_else(|| Some(&DEFAULT_INDEX))
323 }
324 }
325
326 pub fn implicit_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
330 if self.no_index {
331 Either::Left(std::iter::empty())
332 } else {
333 let mut seen = FxHashSet::default();
334 Either::Right(
335 self.indexes
336 .iter()
337 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
338 .filter(|index| !index.default && !index.explicit),
339 )
340 }
341 }
342
343 pub fn explicit_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
347 if self.no_index {
348 Either::Left(std::iter::empty())
349 } else {
350 let mut seen = FxHashSet::default();
351 Either::Right(
352 self.indexes
353 .iter()
354 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
355 .filter(|index| index.explicit),
356 )
357 }
358 }
359
360 pub fn indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
369 self.implicit_indexes()
370 .chain(self.default_index())
371 .filter(|index| !index.explicit)
372 }
373
374 pub fn fetch_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
378 let mut seen = FxHashSet::default();
379 self.indexes()
380 .filter(move |index| seen.insert(index.raw_url()))
381 }
382
383 pub fn simple_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
387 if self.no_index {
388 Either::Left(std::iter::empty())
389 } else {
390 let mut seen = FxHashSet::default();
391 Either::Right(
392 self.indexes
393 .iter()
394 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name))),
395 )
396 }
397 }
398
399 pub fn flat_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
401 self.flat_index.iter()
402 }
403
404 pub fn no_index(&self) -> bool {
406 self.no_index
407 }
408
409 pub fn allowed_indexes(&'a self) -> Vec<&'a Index> {
416 if self.no_index {
417 self.flat_index.iter().rev().collect()
418 } else {
419 let mut indexes = vec![];
420
421 let mut seen = FxHashSet::default();
422 let mut default = false;
423 for index in {
424 self.indexes
425 .iter()
426 .chain(self.flat_index.iter())
427 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
428 } {
429 if index.default {
430 if default {
431 continue;
432 }
433 default = true;
434 }
435 indexes.push(index);
436 }
437 if !default {
438 indexes.push(&*DEFAULT_INDEX);
439 }
440
441 indexes.reverse();
442 indexes
443 }
444 }
445
446 pub fn known_indexes(&'a self) -> impl Iterator<Item = &'a Index> {
455 if self.no_index {
456 Either::Left(self.flat_index.iter().rev())
457 } else {
458 Either::Right(
459 std::iter::once(&*DEFAULT_INDEX)
460 .chain(self.flat_index.iter().rev())
461 .chain(self.indexes.iter().rev()),
462 )
463 }
464 }
465
466 pub fn defined_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
476 if self.no_index {
477 return Either::Left(std::iter::empty());
478 }
479
480 let mut seen = FxHashSet::default();
481 let (non_default, default) = self
482 .indexes
483 .iter()
484 .filter(move |index| {
485 if let Some(name) = &index.name {
486 seen.insert(name)
487 } else {
488 true
489 }
490 })
491 .partition::<Vec<_>, _>(|index| !index.default);
492
493 Either::Right(non_default.into_iter().chain(default))
494 }
495
496 fn index_for_url(&self, url: &IndexUrl) -> Option<&Index> {
498 self.indexes
499 .iter()
500 .find(|index| is_same_index(index.url(), url))
501 }
502
503 pub fn status_code_strategy_for(&self, url: &IndexUrl) -> IndexStatusCodeStrategy {
505 self.index_for_url(url).map_or(
506 IndexStatusCodeStrategy::Default,
507 Index::status_code_strategy,
508 )
509 }
510
511 pub fn ignores_error_code_for(&self, url: &IndexUrl, status_code: StatusCode) -> bool {
513 self.index_for_url(url)
514 .is_some_and(|index| index.ignores_error_code(status_code))
515 }
516
517 pub fn simple_api_cache_control_for(&self, url: &IndexUrl) -> Option<http::HeaderValue> {
519 self.index_for_url(url)
520 .and_then(Index::simple_api_cache_control)
521 }
522
523 pub fn artifact_cache_control_for(&self, url: &IndexUrl) -> Option<http::HeaderValue> {
525 self.index_for_url(url)
526 .and_then(Index::artifact_cache_control)
527 }
528
529 pub fn hash_algorithm_for(&self, url: &IndexUrl) -> Option<HashAlgorithm> {
531 self.index_for_url(url)
532 .and_then(|index| index.hash_algorithm.map(HashAlgorithm::from))
533 }
534
535 pub fn exclude_newer_for(&self, url: &IndexUrl) -> Option<&ExcludeNewerOverride> {
537 self.index_for_url(url).and_then(Index::exclude_newer)
538 }
539}
540
541impl From<&IndexLocations> for uv_auth::Indexes {
542 fn from(index_locations: &IndexLocations) -> Self {
543 Self::from_indexes(index_locations.allowed_indexes().into_iter().map(|index| {
544 let mut url = index.url().url().clone();
545 url.set_username("").ok();
546 url.set_password(None).ok();
547 let mut root_url = index.url().root().unwrap_or_else(|| url.clone());
548 root_url.set_username("").ok();
549 root_url.set_password(None).ok();
550 uv_auth::Index {
551 url,
552 root_url,
553 auth_policy: index.authenticate,
554 }
555 }))
556 }
557}
558
559bitflags::bitflags! {
560 #[derive(Debug, Copy, Clone)]
561 struct Flags: u8 {
562 const NO_RANGE_REQUESTS = 1;
564 const UNAUTHORIZED = 1 << 2;
566 const FORBIDDEN = 1 << 1;
568 }
569}
570
571#[derive(Debug, Default, Clone)]
577pub struct IndexCapabilities(Arc<RwLock<FxHashMap<IndexUrl, Flags>>>);
578
579impl IndexCapabilities {
580 pub fn supports_range_requests(&self, index_url: &IndexUrl) -> bool {
582 !self
583 .0
584 .read()
585 .unwrap()
586 .get(index_url)
587 .is_some_and(|flags| flags.intersects(Flags::NO_RANGE_REQUESTS))
588 }
589
590 pub fn set_no_range_requests(&self, index_url: IndexUrl) {
592 self.0
593 .write()
594 .unwrap()
595 .entry(index_url)
596 .or_insert(Flags::empty())
597 .insert(Flags::NO_RANGE_REQUESTS);
598 }
599
600 pub fn unauthorized(&self, index_url: &IndexUrl) -> bool {
602 self.0
603 .read()
604 .unwrap()
605 .get(index_url)
606 .is_some_and(|flags| flags.intersects(Flags::UNAUTHORIZED))
607 }
608
609 pub(crate) fn set_unauthorized(&self, index_url: IndexUrl) {
611 self.0
612 .write()
613 .unwrap()
614 .entry(index_url)
615 .or_insert(Flags::empty())
616 .insert(Flags::UNAUTHORIZED);
617 }
618
619 pub fn forbidden(&self, index_url: &IndexUrl) -> bool {
621 self.0
622 .read()
623 .unwrap()
624 .get(index_url)
625 .is_some_and(|flags| flags.intersects(Flags::FORBIDDEN))
626 }
627
628 pub(crate) fn set_forbidden(&self, index_url: IndexUrl) {
630 self.0
631 .write()
632 .unwrap()
633 .entry(index_url)
634 .or_insert(Flags::empty())
635 .insert(Flags::FORBIDDEN);
636 }
637}
638
639#[cfg(test)]
640mod tests {
641 use super::*;
642 use crate::{IndexCacheControl, IndexFormat, IndexName};
643 use http::HeaderValue;
644
645 #[test]
646 fn test_index_url_parse_valid_paths() {
647 assert!(is_disambiguated_path("/absolute/path"));
649 assert!(is_disambiguated_path("./relative/path"));
651 assert!(is_disambiguated_path("../../relative/path"));
652 if cfg!(windows) {
653 assert!(is_disambiguated_path("C:/absolute/path"));
655 assert!(is_disambiguated_path(".\\relative\\path"));
657 assert!(is_disambiguated_path("..\\..\\relative\\path"));
658 }
659 }
660
661 #[test]
662 fn test_index_url_parse_ambiguous_paths() {
663 assert!(!is_disambiguated_path("index"));
665 assert!(!is_disambiguated_path("relative/path"));
667 }
668
669 #[test]
670 fn test_index_url_parse_with_schemes() {
671 assert!(is_disambiguated_path("file:///absolute/path"));
672 assert!(is_disambiguated_path("https://registry.com/simple/"));
673 assert!(is_disambiguated_path(
674 "git+https://github.com/example/repo.git"
675 ));
676 }
677
678 #[test]
679 fn fetch_indexes_deduplicates_raw_urls() {
680 let url = IndexUrl::from_str("https://index.example.com/simple").unwrap();
681 let mut first = Index::from(url.clone());
682 first.name = Some(IndexName::from_str("first").unwrap());
683 let mut second = Index::from(url);
684 second.name = Some(IndexName::from_str("second").unwrap());
685 second.default = true;
686 let locations = IndexLocations::new(vec![first, second], Vec::new(), false);
687
688 assert_eq!(locations.indexes().count(), 2);
689 assert_eq!(locations.fetch_indexes().count(), 1);
690 }
691
692 #[test]
693 fn test_cache_control_lookup() {
694 use std::str::FromStr;
695
696 use crate::IndexFormat;
697 use crate::index_name::IndexName;
698
699 let indexes = vec![
700 Index {
701 name: Some(IndexName::from_str("index1").unwrap()),
702 url: IndexUrl::from_str("https://index1.example.com/simple").unwrap(),
703 cache_control: Some(crate::IndexCacheControl {
704 api: Some(HeaderValue::from_static("max-age=300")),
705 files: Some(HeaderValue::from_static("max-age=1800")),
706 }),
707 explicit: false,
708 default: false,
709 origin: None,
710 format: IndexFormat::Simple,
711 publish_url: None,
712 authenticate: uv_auth::AuthPolicy::default(),
713 ignore_error_codes: None,
714 hash_algorithm: None,
715 exclude_newer: None,
716 },
717 Index {
718 name: Some(IndexName::from_str("index2").unwrap()),
719 url: IndexUrl::from_str("https://index2.example.com/simple").unwrap(),
720 cache_control: None,
721 explicit: false,
722 default: false,
723 origin: None,
724 format: IndexFormat::Simple,
725 publish_url: None,
726 authenticate: uv_auth::AuthPolicy::default(),
727 ignore_error_codes: None,
728 hash_algorithm: None,
729 exclude_newer: None,
730 },
731 ];
732
733 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
734
735 let url1 = IndexUrl::from_str("https://index1.example.com/simple").unwrap();
736 assert_eq!(
737 index_locations.simple_api_cache_control_for(&url1),
738 Some(HeaderValue::from_static("max-age=300"))
739 );
740 assert_eq!(
741 index_locations.artifact_cache_control_for(&url1),
742 Some(HeaderValue::from_static("max-age=1800"))
743 );
744
745 let url2 = IndexUrl::from_str("https://index2.example.com/simple").unwrap();
746 assert_eq!(index_locations.simple_api_cache_control_for(&url2), None);
747 assert_eq!(index_locations.artifact_cache_control_for(&url2), None);
748
749 let url3 = IndexUrl::from_str("https://index3.example.com/simple").unwrap();
750 assert_eq!(index_locations.simple_api_cache_control_for(&url3), None);
751 assert_eq!(index_locations.artifact_cache_control_for(&url3), None);
752 }
753
754 #[test]
755 fn test_pytorch_default_cache_control() {
756 let indexes = vec![Index {
758 name: Some(IndexName::from_str("pytorch").unwrap()),
759 url: IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap(),
760 cache_control: None, explicit: false,
762 default: false,
763 origin: None,
764 format: IndexFormat::Simple,
765 publish_url: None,
766 authenticate: uv_auth::AuthPolicy::default(),
767 ignore_error_codes: None,
768 hash_algorithm: None,
769 exclude_newer: None,
770 }];
771
772 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
773
774 let pytorch_url = IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap();
775
776 assert_eq!(
777 index_locations.simple_api_cache_control_for(&pytorch_url),
778 None
779 );
780 assert_eq!(
781 index_locations.artifact_cache_control_for(&pytorch_url),
782 Some(HeaderValue::from_static(
783 "max-age=365000000, immutable, public",
784 ))
785 );
786 }
787
788 #[test]
789 fn test_pytorch_user_override_cache_control() {
790 let indexes = vec![Index {
792 name: Some(IndexName::from_str("pytorch").unwrap()),
793 url: IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap(),
794 cache_control: Some(IndexCacheControl {
795 api: Some(HeaderValue::from_static("no-cache")),
796 files: Some(HeaderValue::from_static("max-age=3600")),
797 }),
798 explicit: false,
799 default: false,
800 origin: None,
801 format: IndexFormat::Simple,
802 publish_url: None,
803 authenticate: uv_auth::AuthPolicy::default(),
804 ignore_error_codes: None,
805 hash_algorithm: None,
806 exclude_newer: None,
807 }];
808
809 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
810
811 let pytorch_url = IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap();
812
813 assert_eq!(
814 index_locations.simple_api_cache_control_for(&pytorch_url),
815 Some(HeaderValue::from_static("no-cache"))
816 );
817 assert_eq!(
818 index_locations.artifact_cache_control_for(&pytorch_url),
819 Some(HeaderValue::from_static("max-age=3600"))
820 );
821 }
822
823 #[test]
824 fn test_nvidia_default_cache_control() {
825 let indexes = vec![Index {
827 name: Some(IndexName::from_str("nvidia").unwrap()),
828 url: IndexUrl::from_str("https://pypi.nvidia.com").unwrap(),
829 cache_control: None, explicit: false,
831 default: false,
832 origin: None,
833 format: IndexFormat::Simple,
834 publish_url: None,
835 authenticate: uv_auth::AuthPolicy::default(),
836 ignore_error_codes: None,
837 hash_algorithm: None,
838 exclude_newer: None,
839 }];
840
841 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
842
843 let nvidia_url = IndexUrl::from_str("https://pypi.nvidia.com").unwrap();
844
845 assert_eq!(
846 index_locations.simple_api_cache_control_for(&nvidia_url),
847 None
848 );
849 assert_eq!(
850 index_locations.artifact_cache_control_for(&nvidia_url),
851 Some(HeaderValue::from_static(
852 "max-age=365000000, immutable, public",
853 ))
854 );
855 }
856}