1use std::path::Path;
2use std::str::FromStr;
3
4use http::HeaderValue;
5use serde::{Deserialize, Serialize, Serializer};
6use thiserror::Error;
7use url::Url;
8
9use uv_auth::{AuthPolicy, Credentials};
10use uv_redacted::DisplaySafeUrl;
11use uv_small_str::SmallString;
12
13use crate::index_name::{IndexName, IndexNameError};
14use crate::origin::Origin;
15use crate::{IndexStatusCodeStrategy, IndexUrl, IndexUrlError, SerializableStatusCode};
16
17#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Default)]
19#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
20pub struct IndexCacheControl {
21 #[cfg_attr(feature = "schemars", schemars(with = "Option<String>"))]
23 pub api: Option<HeaderValue>,
24 #[cfg_attr(feature = "schemars", schemars(with = "Option<String>"))]
26 pub files: Option<HeaderValue>,
27}
28
29impl IndexCacheControl {
30 pub fn simple_api_cache_control(_url: &Url) -> Option<HeaderValue> {
32 None
33 }
34
35 pub fn artifact_cache_control(url: &Url) -> Option<HeaderValue> {
37 let dominated_by_pytorch_or_nvidia = url.host_str().is_some_and(|host| {
38 host.eq_ignore_ascii_case("download.pytorch.org")
39 || host.eq_ignore_ascii_case("pypi.nvidia.com")
40 });
41 if dominated_by_pytorch_or_nvidia {
42 Some(HeaderValue::from_static(
50 "max-age=365000000, immutable, public",
51 ))
52 } else {
53 None
54 }
55 }
56}
57
58#[derive(Serialize)]
59#[serde(rename_all = "kebab-case")]
60struct IndexCacheControlRef<'a> {
61 #[serde(skip_serializing_if = "Option::is_none")]
62 api: Option<&'a str>,
63 #[serde(skip_serializing_if = "Option::is_none")]
64 files: Option<&'a str>,
65}
66
67impl Serialize for IndexCacheControl {
68 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
69 where
70 S: Serializer,
71 {
72 IndexCacheControlRef {
73 api: self.api.as_ref().map(|api| {
74 api.to_str()
75 .expect("cache-control.api is always parsed from a string")
76 }),
77 files: self.files.as_ref().map(|files| {
78 files
79 .to_str()
80 .expect("cache-control.files is always parsed from a string")
81 }),
82 }
83 .serialize(serializer)
84 }
85}
86
87#[derive(Debug, Clone, Deserialize)]
88#[serde(rename_all = "kebab-case")]
89struct IndexCacheControlWire {
90 api: Option<SmallString>,
91 files: Option<SmallString>,
92}
93
94impl<'de> Deserialize<'de> for IndexCacheControl {
95 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
96 where
97 D: serde::Deserializer<'de>,
98 {
99 let wire = IndexCacheControlWire::deserialize(deserializer)?;
100
101 let api = wire
102 .api
103 .map(|api| {
104 HeaderValue::from_str(api.as_ref()).map_err(|_| {
105 serde::de::Error::custom(
106 "`cache-control.api` must be a valid HTTP header value",
107 )
108 })
109 })
110 .transpose()?;
111 let files = wire
112 .files
113 .map(|files| {
114 HeaderValue::from_str(files.as_ref()).map_err(|_| {
115 serde::de::Error::custom(
116 "`cache-control.files` must be a valid HTTP header value",
117 )
118 })
119 })
120 .transpose()?;
121
122 Ok(Self { api, files })
123 }
124}
125
126#[derive(Debug, Clone, Serialize)]
127#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
128#[serde(rename_all = "kebab-case")]
129pub struct Index {
130 pub name: Option<IndexName>,
144 pub url: IndexUrl,
148 #[serde(default)]
163 pub explicit: bool,
164 #[serde(default)]
174 pub default: bool,
175 #[serde(skip)]
177 pub origin: Option<Origin>,
178 #[serde(default)]
184 pub format: IndexFormat,
185 pub publish_url: Option<DisplaySafeUrl>,
198 #[serde(default)]
207 pub authenticate: AuthPolicy,
208 #[serde(default)]
218 pub ignore_error_codes: Option<Vec<SerializableStatusCode>>,
219 #[serde(default)]
231 pub cache_control: Option<IndexCacheControl>,
232}
233
234impl PartialEq for Index {
235 fn eq(&self, other: &Self) -> bool {
236 let Self {
237 name,
238 url,
239 explicit,
240 default,
241 origin: _,
242 format,
243 publish_url,
244 authenticate,
245 ignore_error_codes,
246 cache_control,
247 } = self;
248 *url == other.url
249 && *name == other.name
250 && *explicit == other.explicit
251 && *default == other.default
252 && *format == other.format
253 && *publish_url == other.publish_url
254 && *authenticate == other.authenticate
255 && *ignore_error_codes == other.ignore_error_codes
256 && *cache_control == other.cache_control
257 }
258}
259
260impl Eq for Index {}
261
262impl PartialOrd for Index {
263 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
264 Some(self.cmp(other))
265 }
266}
267
268impl Ord for Index {
269 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
270 let Self {
271 name,
272 url,
273 explicit,
274 default,
275 origin: _,
276 format,
277 publish_url,
278 authenticate,
279 ignore_error_codes,
280 cache_control,
281 } = self;
282 url.cmp(&other.url)
283 .then_with(|| name.cmp(&other.name))
284 .then_with(|| explicit.cmp(&other.explicit))
285 .then_with(|| default.cmp(&other.default))
286 .then_with(|| format.cmp(&other.format))
287 .then_with(|| publish_url.cmp(&other.publish_url))
288 .then_with(|| authenticate.cmp(&other.authenticate))
289 .then_with(|| ignore_error_codes.cmp(&other.ignore_error_codes))
290 .then_with(|| cache_control.cmp(&other.cache_control))
291 }
292}
293
294impl std::hash::Hash for Index {
295 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
296 let Self {
297 name,
298 url,
299 explicit,
300 default,
301 origin: _,
302 format,
303 publish_url,
304 authenticate,
305 ignore_error_codes,
306 cache_control,
307 } = self;
308 url.hash(state);
309 name.hash(state);
310 explicit.hash(state);
311 default.hash(state);
312 format.hash(state);
313 publish_url.hash(state);
314 authenticate.hash(state);
315 ignore_error_codes.hash(state);
316 cache_control.hash(state);
317 }
318}
319
320#[derive(
321 Default,
322 Debug,
323 Copy,
324 Clone,
325 Hash,
326 Eq,
327 PartialEq,
328 Ord,
329 PartialOrd,
330 serde::Serialize,
331 serde::Deserialize,
332)]
333#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
334#[serde(rename_all = "kebab-case")]
335pub enum IndexFormat {
336 #[default]
338 Simple,
339 Flat,
341}
342
343impl Index {
344 pub fn from_index_url(url: IndexUrl) -> Self {
346 Self {
347 url,
348 name: None,
349 explicit: false,
350 default: true,
351 origin: None,
352 format: IndexFormat::Simple,
353 publish_url: None,
354 authenticate: AuthPolicy::default(),
355 ignore_error_codes: None,
356 cache_control: None,
357 }
358 }
359
360 pub fn from_extra_index_url(url: IndexUrl) -> Self {
362 Self {
363 url,
364 name: None,
365 explicit: false,
366 default: false,
367 origin: None,
368 format: IndexFormat::Simple,
369 publish_url: None,
370 authenticate: AuthPolicy::default(),
371 ignore_error_codes: None,
372 cache_control: None,
373 }
374 }
375
376 pub fn from_find_links(url: IndexUrl) -> Self {
378 Self {
379 url,
380 name: None,
381 explicit: false,
382 default: false,
383 origin: None,
384 format: IndexFormat::Flat,
385 publish_url: None,
386 authenticate: AuthPolicy::default(),
387 ignore_error_codes: None,
388 cache_control: None,
389 }
390 }
391
392 #[must_use]
394 pub fn with_origin(mut self, origin: Origin) -> Self {
395 self.origin = Some(origin);
396 self
397 }
398
399 pub fn url(&self) -> &IndexUrl {
401 &self.url
402 }
403
404 pub fn into_url(self) -> IndexUrl {
406 self.url
407 }
408
409 pub fn raw_url(&self) -> &DisplaySafeUrl {
411 self.url.url()
412 }
413
414 pub fn root_url(&self) -> Option<DisplaySafeUrl> {
419 self.url.root()
420 }
421
422 #[must_use]
427 pub fn with_promoted_auth_policy(mut self) -> Self {
428 if matches!(self.authenticate, AuthPolicy::Auto) && self.credentials().is_some() {
429 self.authenticate = AuthPolicy::Always;
430 }
431 self
432 }
433
434 pub fn credentials(&self) -> Option<Credentials> {
436 if let Some(name) = self.name.as_ref() {
438 if let Some(credentials) = Credentials::from_env(name.to_env_var()) {
439 return Some(credentials);
440 }
441 }
442
443 Credentials::from_url(self.url.url())
445 }
446
447 pub fn relative_to(mut self, root_dir: &Path) -> Result<Self, IndexUrlError> {
449 if let IndexUrl::Path(ref url) = self.url {
450 if let Some(given) = url.given() {
451 self.url = IndexUrl::parse(given, Some(root_dir))?;
452 }
453 }
454 Ok(self)
455 }
456
457 pub fn status_code_strategy(&self) -> IndexStatusCodeStrategy {
459 if let Some(ignore_error_codes) = &self.ignore_error_codes {
460 IndexStatusCodeStrategy::from_ignored_error_codes(ignore_error_codes)
461 } else {
462 IndexStatusCodeStrategy::from_index_url(self.url.url())
463 }
464 }
465
466 pub fn artifact_cache_control(&self) -> Option<HeaderValue> {
468 self.cache_control
469 .as_ref()
470 .and_then(|cache_control| cache_control.files.clone())
471 .or_else(|| IndexCacheControl::artifact_cache_control(self.url.url()))
472 }
473
474 pub fn simple_api_cache_control(&self) -> Option<HeaderValue> {
476 self.cache_control
477 .as_ref()
478 .and_then(|cache_control| cache_control.api.clone())
479 .or_else(|| IndexCacheControl::simple_api_cache_control(self.url.url()))
480 }
481}
482
483impl From<IndexUrl> for Index {
484 fn from(value: IndexUrl) -> Self {
485 Self {
486 name: None,
487 url: value,
488 explicit: false,
489 default: false,
490 origin: None,
491 format: IndexFormat::Simple,
492 publish_url: None,
493 authenticate: AuthPolicy::default(),
494 ignore_error_codes: None,
495 cache_control: None,
496 }
497 }
498}
499
500impl FromStr for Index {
501 type Err = IndexSourceError;
502
503 fn from_str(s: &str) -> Result<Self, Self::Err> {
504 if let Some((name, url)) = s.split_once('=') {
506 if !name.chars().any(|c| c == ':') {
507 let name = IndexName::from_str(name)?;
508 let url = IndexUrl::from_str(url)?;
509 return Ok(Self {
510 name: Some(name),
511 url,
512 explicit: false,
513 default: false,
514 origin: None,
515 format: IndexFormat::Simple,
516 publish_url: None,
517 authenticate: AuthPolicy::default(),
518 ignore_error_codes: None,
519 cache_control: None,
520 });
521 }
522 }
523
524 let url = IndexUrl::from_str(s)?;
526 Ok(Self {
527 name: None,
528 url,
529 explicit: false,
530 default: false,
531 origin: None,
532 format: IndexFormat::Simple,
533 publish_url: None,
534 authenticate: AuthPolicy::default(),
535 ignore_error_codes: None,
536 cache_control: None,
537 })
538 }
539}
540
541#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
543pub struct IndexMetadata {
544 pub url: IndexUrl,
546 pub format: IndexFormat,
548}
549
550impl IndexMetadata {
551 pub fn as_ref(&self) -> IndexMetadataRef<'_> {
553 let Self { url, format: kind } = self;
554 IndexMetadataRef { url, format: *kind }
555 }
556
557 pub fn into_url(self) -> IndexUrl {
559 self.url
560 }
561}
562
563#[derive(Debug, Copy, Clone)]
565pub struct IndexMetadataRef<'a> {
566 pub url: &'a IndexUrl,
568 pub format: IndexFormat,
570}
571
572impl IndexMetadata {
573 pub fn url(&self) -> &IndexUrl {
575 &self.url
576 }
577}
578
579impl IndexMetadataRef<'_> {
580 pub fn url(&self) -> &IndexUrl {
582 self.url
583 }
584}
585
586impl<'a> From<&'a Index> for IndexMetadataRef<'a> {
587 fn from(value: &'a Index) -> Self {
588 Self {
589 url: &value.url,
590 format: value.format,
591 }
592 }
593}
594
595impl<'a> From<&'a IndexMetadata> for IndexMetadataRef<'a> {
596 fn from(value: &'a IndexMetadata) -> Self {
597 Self {
598 url: &value.url,
599 format: value.format,
600 }
601 }
602}
603
604impl From<IndexUrl> for IndexMetadata {
605 fn from(value: IndexUrl) -> Self {
606 Self {
607 url: value,
608 format: IndexFormat::Simple,
609 }
610 }
611}
612
613impl<'a> From<&'a IndexUrl> for IndexMetadataRef<'a> {
614 fn from(value: &'a IndexUrl) -> Self {
615 Self {
616 url: value,
617 format: IndexFormat::Simple,
618 }
619 }
620}
621
622#[derive(Deserialize)]
624#[serde(rename_all = "kebab-case")]
625struct IndexWire {
626 name: Option<IndexName>,
627 url: IndexUrl,
628 #[serde(default)]
629 explicit: bool,
630 #[serde(default)]
631 default: bool,
632 #[serde(default)]
633 format: IndexFormat,
634 publish_url: Option<DisplaySafeUrl>,
635 #[serde(default)]
636 authenticate: AuthPolicy,
637 #[serde(default)]
638 ignore_error_codes: Option<Vec<SerializableStatusCode>>,
639 #[serde(default)]
640 cache_control: Option<IndexCacheControl>,
641}
642
643impl<'de> Deserialize<'de> for Index {
644 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
645 where
646 D: serde::Deserializer<'de>,
647 {
648 let wire = IndexWire::deserialize(deserializer)?;
649
650 if wire.explicit && wire.name.is_none() {
651 return Err(serde::de::Error::custom(format!(
652 "An index with `explicit = true` requires a `name`: {}",
653 wire.url
654 )));
655 }
656
657 Ok(Self {
658 name: wire.name,
659 url: wire.url,
660 explicit: wire.explicit,
661 default: wire.default,
662 origin: None,
663 format: wire.format,
664 publish_url: wire.publish_url,
665 authenticate: wire.authenticate,
666 ignore_error_codes: wire.ignore_error_codes,
667 cache_control: wire.cache_control,
668 })
669 }
670}
671
672#[derive(Error, Debug)]
674pub enum IndexSourceError {
675 #[error(transparent)]
676 Url(#[from] IndexUrlError),
677 #[error(transparent)]
678 IndexName(#[from] IndexNameError),
679 #[error("Index included a name, but the name was empty")]
680 EmptyName,
681}
682
683#[cfg(test)]
684mod tests {
685 use super::*;
686 use http::HeaderValue;
687
688 #[test]
689 fn test_index_cache_control_headers() {
690 let toml_str = r#"
692 name = "test-index"
693 url = "https://test.example.com/simple"
694 cache-control = { api = "max-age=600", files = "max-age=3600" }
695 "#;
696
697 let index: Index = toml::from_str(toml_str).unwrap();
698 assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
699 assert!(index.cache_control.is_some());
700 let cache_control = index.cache_control.as_ref().unwrap();
701 assert_eq!(
702 cache_control.api,
703 Some(HeaderValue::from_static("max-age=600"))
704 );
705 assert_eq!(
706 cache_control.files,
707 Some(HeaderValue::from_static("max-age=3600"))
708 );
709 }
710
711 #[test]
712 fn test_index_without_cache_control() {
713 let toml_str = r#"
715 name = "test-index"
716 url = "https://test.example.com/simple"
717 "#;
718
719 let index: Index = toml::from_str(toml_str).unwrap();
720 assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
721 assert_eq!(index.cache_control, None);
722 }
723
724 #[test]
725 fn test_index_partial_cache_control() {
726 let toml_str = r#"
728 name = "test-index"
729 url = "https://test.example.com/simple"
730 cache-control = { api = "max-age=300" }
731 "#;
732
733 let index: Index = toml::from_str(toml_str).unwrap();
734 assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
735 assert!(index.cache_control.is_some());
736 let cache_control = index.cache_control.as_ref().unwrap();
737 assert_eq!(
738 cache_control.api,
739 Some(HeaderValue::from_static("max-age=300"))
740 );
741 assert_eq!(cache_control.files, None);
742 }
743
744 #[test]
745 fn test_index_invalid_api_cache_control() {
746 let toml_str = r#"
747 name = "test-index"
748 url = "https://test.example.com/simple"
749 cache-control = { api = "max-age=600\n" }
750 "#;
751
752 let err = toml::from_str::<Index>(toml_str).unwrap_err();
753 assert!(
754 err.to_string()
755 .contains("`cache-control.api` must be a valid HTTP header value")
756 );
757 }
758
759 #[test]
760 fn test_index_invalid_files_cache_control() {
761 let toml_str = r#"
762 name = "test-index"
763 url = "https://test.example.com/simple"
764 cache-control = { files = "max-age=3600\n" }
765 "#;
766
767 let err = toml::from_str::<Index>(toml_str).unwrap_err();
768 assert!(
769 err.to_string()
770 .contains("`cache-control.files` must be a valid HTTP header value")
771 );
772 }
773}