Skip to main content

uv_distribution/
distribution_database.rs

1use std::cmp::Reverse;
2use std::future::Future;
3use std::io;
4use std::path::Path;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context, Poll};
8
9use futures::{FutureExt, TryStreamExt};
10use http_content_range::{ContentRange, ContentRangeBytes, ContentRangeUnbound};
11use rayon::in_place_scope;
12use rayon::prelude::*;
13use rustc_hash::FxHashMap;
14use tokio::io::{AsyncRead, AsyncSeekExt, AsyncWriteExt, ReadBuf};
15use tokio::sync::Semaphore;
16use tokio_util::compat::FuturesAsyncReadCompatExt;
17use tracing::{Instrument, debug, info_span, instrument, warn};
18
19use uv_cache::{ArchiveFileId, ArchiveId, Cache, CacheBucket, CacheEntry, WheelCache};
20use uv_cache_info::{CacheInfo, Timestamp};
21use uv_client::{
22    CacheControl, CachedClientError, Connectivity, DataWithCachePolicy, RegistryClient,
23    RequestBuilder, RetryState,
24};
25use uv_distribution_filename::WheelFilename;
26use uv_distribution_types::{
27    ArchiveHashPolicy, BuildInfo, BuildableSource, BuiltDist, Dist, DistRef, HashCollection,
28    HashValidation, Hashed, IndexUrl, InstalledDist, MetadataHashPolicy, Name, ResolutionRecorder,
29    SourceDist, SourceUrl, parse_url_hashes,
30};
31use uv_extract::dirhash::{DirectoryDigest, HashedFile};
32use uv_extract::hash::Hasher;
33use uv_fs::{LockedFile, write_atomic};
34use uv_git::{GIT_LFS, GitError};
35use uv_normalize::PackageName;
36use uv_pep440::Version;
37use uv_platform_tags::Tags;
38use uv_preview::PreviewFeature;
39use uv_pypi_types::{HashDigest, HashDigests, PyProjectToml, ResolutionMetadata};
40use uv_python_types::PythonVariant;
41use uv_redacted::SafeUrl;
42use uv_threads::initialize_rayon_once;
43use uv_types::{BuildContext, BuildStack};
44
45use crate::archive::Archive;
46use crate::download::variant_compatible;
47use crate::error::PythonVersion;
48use crate::extracted_wheel::{ExtractedWheel, HashedWheel, WheelExtractor};
49use crate::hash::http_hash_algorithms;
50use crate::metadata::{ArchiveMetadata, Metadata};
51use crate::source::SourceDistributionBuilder;
52use crate::{
53    Error, FirstPartyPackages, LocalWheel, Reporter, RequiresDist, check_supported_wheel_variant,
54};
55
56/// A cached high-level interface to convert distributions (a requirement resolved to a location)
57/// to a wheel or wheel metadata.
58///
59/// For wheel metadata, this happens by either fetching the metadata from the remote wheel or by
60/// building the source distribution. For wheel files, either the wheel is downloaded or a source
61/// distribution is downloaded, built and the new wheel gets returned.
62///
63/// All kinds of wheel sources (index, URL, path) and source distribution source (index, URL, path,
64/// Git) are supported.
65///
66/// This struct also has the task of acquiring locks around source dist builds in general and git
67/// operation especially, as well as respecting concurrency limits.
68pub struct DistributionDatabase<'a, Context: BuildContext> {
69    build_context: &'a Context,
70    recorder: Option<ResolutionRecorder>,
71    builder: SourceDistributionBuilder<'a, Context>,
72    client: ManagedClient<'a>,
73    reporter: Option<Arc<dyn Reporter>>,
74    content_addressed_cache: bool,
75    first_party_packages: Option<&'a FirstPartyPackages>,
76}
77
78impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
79    pub fn new(
80        client: &'a RegistryClient,
81        build_context: &'a Context,
82        downloads_semaphore: Arc<Semaphore>,
83    ) -> Self {
84        // When ZIP validation is disabled, the extracted tree can contain files that aren't
85        // represented in the central directory and therefore aren't included in its digest.
86        // Avoid using an incomplete digest as a content-addressed archive ID.
87        let content_addressed_cache = uv_preview::is_enabled(PreviewFeature::ContentAddressedCache)
88            && !uv_extract::insecure_no_validate();
89        Self {
90            recorder: None,
91            build_context,
92            builder: SourceDistributionBuilder::new(build_context),
93            client: ManagedClient::new(client, downloads_semaphore),
94            reporter: None,
95            content_addressed_cache,
96            first_party_packages: None,
97        }
98    }
99
100    /// Allow metadata builds for the given first-party workspace source trees.
101    #[must_use]
102    pub fn with_first_party_packages(
103        mut self,
104        first_party_packages: &'a FirstPartyPackages,
105    ) -> Self {
106        self.first_party_packages = Some(first_party_packages);
107        self
108    }
109
110    /// Return whether the name and path identify an eligible workspace member.
111    pub fn is_first_party(&self, name: &PackageName, path: &Path) -> bool {
112        self.first_party_packages
113            .is_some_and(|packages| packages.contains(name, path))
114    }
115
116    /// Record which static metadata entries are consulted while resolving runtime dependencies.
117    #[must_use]
118    pub fn with_recorder(mut self, recorder: Option<ResolutionRecorder>) -> Self {
119        self.recorder = recorder;
120        self
121    }
122
123    /// Record a metadata lookup before reading an in-memory cache.
124    pub fn record_metadata(&self, dist: &Dist) {
125        if let Some(recorder) = &self.recorder {
126            recorder.dependency_metadata(dist.name());
127        }
128    }
129
130    /// Look up user-provided metadata, recording misses as well as matches.
131    pub fn dependency_metadata(
132        &self,
133        name: &PackageName,
134        version: Option<&Version>,
135    ) -> Option<ResolutionMetadata> {
136        if let Some(recorder) = &self.recorder {
137            recorder.dependency_metadata(name);
138        }
139        self.build_context.dependency_metadata().get(name, version)
140    }
141
142    /// Set the build stack to use for the [`DistributionDatabase`].
143    #[must_use]
144    pub fn with_build_stack(self, build_stack: &'a BuildStack) -> Self {
145        Self {
146            builder: self.builder.with_build_stack(build_stack),
147            ..self
148        }
149    }
150
151    /// Set the [`Reporter`] to use for the [`DistributionDatabase`].
152    #[must_use]
153    pub fn with_reporter(self, reporter: Arc<dyn Reporter>) -> Self {
154        Self {
155            builder: self.builder.with_reporter(reporter.clone()),
156            reporter: Some(reporter),
157            ..self
158        }
159    }
160
161    /// Handle a specific `reqwest` error, and convert it to [`io::Error`].
162    fn handle_response_errors(&self, err: reqwest::Error) -> io::Error {
163        if err.is_timeout() {
164            // Assumption: The connect timeout with the 10s default is not the culprit.
165            io::Error::new(
166                io::ErrorKind::TimedOut,
167                format!(
168                    "Failed to download distribution due to network timeout. Try increasing UV_HTTP_TIMEOUT (current value: {}s).",
169                    self.client.unmanaged.read_timeout().as_secs()
170                ),
171            )
172        } else {
173            io::Error::other(err)
174        }
175    }
176
177    /// Acquire an advisory lock for a wheel cache entry.
178    ///
179    /// A remote wheel's content hash is not always available until after the download, so
180    /// concurrent cache fills coordinate on the wheel cache entry instead. The entry is already
181    /// scoped to the distribution's source and wheel filename.
182    ///
183    /// Callers hold the returned lock across cache lookup, download or extraction, and publication.
184    /// A process that waited for another cache fill therefore rechecks and reuses the completed entry.
185    async fn lock_wheel(
186        wheel_entry: &CacheEntry,
187        filename: &WheelFilename,
188    ) -> Result<LockedFile, Error> {
189        // Use the bounded cache key to avoid filesystem filename limits.
190        let lock_key = filename.cache_key();
191
192        let lock_entry = wheel_entry.with_file(format!("{lock_key}.lock"));
193        lock_entry.lock().await.map_err(Error::CacheLock)
194    }
195
196    /// Validate dist hashes, and return them.
197    fn validate_hashes(
198        dist: &BuiltDist,
199        hashes: ArchiveHashPolicy<'_>,
200        hashers: Vec<Hasher>,
201    ) -> Result<HashDigests, Error> {
202        let computed_hashes: HashDigests = hashers.into_iter().map(HashDigest::from).collect();
203        if hashes.requires_validation() && !hashes.matches(computed_hashes.as_slice()) {
204            return Err(Error::hash_mismatch(
205                dist.to_string(),
206                hashes.digests(),
207                computed_hashes.as_slice(),
208            ));
209        }
210        Ok(computed_hashes)
211    }
212
213    /// Either fetch the wheel or fetch and build the source distribution
214    ///
215    /// Returns a wheel that's compliant with the given platform tags.
216    ///
217    /// Applicable hash checks are enforced before newly fetched archives are published to the cache.
218    /// Callers must enforce their hash policy when reusing cached wheels.
219    #[instrument(skip_all, fields(%dist))]
220    pub async fn get_or_build_wheel(
221        &self,
222        dist: &Dist,
223        tags: &Tags,
224        hashes: ArchiveHashPolicy<'_>,
225    ) -> Result<LocalWheel, Error> {
226        match dist {
227            Dist::Built(built) => {
228                check_supported_wheel_variant(built)?;
229                let wheel = self.get_wheel(built, hashes).await?;
230                if !variant_compatible(&wheel.filename, &wheel.archive)? {
231                    return Err(Error::IncompatibleWheelVariant(Box::new(wheel.filename)));
232                }
233                Ok(wheel)
234            }
235            // TODO(konsti): Support source distributions that build variant wheels.
236            // For now, `build_wheel` assumes a non-variant wheel compatible with `tags`.
237            Dist::Source(source) => self.build_wheel(source, tags, hashes).await,
238        }
239    }
240
241    /// Either fetch the only wheel metadata (directly from the index or with range requests) or
242    /// fetch and build the source distribution.
243    ///
244    /// While hashes will be generated in some cases, hash-checking is only enforced for source
245    /// distributions, and should be enforced by the caller for wheels.
246    #[instrument(skip_all, fields(%dist))]
247    pub async fn get_installed_metadata(
248        &self,
249        dist: &InstalledDist,
250    ) -> Result<ArchiveMetadata, Error> {
251        // If the metadata was provided by the user directly, prefer it.
252        if let Some(metadata) = self.dependency_metadata(dist.name(), Some(dist.version())) {
253            return Ok(Metadata::from_dependency_metadata(metadata).into());
254        }
255
256        let metadata = dist
257            .read_metadata()
258            .map_err(|err| Error::ReadInstalled(Box::new(dist.clone()), err))?;
259
260        Ok(ArchiveMetadata::from_metadata23(metadata.clone()))
261    }
262
263    /// Retrieve distribution metadata and any hashes requested for resolution.
264    ///
265    /// Source archive hashes are validated before executing build backends; wheel archive hash
266    /// validation is deferred to installation.
267    #[instrument(skip_all, fields(%dist))]
268    pub async fn get_or_build_wheel_metadata(
269        &self,
270        dist: &Dist,
271        hashes: MetadataHashPolicy<'_>,
272    ) -> Result<ArchiveMetadata, Error> {
273        match dist {
274            Dist::Built(built) => self.get_wheel_metadata(built, hashes).await,
275            Dist::Source(source) => {
276                self.build_wheel_metadata(&BuildableSource::Dist(source), hashes)
277                    .await
278            }
279        }
280    }
281
282    /// Fetch a wheel from the cache or download it from the index.
283    ///
284    /// Applicable hash checks are enforced before newly fetched wheels are published to the cache.
285    /// Registry hashes are used when available and no explicit verification policy is provided.
286    /// Callers must enforce their hash policy when reusing cached wheels.
287    async fn get_wheel(
288        &self,
289        dist: &BuiltDist,
290        hashes: ArchiveHashPolicy<'_>,
291    ) -> Result<LocalWheel, Error> {
292        match dist {
293            BuiltDist::Registry(wheels) => {
294                let wheel = wheels.best_wheel();
295                let url = wheel.file.url.to_url()?;
296                let size = wheel.file.size;
297
298                // Create a cache entry for the wheel.
299                let wheel_entry = self.build_context.cache().entry(
300                    CacheBucket::Wheels,
301                    WheelCache::Index(&wheel.index).wheel_dir(wheel.name().as_ref()),
302                    wheel.filename.cache_key(),
303                );
304
305                // If the URL is a file URL, load the wheel directly.
306                if url.scheme() == "file" {
307                    let path = url
308                        .to_file_path()
309                        .map_err(|()| Error::NonFileUrl(url.clone()))?;
310                    return self
311                        .load_wheel(&path, &wheel.filename, wheel_entry, dist, hashes)
312                        .await;
313                }
314
315                // Download and unzip.
316                match self
317                    .stream_wheel(
318                        url.clone(),
319                        dist.index(),
320                        &wheel.filename,
321                        size,
322                        &wheel_entry,
323                        dist,
324                        hashes,
325                    )
326                    .await
327                {
328                    Ok(archive) => Ok(LocalWheel {
329                        dist: Dist::Built(dist.clone()),
330                        archive: self
331                            .build_context
332                            .cache()
333                            .archive(&archive.id)
334                            .into_boxed_path(),
335                        hashes: archive.hashes,
336                        filename: wheel.filename.clone(),
337                        cache: CacheInfo::default(),
338                        build: None,
339                    }),
340                    Err(Error::Extract(name, err)) => {
341                        if err.is_http_streaming_unsupported() {
342                            warn!(
343                                "Streaming unsupported for `{dist}`; downloading wheel to disk ({err})"
344                            );
345                        } else if err.is_http_streaming_failed() {
346                            warn!(
347                                "Streaming failed for `{dist}`; downloading wheel to disk ({err})"
348                            );
349                        } else {
350                            return Err(Error::Extract(name, err));
351                        }
352
353                        // If the request failed because streaming was unsupported or failed,
354                        // download the wheel directly.
355                        let archive = self
356                            .download_wheel(
357                                url,
358                                dist.index(),
359                                &wheel.filename,
360                                size,
361                                &wheel_entry,
362                                dist,
363                                hashes,
364                            )
365                            .await?;
366
367                        Ok(LocalWheel {
368                            dist: Dist::Built(dist.clone()),
369                            archive: self
370                                .build_context
371                                .cache()
372                                .archive(&archive.id)
373                                .into_boxed_path(),
374                            hashes: archive.hashes,
375                            filename: wheel.filename.clone(),
376                            cache: CacheInfo::default(),
377                            build: None,
378                        })
379                    }
380                    Err(err) => Err(err),
381                }
382            }
383
384            BuiltDist::DirectUrl(wheel) => {
385                // Create a cache entry for the wheel.
386                let wheel_entry = self.build_context.cache().entry(
387                    CacheBucket::Wheels,
388                    WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
389                    wheel.filename.cache_key(),
390                );
391
392                // Download and unzip.
393                match self
394                    .stream_wheel(
395                        wheel.url.raw().clone(),
396                        None,
397                        &wheel.filename,
398                        wheel.size,
399                        &wheel_entry,
400                        dist,
401                        hashes,
402                    )
403                    .await
404                {
405                    Ok(archive) => Ok(LocalWheel {
406                        dist: Dist::Built(dist.clone()),
407                        archive: self
408                            .build_context
409                            .cache()
410                            .archive(&archive.id)
411                            .into_boxed_path(),
412                        hashes: archive.hashes,
413                        filename: wheel.filename.clone(),
414                        cache: CacheInfo::default(),
415                        build: None,
416                    }),
417                    Err(Error::Extract(name, err)) => {
418                        if err.is_http_streaming_unsupported() {
419                            warn!(
420                                "Streaming unsupported for `{dist}`; downloading wheel to disk ({err})"
421                            );
422                        } else if err.is_http_streaming_failed() {
423                            warn!(
424                                "Streaming failed for `{dist}`; downloading wheel to disk ({err})"
425                            );
426                        } else {
427                            return Err(Error::Extract(name, err));
428                        }
429
430                        // If the request failed because streaming was unsupported or failed,
431                        // download the wheel directly.
432                        let archive = self
433                            .download_wheel(
434                                wheel.url.raw().clone(),
435                                None,
436                                &wheel.filename,
437                                wheel.size,
438                                &wheel_entry,
439                                dist,
440                                hashes,
441                            )
442                            .await?;
443                        Ok(LocalWheel {
444                            dist: Dist::Built(dist.clone()),
445                            archive: self
446                                .build_context
447                                .cache()
448                                .archive(&archive.id)
449                                .into_boxed_path(),
450                            hashes: archive.hashes,
451                            filename: wheel.filename.clone(),
452                            cache: CacheInfo::default(),
453                            build: None,
454                        })
455                    }
456                    Err(err) => Err(err),
457                }
458            }
459
460            BuiltDist::GitPath(wheel) => {
461                // Fetch the Git repository.
462                let fetch = self
463                    .build_context
464                    .git()
465                    .fetch(
466                        &wheel.git,
467                        self.client.unmanaged.git_http_settings(wheel.git.url()),
468                        self.build_context.cache().bucket(CacheBucket::Git),
469                        self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
470                    )
471                    .await?;
472
473                if wheel.git.lfs().enabled() && !fetch.lfs_ready() {
474                    if GIT_LFS.is_err() {
475                        return Err(Error::MissingWheelGitLfsArtifacts(
476                            wheel.url.to_url(),
477                            GitError::GitLfsNotFound,
478                        ));
479                    }
480                    return Err(Error::MissingWheelGitLfsArtifacts(
481                        wheel.url.to_url(),
482                        GitError::GitLfsNotConfigured,
483                    ));
484                }
485
486                let git_sha = fetch.git().precise().expect("Exact commit after checkout");
487                let cache_entry = self.build_context.cache().entry(
488                    CacheBucket::Wheels,
489                    WheelCache::Git(&wheel.url, git_sha.as_short_str()).root(),
490                    wheel.filename.stem(),
491                );
492
493                let install_path = fetch.path().join(&wheel.install_path);
494
495                self.load_wheel(&install_path, &wheel.filename, cache_entry, dist, hashes)
496                    .await
497            }
498
499            BuiltDist::Path(wheel) => {
500                let cache_entry = self.build_context.cache().entry(
501                    CacheBucket::Wheels,
502                    WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
503                    wheel.filename.cache_key(),
504                );
505
506                self.load_wheel(
507                    &wheel.install_path,
508                    &wheel.filename,
509                    cache_entry,
510                    dist,
511                    hashes,
512                )
513                .await
514            }
515        }
516    }
517
518    /// Convert a source distribution into a wheel, fetching it from the cache or building it if
519    /// necessary.
520    ///
521    /// The returned wheel is guaranteed to come from a distribution with a matching hash, and
522    /// no build processes will be executed for distributions with mismatched hashes.
523    async fn build_wheel(
524        &self,
525        dist: &SourceDist,
526        tags: &Tags,
527        hashes: ArchiveHashPolicy<'_>,
528    ) -> Result<LocalWheel, Error> {
529        let built_wheel = self
530            .builder
531            .download_and_build(&BuildableSource::Dist(dist), tags, hashes, &self.client)
532            .boxed_local()
533            .await?;
534
535        // Check that the wheel is compatible with its install target.
536        //
537        // When building a build dependency for a cross-install, the build dependency needs
538        // to install and run on the host instead of the target. In this case the `tags` are already
539        // for the host instead of the target, so this check passes.
540        if !built_wheel.filename.is_compatible(tags) {
541            return if tags.is_cross() {
542                Err(Error::BuiltWheelIncompatibleTargetPlatform {
543                    filename: built_wheel.filename,
544                    python_platform: tags.python_platform().clone(),
545                    python_version: PythonVersion {
546                        version: tags.python_version(),
547                        variant: if tags.is_freethreaded() {
548                            PythonVariant::Freethreaded
549                        } else {
550                            PythonVariant::Default
551                        },
552                    },
553                })
554            } else {
555                Err(Error::BuiltWheelIncompatibleHostPlatform {
556                    filename: built_wheel.filename,
557                    python_platform: tags.python_platform().clone(),
558                    python_version: PythonVersion {
559                        version: tags.python_version(),
560                        variant: if tags.is_freethreaded() {
561                            PythonVariant::Freethreaded
562                        } else {
563                            PythonVariant::Default
564                        },
565                    },
566                })
567            };
568        }
569
570        // Acquire the advisory lock.
571        let wheel_entry = CacheEntry::from_path(built_wheel.target.as_ref());
572        let _lock = Self::lock_wheel(&wheel_entry, &built_wheel.filename).await?;
573
574        // If the wheel was unzipped previously, respect it. Source distributions are
575        // cached under a unique revision ID, so unzipped directories are never stale.
576        match self.build_context.cache().resolve_link(&built_wheel.target) {
577            Ok(archive) => {
578                return Ok(LocalWheel {
579                    dist: Dist::Source(dist.clone()),
580                    archive: archive.into_boxed_path(),
581                    filename: built_wheel.filename,
582                    hashes: built_wheel.hashes,
583                    cache: built_wheel.cache_info,
584                    build: Some(built_wheel.build_info),
585                });
586            }
587            Err(err) if err.kind() == io::ErrorKind::NotFound => {}
588            Err(err) => return Err(Error::CacheRead(err)),
589        }
590
591        // Otherwise, unzip the wheel.
592        let id = self
593            .unzip_wheel(
594                &built_wheel.path,
595                &built_wheel.target,
596                DistRef::Source(dist),
597            )
598            .await?;
599
600        Ok(LocalWheel {
601            dist: Dist::Source(dist.clone()),
602            archive: self.build_context.cache().archive(&id).into_boxed_path(),
603            hashes: built_wheel.hashes,
604            filename: built_wheel.filename,
605            cache: built_wheel.cache_info,
606            build: Some(built_wheel.build_info),
607        })
608    }
609
610    /// Fetch wheel metadata, along with any hashes requested for resolution.
611    ///
612    /// Wheel archive hash validation is deferred to installation. Metadata sidecar hashes are
613    /// checked by the client when downloading sidecars with index-provided hashes.
614    async fn get_wheel_metadata(
615        &self,
616        dist: &BuiltDist,
617        hashes: MetadataHashPolicy<'_>,
618    ) -> Result<ArchiveMetadata, Error> {
619        let hash_policy = match hashes.validation {
620            HashValidation::None => {
621                let compute_hashes = match hashes.collection {
622                    HashCollection::None => false,
623                    HashCollection::Url | HashCollection::All => match dist {
624                        BuiltDist::Registry(dist) => {
625                            // Generate missing hashes for indexes and `--find-links` without hashes.
626                            // This hashes only the selected wheel, not every distribution for the version.
627                            hashes.collection == HashCollection::All
628                                && dist.best_wheel().file.hashes.is_empty()
629                        }
630                        BuiltDist::DirectUrl(dist) => parse_url_hashes(&dist.url).is_none(),
631                        BuiltDist::Path(_) | BuiltDist::GitPath(_) => true,
632                    },
633                };
634                if compute_hashes {
635                    ArchiveHashPolicy::Generate
636                } else {
637                    ArchiveHashPolicy::None
638                }
639            }
640            HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
641        };
642
643        // Fetch the entire wheel only when we need to compute a hash for resolution.
644        // TODO(charlie): Request the hashes via a separate method, to reduce the coupling in this API.
645        if hash_policy == ArchiveHashPolicy::Generate {
646            let wheel = self.get_wheel(dist, hash_policy).await?;
647            // If the metadata was provided by the user directly, prefer it.
648            let metadata = if let Some(metadata) =
649                self.dependency_metadata(dist.name(), Some(dist.version()))
650            {
651                Metadata::from_dependency_metadata(metadata)
652            } else {
653                Metadata::from_metadata23(wheel.metadata()?)
654            };
655            let hashes = wheel.hashes;
656            return Ok(ArchiveMetadata { metadata, hashes });
657        }
658
659        // If the metadata was provided by the user directly, prefer it.
660        if let Some(metadata) = self.dependency_metadata(dist.name(), Some(dist.version())) {
661            return Ok(Metadata::from_dependency_metadata(metadata).into());
662        }
663
664        let result = self
665            .client
666            .managed(|client| {
667                client
668                    .wheel_metadata(
669                        dist,
670                        self.build_context.git(),
671                        self.build_context.capabilities(),
672                        self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
673                    )
674                    .boxed_local()
675            })
676            .await;
677
678        match result {
679            Ok(metadata) => {
680                // Validate that the metadata is consistent with the distribution.
681                Ok(ArchiveMetadata::from_metadata23(metadata))
682            }
683            Err(err) if err.is_http_streaming_unsupported() => {
684                warn!(
685                    "Streaming unsupported when fetching metadata for `{dist}`; downloading wheel directly ({err})"
686                );
687
688                // If the request failed due to an error that could be resolved by
689                // downloading the wheel directly, try that.
690                let wheel = self.get_wheel(dist, hash_policy).await?;
691                let metadata = wheel.metadata()?;
692                let hashes = wheel.hashes;
693                Ok(ArchiveMetadata {
694                    metadata: Metadata::from_metadata23(metadata),
695                    hashes,
696                })
697            }
698            Err(err) => Err(err.into()),
699        }
700    }
701
702    /// Build the wheel metadata for a source distribution, or fetch it from the cache if possible.
703    ///
704    /// Requested hashes are validated before executing a build backend. User-provided metadata
705    /// can avoid downloading or validating the source archive.
706    pub async fn build_wheel_metadata(
707        &self,
708        source: &BuildableSource<'_>,
709        hashes: MetadataHashPolicy<'_>,
710    ) -> Result<ArchiveMetadata, Error> {
711        // If the metadata was provided by the user directly, prefer it.
712        if let Some(dist) = source.as_dist() {
713            if let Some(metadata) = self.dependency_metadata(dist.name(), dist.version()) {
714                // If we skipped the build, we should still resolve any Git dependencies to precise
715                // commits.
716                self.builder.resolve_revision(source, &self.client).await?;
717
718                return Ok(Metadata::from_dependency_metadata(metadata).into());
719            }
720        }
721
722        let url_hashes = if let BuildableSource::Dist(SourceDist::DirectUrl(dist)) = source {
723            parse_url_hashes(&dist.url)
724        } else if let BuildableSource::Url(SourceUrl::Direct(url)) = source {
725            parse_url_hashes(url.url)
726        } else {
727            None
728        };
729
730        let build_hash_policy = match hashes.validation {
731            HashValidation::None => match hashes.collection {
732                HashCollection::None => ArchiveHashPolicy::None,
733                // If resolving metadata requires a build, validate any URL hash before executing
734                // the backend, even when the caller only requested hash collection.
735                HashCollection::Url | HashCollection::All => match url_hashes.as_ref() {
736                    Some(digests) => ArchiveHashPolicy::All(digests.as_slice()),
737                    None => ArchiveHashPolicy::Generate,
738                },
739            },
740            HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
741        };
742        let ArchiveMetadata { metadata, hashes } = self
743            .builder
744            .for_metadata(self.first_party_packages)
745            .download_and_build_metadata(source, build_hash_policy, &self.client)
746            .boxed_local()
747            .await?;
748
749        // This also handles external dependencies, so force relative paths by default.
750        // For explicit inputs and current project/workspace metadata, this is later overridden with
751        // `with_force_relative(false)` to respect the user's relative/absolute path preference.
752        Ok(ArchiveMetadata {
753            metadata: metadata.with_force_relative(true),
754            hashes,
755        })
756    }
757
758    /// Return the [`RequiresDist`] from a `pyproject.toml`, if it can be statically extracted.
759    pub async fn requires_dist(
760        &self,
761        path: &Path,
762        pyproject_toml: &PyProjectToml,
763    ) -> Result<Option<RequiresDist>, Error> {
764        self.builder
765            .source_tree_requires_dist(
766                path,
767                pyproject_toml,
768                self.client.unmanaged.credentials_cache(),
769            )
770            .await
771    }
772
773    /// Stream a wheel from a URL, unzipping it into the cache as it's downloaded.
774    ///
775    /// Note that `progress_size_hint` is a size hint for the progress reporter,
776    /// and is not guaranteed to be the true size of the wheel if/when fetched from the
777    /// origin.
778    async fn stream_wheel(
779        &self,
780        url: SafeUrl,
781        index: Option<&IndexUrl>,
782        filename: &WheelFilename,
783        progress_size_hint: Option<u64>,
784        wheel_entry: &CacheEntry,
785        dist: &BuiltDist,
786        hashes: ArchiveHashPolicy<'_>,
787    ) -> Result<Archive, Error> {
788        // Sometimes we can promote the size hint to a trusted effective size.
789        let expected_size = match dist {
790            BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => {
791                progress_size_hint
792            }
793            BuiltDist::DirectUrl(_) => progress_size_hint,
794            _ => None,
795        };
796
797        // Acquire an advisory lock, to guard against concurrent writes.
798        let _lock = Self::lock_wheel(wheel_entry, filename).await?;
799
800        // Create an entry for the HTTP cache.
801        let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
802
803        let download = |response: reqwest::Response, _: &mut RetryState| {
804            async {
805                let hashes = if let BuiltDist::Registry(wheels) = dist {
806                    hashes.with_index_hashes(wheels.best_wheel().file.hashes.as_slice())
807                } else {
808                    hashes
809                };
810
811                let progress_size_hint = progress_size_hint.or_else(|| content_length(&response));
812
813                let progress = self.reporter.as_ref().map(|reporter| {
814                    (
815                        reporter,
816                        reporter.on_download_start(dist.name(), progress_size_hint),
817                    )
818                });
819
820                let reader = response
821                    .bytes_stream()
822                    .map_err(|err| self.handle_response_errors(err))
823                    .into_async_read();
824
825                // Create a hasher for each hash algorithm.
826                let algorithms = http_hash_algorithms(hashes);
827                let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
828                let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
829
830                // Download and unzip the wheel to a temporary directory.
831                let extractor = WheelExtractor::new(
832                    self.build_context.cache().root(),
833                    self.content_addressed_cache,
834                )
835                .map_err(Error::CacheWrite)?;
836
837                let mut extracted = match progress {
838                    Some((reporter, progress)) => {
839                        let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
840                        extractor
841                            .extract_streaming(&mut reader)
842                            .await
843                            .map_err(|err| Error::Extract(filename.to_string(), err))?
844                    }
845                    None => extractor
846                        .extract_streaming(&mut hasher)
847                        .await
848                        .map_err(|err| Error::Extract(filename.to_string(), err))?,
849                };
850                // Exhaust the reader to compute the hashes.
851                hasher.finish().await.map_err(Error::HashExhaustion)?;
852                let actual_size = hasher.bytes_read();
853                if let Some(expected) = expected_size
854                    && actual_size != expected
855                {
856                    return Err(Error::MismatchedSize {
857                        distribution: dist.to_string(),
858                        expected,
859                        actual: actual_size,
860                    });
861                }
862
863                let computed_hashes = Self::validate_hashes(dist, hashes, hashers)?;
864
865                // Before we make the wheel accessible by persisting it, ensure that the RECORD is
866                // valid.
867                extracted.validate_and_heal_record(dist)?;
868
869                // Persist the temporary directory to the directory store.
870                let id = self
871                    .persist_extracted_wheel(extracted, wheel_entry.path())
872                    .await?;
873
874                if let Some((reporter, progress)) = progress {
875                    reporter.on_download_complete(dist.name(), progress);
876                }
877
878                Ok(Archive::new(
879                    id,
880                    computed_hashes,
881                    filename.clone(),
882                    Some(actual_size),
883                ))
884            }
885            .instrument(info_span!("wheel", wheel = %dist))
886        };
887
888        // Fetch the archive from the cache, or download it if necessary.
889        let req = self.request(url.clone())?;
890
891        // Determine the cache control policy for the URL.
892        let cache_control = match self.client.unmanaged.connectivity() {
893            Connectivity::Online
894                if let Some(header) = index.and_then(|index| {
895                    self.build_context
896                        .locations()
897                        .artifact_cache_control_for(index)
898                }) =>
899            {
900                CacheControl::Override(header)
901            }
902            Connectivity::Online => CacheControl::from(
903                self.build_context
904                    .cache()
905                    .freshness(&http_entry, Some(&filename.name), None)
906                    .map_err(Error::CacheRead)?,
907            ),
908            Connectivity::Offline => CacheControl::AllowStale,
909        };
910
911        let archive = self
912            .client
913            .managed(|client| {
914                client.cached_client().get_serde_with_retry(
915                    req,
916                    &http_entry,
917                    cache_control.clone(),
918                    download,
919                )
920            })
921            .await
922            .map_err(|err| match err {
923                CachedClientError::Callback { err, .. } => err,
924                CachedClientError::Client(err) => Error::Client(err),
925            })?;
926
927        if let (Some(expected), Some(actual)) = (expected_size, archive.size)
928            && expected != actual
929        {
930            return Err(Error::MismatchedSize {
931                distribution: dist.to_string(),
932                expected,
933                actual,
934            });
935        }
936
937        // If the archive is missing the required hashes or size, or has since been removed, force a refresh.
938        let archive = Some(archive)
939            .filter(|archive| archive.has_digests(hashes))
940            .filter(|archive| archive.exists(self.build_context.cache()))
941            .filter(|archive| expected_size.is_none() || archive.size.is_some());
942
943        let archive = if let Some(archive) = archive {
944            archive
945        } else {
946            self.client
947                .managed(async |client| {
948                    client
949                        .cached_client()
950                        .skip_cache_with_retry(
951                            self.request(url)?,
952                            &http_entry,
953                            cache_control,
954                            download,
955                        )
956                        .await
957                        .map_err(|err| match err {
958                            CachedClientError::Callback { err, .. } => err,
959                            CachedClientError::Client(err) => Error::Client(err),
960                        })
961                })
962                .await?
963        };
964
965        Ok(archive)
966    }
967
968    /// Download a wheel from a URL, then unzip it into the cache.
969    ///
970    /// Note that, like [`Self::stream_wheel`], `progress_size_hint` is
971    /// a size hint for the progress reporter, and is not guaranteed to be the size
972    /// of the wheel if/when fetched from the origin.
973    async fn download_wheel(
974        &self,
975        url: SafeUrl,
976        index: Option<&IndexUrl>,
977        filename: &WheelFilename,
978        progress_size_hint: Option<u64>,
979        wheel_entry: &CacheEntry,
980        dist: &BuiltDist,
981        hashes: ArchiveHashPolicy<'_>,
982    ) -> Result<Archive, Error> {
983        // Sometimes we can promote the size hint to a trusted effective size.
984        let expected_size = match dist {
985            BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => {
986                progress_size_hint
987            }
988            BuiltDist::DirectUrl(_) => progress_size_hint,
989            _ => None,
990        };
991
992        // Acquire an advisory lock, to guard against concurrent writes.
993        let _lock = Self::lock_wheel(wheel_entry, filename).await?;
994
995        // Create an entry for the HTTP cache.
996        let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
997
998        let download_url = url.clone();
999
1000        let download = async |response, retry_state: &mut RetryState| {
1001            self.download_wheel_response(
1002                response,
1003                &download_url,
1004                retry_state,
1005                filename,
1006                progress_size_hint,
1007                expected_size,
1008                wheel_entry,
1009                dist,
1010                hashes,
1011            )
1012            .instrument(info_span!("wheel", wheel = %dist))
1013            .await
1014        };
1015
1016        // Fetch the archive from the cache, or download it if necessary.
1017        let req = self.request(url.clone())?;
1018
1019        // Determine the cache control policy for the URL.
1020        let cache_control = match self.client.unmanaged.connectivity() {
1021            Connectivity::Online
1022                if let Some(header) = index.and_then(|index| {
1023                    self.build_context
1024                        .locations()
1025                        .artifact_cache_control_for(index)
1026                }) =>
1027            {
1028                CacheControl::Override(header)
1029            }
1030            Connectivity::Online => CacheControl::from(
1031                self.build_context
1032                    .cache()
1033                    .freshness(&http_entry, Some(&filename.name), None)
1034                    .map_err(Error::CacheRead)?,
1035            ),
1036            Connectivity::Offline => CacheControl::AllowStale,
1037        };
1038
1039        let archive = self
1040            .client
1041            .managed(|client| {
1042                client.cached_client().get_serde_with_retry(
1043                    req,
1044                    &http_entry,
1045                    cache_control.clone(),
1046                    download,
1047                )
1048            })
1049            .await
1050            .map_err(|err| match err {
1051                CachedClientError::Callback { err, .. } => err,
1052                CachedClientError::Client(err) => Error::Client(err),
1053            })?;
1054
1055        if let (Some(expected), Some(actual)) = (expected_size, archive.size)
1056            && expected != actual
1057        {
1058            return Err(Error::MismatchedSize {
1059                distribution: dist.to_string(),
1060                expected,
1061                actual,
1062            });
1063        }
1064
1065        // If the archive is missing the required hashes or size, or has since been removed, force a refresh.
1066        let archive = Some(archive)
1067            .filter(|archive| archive.has_digests(hashes))
1068            .filter(|archive| archive.exists(self.build_context.cache()))
1069            .filter(|archive| expected_size.is_none() || archive.size.is_some());
1070
1071        let archive = if let Some(archive) = archive {
1072            archive
1073        } else {
1074            self.client
1075                .managed(async |client| {
1076                    client
1077                        .cached_client()
1078                        .skip_cache_with_retry(
1079                            self.request(url)?,
1080                            &http_entry,
1081                            cache_control,
1082                            download,
1083                        )
1084                        .await
1085                        .map_err(|err| match err {
1086                            CachedClientError::Callback { err, .. } => err,
1087                            CachedClientError::Client(err) => Error::Client(err),
1088                        })
1089                })
1090                .await?
1091        };
1092
1093        Ok(archive)
1094    }
1095
1096    /// Download and extract a wheel from an HTTP response.
1097    ///
1098    /// This helper is called by [`Self::download_wheel`] and handles partial content/resumptions
1099    /// of interrupted downloads if the origin supports it. Hard failures are propagated
1100    /// back to the caller and thus to our general retry machinery.
1101    ///
1102    /// The caller is responsible for obtaining a lock on the wheel cache.
1103    async fn download_wheel_response(
1104        &self,
1105        mut response: reqwest::Response,
1106        url: &SafeUrl,
1107        retry_state: &mut RetryState,
1108        filename: &WheelFilename,
1109        progress_size_hint: Option<u64>,
1110        expected_size: Option<u64>,
1111        wheel_entry: &CacheEntry,
1112        dist: &BuiltDist,
1113        hashes: ArchiveHashPolicy<'_>,
1114    ) -> Result<Archive, Error> {
1115        let hashes = if let BuiltDist::Registry(wheels) = dist {
1116            hashes.with_index_hashes(wheels.best_wheel().file.hashes.as_slice())
1117        } else {
1118            hashes
1119        };
1120
1121        let progress_size_hint = progress_size_hint.or_else(|| content_length(&response));
1122        let mut download_size = content_length(&response).or(expected_size);
1123
1124        let progress = self.reporter.as_ref().map(|reporter| {
1125            (
1126                reporter,
1127                reporter.on_download_start(dist.name(), progress_size_hint),
1128            )
1129        });
1130
1131        let algorithms = http_hash_algorithms(hashes);
1132
1133        // Download the wheel to a temporary file.
1134        let temp_file =
1135            tempfile::tempfile_in(self.build_context.cache().root()).map_err(Error::CacheWrite)?;
1136        let mut writer = tokio::io::BufWriter::new(fs_err::tokio::File::from_std(
1137            // It's an unnamed file on Linux so that's the best approximation.
1138            fs_err::File::from_parts(temp_file, self.build_context.cache().root()),
1139        ));
1140
1141        // States for the download loop below.
1142        let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
1143        // The total number of bytes retrieved, accumulated over individual requests.
1144        // Note that this is *not* the same as the number of bytes actually written
1145        // to the temporary file, since a copy from the response to the file can fail.
1146        let mut bytes_retrieved = 0;
1147        // The most recent range request's starting offset.
1148        let mut resumed_at = None;
1149        // The `Content-Range` that the download was last resumed at.
1150        let mut resumed_range: Option<ContentRangeBytes> = None;
1151
1152        // This loop is where we handle resumption of interrupted downloads, as well as
1153        // range requests (if the origin supports them).
1154        //
1155        // Errors returned from this loop reach the outer retry classifier, which may restart
1156        // the full download.
1157        loop {
1158            // Reject conflicting full-response lengths before reading the body. A range response's
1159            // Content-Length describes only that range, so it cannot be compared to the wheel size.
1160            if response.status() == reqwest::StatusCode::OK
1161                && let (Some(expected), Some(actual)) = (expected_size, content_length(&response))
1162                && expected != actual
1163            {
1164                return Err(Error::MismatchedContentLength {
1165                    distribution: dist.to_string(),
1166                    expected,
1167                    actual,
1168                });
1169            }
1170
1171            // Check whether the response indicates range request support. A `206 Partial Content`
1172            // implies range support while an `Accept-Ranges: bytes` header explicitly advertises it.
1173            let supports_range_requests = response.status() == reqwest::StatusCode::PARTIAL_CONTENT
1174                || response
1175                    .headers()
1176                    .get(reqwest::header::ACCEPT_RANGES)
1177                    .is_some_and(|value| value == "bytes");
1178
1179            // A server can advertise range requests but ignore one. In that case, the
1180            // response is a complete download and must replace the partial bytes.
1181            let replaces_partial_download =
1182                resumed_at.is_some() && response.status() != reqwest::StatusCode::PARTIAL_CONTENT;
1183            if replaces_partial_download {
1184                writer
1185                    .get_mut()
1186                    .set_len(0)
1187                    .await
1188                    .map_err(Error::CacheWrite)?;
1189                writer
1190                    .seek(io::SeekFrom::Start(0))
1191                    .await
1192                    .map_err(Error::CacheWrite)?;
1193                hashers = http_hash_algorithms(hashes)
1194                    .into_iter()
1195                    .map(Hasher::from)
1196                    .collect();
1197                bytes_retrieved = 0;
1198            }
1199
1200            let reader = response
1201                .bytes_stream()
1202                .map_err(|err| self.handle_response_errors(err))
1203                .into_async_read();
1204            let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
1205
1206            // Drain the response. This could be a partial response or a full one.
1207            // Note that the partial response here can take several forms: it can be an interrupted
1208            // request *or* it can be a `206 Partial Content`.
1209            let copy_result = match progress {
1210                Some((reporter, progress)) => {
1211                    // Wrap the reader in a progress reporter. This will report 100%
1212                    // progress once the download is complete, before the wheel is unzipped.
1213                    let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
1214
1215                    tokio::io::copy(&mut reader, &mut writer)
1216                        .await
1217                        .map_err(Error::CacheWrite)
1218                }
1219                None => tokio::io::copy(&mut hasher, &mut writer)
1220                    .await
1221                    .map_err(Error::CacheWrite),
1222            };
1223
1224            bytes_retrieved += hasher.bytes_read();
1225
1226            let interrupted = copy_result.is_err();
1227            let err = match copy_result {
1228                // Draining the response succeeded. However, the response could be a `206 Partial Content`,
1229                // so we can't assume that we're done.
1230                Ok(_) => {
1231                    // No `resumed_range` means this was a normal full response, so there's
1232                    // no resumption to do. We can leave the loop.
1233                    let Some(range) = resumed_range else {
1234                        break;
1235                    };
1236
1237                    // We should be in sync with the origin. Failure here means that the
1238                    // origin actually sent us an under- or over-length response. Treat this
1239                    // as a non-retryable error, since the HTTP body was received successfully.
1240                    //
1241                    // Byte ranges are inclusive, not exclusive.
1242                    if bytes_retrieved != range.last_byte + 1 {
1243                        return Err(Error::MismatchedRangeSize {
1244                            distribution: dist.to_string(),
1245                            expected: range.last_byte - range.first_byte + 1,
1246                            actual: hasher.bytes_read(),
1247                        });
1248                    }
1249
1250                    // We've successfully performed the download over one or more
1251                    // range requests. We can leave the loop.
1252                    if bytes_retrieved == range.complete_length {
1253                        break;
1254                    }
1255
1256                    // A successful range response may cover only part of the requested
1257                    // bytes. Keep requesting the remainder before extracting the wheel.
1258                    //
1259                    // Observe that we don't return this error; we bind it for handling
1260                    // below.
1261                    Error::CacheWrite(io::Error::new(
1262                        io::ErrorKind::UnexpectedEof,
1263                        "Range response did not complete the download",
1264                    ))
1265                }
1266                // Draining the response failed, potentially because the connection was interrupted.
1267                // Like above, we don't return this error until we know how we want to treat it.
1268                Err(err) => err,
1269            };
1270
1271            // At this point we have *some* error state: either a synthetic error (from a
1272            // partial response that did not complete the download) or a real error (e.g.
1273            // from an interrupted download). We can only resume if certain conditions
1274            // below are met.
1275
1276            // If the response replaces an already in-flight partial download or
1277            // indicates a lack of range-request support, then we can't resume.
1278            // Return the error back to the retry stack.
1279            if replaces_partial_download || !supports_range_requests {
1280                return Err(err);
1281            }
1282
1283            // Some sanity checks: we should have made *some* progress as part of
1284            // partial response, plus our writer's offset should agree with the number
1285            // of bytes retrieved, *and* we should be making forward progress (i.e. not
1286            // resuming behind where we already are).
1287            writer.flush().await.map_err(Error::CacheWrite)?;
1288            let offset = writer
1289                .get_mut()
1290                .stream_position()
1291                .await
1292                .map_err(Error::CacheWrite)?;
1293            if offset == 0
1294                || offset != bytes_retrieved
1295                || resumed_at.is_some_and(|previous| offset <= previous)
1296            {
1297                return Err(err);
1298            }
1299
1300            // Recovering from a failed body consumes the same budget as a full restart.
1301            // A successfully completed range needs no retry, even if more bytes remain.
1302            if interrupted {
1303                let Some(backoff) = retry_state.should_retry(&err, 0) else {
1304                    return Err(err);
1305                };
1306                retry_state.sleep_backoff(backoff).await;
1307            }
1308
1309            // Finally our resumption, which is a range request.
1310            debug!("Resuming download of `{url}` at byte {offset}");
1311            let resumed_response = retry_state
1312                .send(self.request_with_offset(url.clone(), offset))
1313                .await?;
1314
1315            // A chunked response can fail after all wheel bytes arrive, leaving no satisfiable
1316            // range. Return the original error so the outer retry policy can restart in full.
1317            if resumed_response.status() == reqwest::StatusCode::RANGE_NOT_SATISFIABLE {
1318                debug!("Range not satisfiable while resuming `{url}`; abandoning resumed download");
1319                return Err(err);
1320            }
1321            resumed_response.error_for_status_ref()?;
1322
1323            resumed_range = if resumed_response.status() == reqwest::StatusCode::PARTIAL_CONTENT {
1324                // The origin honored our range request, so we update `resumed_range`
1325                // for the next iteration.
1326                let Some(range) = content_range(&resumed_response, offset, download_size) else {
1327                    warn!(
1328                        "Invalid range request response from server that declares HTTP range \
1329                         request support, abandoning resumed download: {url}"
1330                    );
1331                    return Err(err);
1332                };
1333                download_size = Some(range.complete_length);
1334                Some(range)
1335            } else {
1336                // The origin is allowed to ignore our range request and send a full response instead.
1337                // That means we have no `resumed_range` to honor on the next iteration.
1338                None
1339            };
1340
1341            response = resumed_response;
1342            resumed_at = Some(offset);
1343        }
1344
1345        // We've left the resumption loop.
1346        // Sanity check: we should have written as many bytes as we expected.
1347        if let Some(expected) = expected_size
1348            && bytes_retrieved != expected
1349        {
1350            return Err(Error::MismatchedSize {
1351                distribution: dist.to_string(),
1352                expected,
1353                actual: bytes_retrieved,
1354            });
1355        }
1356
1357        // Unzip the wheel to a temporary directory.
1358        let extractor = WheelExtractor::new(
1359            self.build_context.cache().root(),
1360            self.content_addressed_cache,
1361        )
1362        .map_err(Error::CacheWrite)?;
1363        let mut file = writer.into_inner();
1364        file.seek(io::SeekFrom::Start(0))
1365            .await
1366            .map_err(Error::CacheWrite)?;
1367
1368        let file = file.into_std().await;
1369        let mut extracted = tokio::task::spawn_blocking(move || extractor.extract_seekable(file))
1370            .await?
1371            .map_err(|err| Error::Extract(filename.to_string(), err))?;
1372        let computed_hashes = Self::validate_hashes(dist, hashes, hashers)?;
1373
1374        // Before we make the wheel accessible by persisting it, ensure that the RECORD is
1375        // valid.
1376        extracted.validate_and_heal_record(dist)?;
1377
1378        // Persist the temporary directory to the directory store.
1379        let id = self
1380            .persist_extracted_wheel(extracted, wheel_entry.path())
1381            .await?;
1382
1383        if let Some((reporter, progress)) = progress {
1384            reporter.on_download_complete(dist.name(), progress);
1385        }
1386
1387        Ok(Archive::new(
1388            id,
1389            computed_hashes,
1390            filename.clone(),
1391            Some(bytes_retrieved),
1392        ))
1393    }
1394
1395    /// Load a wheel from a local path.
1396    async fn load_wheel(
1397        &self,
1398        path: &Path,
1399        filename: &WheelFilename,
1400        wheel_entry: CacheEntry,
1401        dist: &BuiltDist,
1402        hashes: ArchiveHashPolicy<'_>,
1403    ) -> Result<LocalWheel, Error> {
1404        // Acquire an advisory lock, to guard against concurrent writes.
1405        let _lock = Self::lock_wheel(&wheel_entry, filename).await?;
1406
1407        // Determine the last-modified time of the wheel.
1408        let modified = Timestamp::from_path(path).map_err(Error::CacheRead)?;
1409
1410        // Attempt to read the archive pointer from the cache.
1411        let pointer_entry = wheel_entry.with_file(format!("{}.rev", filename.cache_key()));
1412        let pointer = PathArchivePointer::read_from(&pointer_entry)?;
1413
1414        // Extract the archive from the pointer.
1415        let archive = pointer
1416            .filter(|pointer| pointer.is_up_to_date(modified))
1417            .map(PathArchivePointer::into_archive)
1418            .filter(|archive| archive.has_digests(hashes));
1419
1420        // Index hashes may replace `Generate`, but the caller still needs its SHA-256.
1421        let caller_hashes = hashes;
1422        let hashes = if let BuiltDist::Registry(wheels) = dist {
1423            hashes.with_index_hashes(wheels.best_wheel().file.hashes.as_slice())
1424        } else {
1425            hashes
1426        };
1427
1428        // If the file is already unzipped, and the cache is up-to-date, return it.
1429        if let Some(archive) = archive {
1430            Ok(LocalWheel {
1431                dist: Dist::Built(dist.clone()),
1432                archive: self
1433                    .build_context
1434                    .cache()
1435                    .archive(&archive.id)
1436                    .into_boxed_path(),
1437                hashes: archive.hashes,
1438                filename: filename.clone(),
1439                cache: CacheInfo::from_timestamp(modified),
1440                build: None,
1441            })
1442        } else if hashes.is_none() {
1443            // Otherwise, unzip the wheel.
1444            let archive = Archive::new(
1445                self.unzip_wheel(path, wheel_entry.path(), DistRef::Built(dist))
1446                    .await?,
1447                HashDigests::empty(),
1448                filename.clone(),
1449                None,
1450            );
1451
1452            // Write the archive pointer to the cache.
1453            let pointer = PathArchivePointer {
1454                timestamp: modified,
1455                archive: archive.clone(),
1456            };
1457            pointer.write_to(&pointer_entry).await?;
1458
1459            Ok(LocalWheel {
1460                dist: Dist::Built(dist.clone()),
1461                archive: self
1462                    .build_context
1463                    .cache()
1464                    .archive(&archive.id)
1465                    .into_boxed_path(),
1466                hashes: archive.hashes,
1467                filename: filename.clone(),
1468                cache: CacheInfo::from_timestamp(modified),
1469                build: None,
1470            })
1471        } else {
1472            // If necessary, compute the hashes of the wheel.
1473            let file = fs_err::tokio::File::open(path)
1474                .await
1475                .map_err(Error::CacheRead)?;
1476            let extractor = WheelExtractor::new(
1477                self.build_context.cache().root(),
1478                self.content_addressed_cache,
1479            )
1480            .map_err(Error::CacheWrite)?;
1481
1482            // Include the caller's algorithms alongside any index algorithms selected for validation.
1483            let mut algorithms = caller_hashes.algorithms();
1484            algorithms.extend(hashes.algorithms());
1485            algorithms.sort();
1486            algorithms.dedup();
1487            let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
1488            let mut hasher = uv_extract::hash::HashReader::new(file, &mut hashers);
1489
1490            // Unzip the wheel to a temporary directory.
1491            let mut extracted = extractor
1492                .extract_streaming(&mut hasher)
1493                .await
1494                .map_err(|err| Error::Extract(filename.to_string(), err))?;
1495
1496            // Exhaust the reader to compute the hash.
1497            hasher.finish().await.map_err(Error::HashExhaustion)?;
1498
1499            let computed_hashes = Self::validate_hashes(dist, hashes, hashers)?;
1500
1501            // Before we make the wheel accessible by persisting it, ensure that the RECORD is
1502            // valid.
1503            extracted.validate_and_heal_record(dist)?;
1504
1505            // Persist the temporary directory to the directory store.
1506            let id = self
1507                .persist_extracted_wheel(extracted, wheel_entry.path())
1508                .await?;
1509
1510            // Create an archive.
1511            let archive = Archive::new(id, computed_hashes, filename.clone(), None);
1512
1513            // Write the archive pointer to the cache.
1514            let pointer = PathArchivePointer {
1515                timestamp: modified,
1516                archive: archive.clone(),
1517            };
1518            pointer.write_to(&pointer_entry).await?;
1519
1520            Ok(LocalWheel {
1521                dist: Dist::Built(dist.clone()),
1522                archive: self
1523                    .build_context
1524                    .cache()
1525                    .archive(&archive.id)
1526                    .into_boxed_path(),
1527                hashes: archive.hashes,
1528                filename: filename.clone(),
1529                cache: CacheInfo::from_timestamp(modified),
1530                build: None,
1531            })
1532        }
1533    }
1534
1535    /// Unzip a wheel into the cache, returning the path to the unzipped directory.
1536    async fn unzip_wheel(
1537        &self,
1538        path: &Path,
1539        target: &Path,
1540        dist: DistRef<'_>,
1541    ) -> Result<ArchiveId, Error> {
1542        let content_addressed_cache = self.content_addressed_cache;
1543
1544        let mut extracted = tokio::task::spawn_blocking({
1545            let path = path.to_owned();
1546            let root = self.build_context.cache().root().to_path_buf();
1547            move || -> Result<_, Error> {
1548                // Unzip the wheel into a temporary directory.
1549                let extractor = WheelExtractor::new(&root, content_addressed_cache)
1550                    .map_err(Error::CacheWrite)?;
1551                let reader = fs_err::File::open(&path).map_err(Error::CacheWrite)?;
1552                extractor
1553                    .extract_seekable(reader)
1554                    .map_err(|err| Error::Extract(path.to_string_lossy().into_owned(), err))
1555            }
1556        })
1557        .await??;
1558
1559        // Before we make the wheel accessible by persisting it, ensure that the RECORD is valid.
1560        extracted.validate_and_heal_record(dist)?;
1561
1562        // Persist the temporary directory to the directory store.
1563        let id = self.persist_extracted_wheel(extracted, target).await?;
1564
1565        Ok(id)
1566    }
1567
1568    /// Persist an extracted wheel into the archive store.
1569    ///
1570    /// A hash tree makes identical extracted trees converge on one archive entry. Without one,
1571    /// persistence retains the existing behavior of assigning a unique archive ID.
1572    async fn persist_extracted_wheel(
1573        &self,
1574        extracted: ExtractedWheel,
1575        target: &Path,
1576    ) -> Result<ArchiveId, Error> {
1577        let (temp_dir, hashed_wheel) = extracted.into_parts();
1578        let cache = self.build_context.cache();
1579        let (temp_dir, id) = if let Some(HashedWheel { files, tree }) = hashed_wheel {
1580            let digest = DirectoryDigest::from(tree.hash());
1581            let id = ArchiveId::from_digest(digest.into());
1582            let cache = cache.clone();
1583            let temp_dir = tokio::task::spawn_blocking(move || {
1584                persist_archive_files(&cache, temp_dir.path(), &files)
1585                    .map_err(Error::CacheWrite)?;
1586                Ok::<_, Error>(temp_dir)
1587            })
1588            .await??;
1589            (temp_dir, id)
1590        } else {
1591            (temp_dir, ArchiveId::default())
1592        };
1593
1594        cache
1595            .persist_with_id(temp_dir, target, id)
1596            .await
1597            .map_err(Error::CacheWrite)
1598    }
1599
1600    /// Returns a GET [`reqwest::Request`] for the given URL.
1601    fn request(&self, url: SafeUrl) -> Result<reqwest::Request, reqwest::Error> {
1602        self.client
1603            .unmanaged
1604            .uncached_client(&url)
1605            .get(url)
1606            .header(
1607                // `reqwest` defaults to accepting compressed responses.
1608                // Specify identity encoding to get consistent .whl downloading
1609                // behavior from servers. ref: https://github.com/pypa/pip/pull/1688
1610                "accept-encoding",
1611                reqwest::header::HeaderValue::from_static("identity"),
1612            )
1613            .build()
1614    }
1615
1616    /// Build a GET request with a `Range: bytes=<offset>-` header.
1617    ///
1618    /// Used to resume an interrupted download from `offset` bytes into the file.
1619    fn request_with_offset(&self, url: SafeUrl, offset: u64) -> RequestBuilder<'_> {
1620        self.client
1621            .unmanaged
1622            .uncached_client(&url)
1623            .get(url)
1624            .header(
1625                "accept-encoding",
1626                reqwest::header::HeaderValue::from_static("identity"),
1627            )
1628            .header(reqwest::header::RANGE, format!("bytes={offset}-"))
1629    }
1630
1631    /// Return the [`ManagedClient`] used by this resolver.
1632    pub fn client(&self) -> &ManagedClient<'a> {
1633        &self.client
1634    }
1635}
1636
1637/// Share extracted files other than `RECORD` while keeping the unpublished archive complete.
1638fn persist_archive_files(cache: &Cache, archive: &Path, files: &[HashedFile]) -> io::Result<()> {
1639    initialize_rayon_once();
1640    let targets = files
1641        .par_iter()
1642        // Keep RECORD private, since it may have been healed after hashing.
1643        .filter(|file| !file.path().ends_with("RECORD"))
1644        .map(|file| {
1645            let id = ArchiveFileId::from_digest(&file.object_digest_hex());
1646            (archive.join(file.path()), cache.archive_file(&id))
1647        })
1648        .collect::<Vec<_>>();
1649
1650    // Group files by shard so its directory is created once and its files are linked by the
1651    // same worker, avoiding contention between workers on each shard directory.
1652    let mut shards: FxHashMap<&Path, Vec<_>> = FxHashMap::default();
1653    for (source, target) in &targets {
1654        let Some(parent) = target.parent() else {
1655            return Err(io::Error::new(
1656                io::ErrorKind::InvalidInput,
1657                "archive file path must have a parent directory",
1658            ));
1659        };
1660        shards.entry(parent).or_default().push((source, target));
1661    }
1662
1663    let mut shards = shards
1664        .into_iter()
1665        .map(|(parent, files)| (parent, files, Ok(())))
1666        .collect::<Vec<_>>();
1667    // Start larger shards first so their work can overlap the remaining directory creation.
1668    shards.sort_unstable_by_key(|(_, files, _)| Reverse(files.len()));
1669
1670    // Creating shards concurrently contends on their shared parent. Keep creation on this
1671    // thread, while workers link files in the shards that are already available.
1672    in_place_scope(|scope| -> io::Result<()> {
1673        for (parent, files, result) in &mut shards {
1674            fs_err::create_dir_all(parent)?;
1675            scope.spawn(move |_| {
1676                *result = files
1677                    .iter()
1678                    .try_for_each(|(source, target)| persist_archive_file(source, target));
1679            });
1680        }
1681        Ok(())
1682    })?;
1683
1684    shards.into_iter().try_for_each(|(_, _, result)| result)
1685}
1686
1687/// Publish a shared object and retain a hardlink in the archive, with a copy fallback.
1688fn persist_archive_file(src: &Path, dst: &Path) -> io::Result<()> {
1689    // The shard already exists, and most objects are new, so try linking before checking for an
1690    // existing object. This avoids an extra filesystem lookup for every new object.
1691    match fs_err::hard_link(src, dst) {
1692        Ok(()) => return Ok(()),
1693        Err(_) if dst.try_exists()? => {}
1694        Err(_) => return uv_fs::copy_atomic_sync(src, dst),
1695    }
1696
1697    // This archive is still private, so it is safe to replace its extracted copy before publication.
1698    if let Err(err) = fs_err::remove_file(src)
1699        && err.kind() != io::ErrorKind::NotFound
1700    {
1701        return Err(err);
1702    }
1703
1704    fs_err::hard_link(dst, src).or_else(|_| uv_fs::copy_atomic_sync(dst, src))
1705}
1706
1707/// A wrapper around `RegistryClient` that manages a concurrency limit.
1708pub struct ManagedClient<'a> {
1709    pub(super) unmanaged: &'a RegistryClient,
1710    control: Arc<Semaphore>,
1711}
1712
1713impl<'a> ManagedClient<'a> {
1714    /// Create a new `ManagedClient` using the given client and concurrency semaphore.
1715    fn new(client: &'a RegistryClient, control: Arc<Semaphore>) -> Self {
1716        ManagedClient {
1717            unmanaged: client,
1718            control,
1719        }
1720    }
1721
1722    /// Return whether the client is configured for offline use.
1723    pub fn is_offline(&self) -> bool {
1724        self.unmanaged.connectivity().is_offline()
1725    }
1726
1727    /// Perform a request using the client, respecting the concurrency limit.
1728    ///
1729    /// If the concurrency limit has been reached, this method will wait until a pending
1730    /// operation completes before executing the closure.
1731    pub async fn managed<F, T>(&self, f: impl FnOnce(&'a RegistryClient) -> F) -> T
1732    where
1733        F: Future<Output = T>,
1734    {
1735        let _permit = self.control.acquire().await.unwrap();
1736        f(self.unmanaged).await
1737    }
1738
1739    /// Perform a request using a client that internally manages the concurrency limit.
1740    ///
1741    /// The callback is passed the client and a semaphore. It must acquire the semaphore before
1742    /// any request through the client and drop it after.
1743    ///
1744    /// This method serves as an escape hatch for functions that may want to send multiple requests
1745    /// in parallel.
1746    pub async fn manual<F, T>(&'a self, f: impl FnOnce(&'a RegistryClient, &'a Semaphore) -> F) -> T
1747    where
1748        F: Future<Output = T>,
1749    {
1750        f(self.unmanaged, &self.control).await
1751    }
1752}
1753
1754/// Returns the value of the `Content-Length` header from the [`reqwest::Response`], if present.
1755fn content_length(response: &reqwest::Response) -> Option<u64> {
1756    response
1757        .headers()
1758        .get(reqwest::header::CONTENT_LENGTH)
1759        .and_then(|val| val.to_str().ok())
1760        .and_then(|val| val.parse::<u64>().ok())
1761}
1762
1763/// Return the bounds of a range response starting at `offset` with a known complete length.
1764fn content_range(
1765    response: &reqwest::Response,
1766    offset: u64,
1767    download_size: Option<u64>,
1768) -> Option<ContentRangeBytes> {
1769    let range = response
1770        .headers()
1771        .get(reqwest::header::CONTENT_RANGE)
1772        .and_then(|value| value.to_str().ok())
1773        .and_then(ContentRange::parse)?;
1774
1775    let range = match range {
1776        ContentRange::Bytes(range) => range,
1777        ContentRange::UnboundBytes(ContentRangeUnbound {
1778            first_byte,
1779            last_byte,
1780        }) => ContentRangeBytes {
1781            first_byte,
1782            last_byte,
1783            complete_length: download_size?,
1784        },
1785        ContentRange::Unsatisfied(_) => return None,
1786    };
1787
1788    if range.first_byte != offset
1789        || range.last_byte >= range.complete_length
1790        || download_size.is_some_and(|size| size != range.complete_length)
1791    {
1792        return None;
1793    }
1794
1795    Some(range)
1796}
1797
1798/// An asynchronous reader that reports progress as bytes are read.
1799struct ProgressReader<'a, R> {
1800    reader: R,
1801    index: usize,
1802    reporter: &'a dyn Reporter,
1803}
1804
1805impl<'a, R> ProgressReader<'a, R> {
1806    /// Create a new [`ProgressReader`] that wraps another reader.
1807    fn new(reader: R, index: usize, reporter: &'a dyn Reporter) -> Self {
1808        Self {
1809            reader,
1810            index,
1811            reporter,
1812        }
1813    }
1814}
1815
1816impl<R> AsyncRead for ProgressReader<'_, R>
1817where
1818    R: AsyncRead + Unpin,
1819{
1820    fn poll_read(
1821        mut self: Pin<&mut Self>,
1822        cx: &mut Context<'_>,
1823        buf: &mut ReadBuf<'_>,
1824    ) -> Poll<io::Result<()>> {
1825        Pin::new(&mut self.as_mut().reader)
1826            .poll_read(cx, buf)
1827            .map_ok(|()| {
1828                self.reporter
1829                    .on_download_progress(self.index, buf.filled().len() as u64);
1830            })
1831    }
1832}
1833
1834/// A pointer to an archive in the cache, fetched from an HTTP archive.
1835///
1836/// Encoded with `MsgPack`, and represented on disk by a `.http` file.
1837#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1838pub struct HttpArchivePointer {
1839    archive: Archive,
1840}
1841
1842impl HttpArchivePointer {
1843    /// Read an [`HttpArchivePointer`] from the cache.
1844    pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1845        match fs_err::File::open(path.as_ref()) {
1846            Ok(file) => {
1847                let data = DataWithCachePolicy::from_reader(file)?.into_data();
1848                let archive = rmp_serde::from_slice::<Archive>(&data)?;
1849                Ok(Some(Self { archive }))
1850            }
1851            Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1852            Err(err) => Err(Error::CacheRead(err)),
1853        }
1854    }
1855
1856    /// Return the [`Archive`] from the pointer.
1857    pub fn into_archive(self) -> Archive {
1858        self.archive
1859    }
1860
1861    /// Return the [`CacheInfo`] from the pointer.
1862    pub fn to_cache_info(&self) -> CacheInfo {
1863        CacheInfo::default()
1864    }
1865
1866    /// Return the [`BuildInfo`] from the pointer.
1867    pub fn to_build_info(&self) -> Option<BuildInfo> {
1868        None
1869    }
1870}
1871
1872/// A pointer to an archive in the cache, fetched from a local path.
1873///
1874/// Encoded with `MsgPack`, and represented on disk by a `.rev` file.
1875#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1876pub struct PathArchivePointer {
1877    timestamp: Timestamp,
1878    archive: Archive,
1879}
1880
1881impl PathArchivePointer {
1882    /// Read an [`PathArchivePointer`] from the cache.
1883    pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1884        match fs_err::read(path) {
1885            Ok(cached) => Ok(Some(rmp_serde::from_slice::<Self>(&cached)?)),
1886            Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1887            Err(err) => Err(Error::CacheRead(err)),
1888        }
1889    }
1890
1891    /// Write an [`PathArchivePointer`] to the cache.
1892    async fn write_to(&self, entry: &CacheEntry) -> Result<(), Error> {
1893        write_atomic(entry.path(), rmp_serde::to_vec(&self)?)
1894            .await
1895            .map_err(Error::CacheWrite)
1896    }
1897
1898    /// Returns `true` if the archive is up-to-date with the given modified timestamp.
1899    pub fn is_up_to_date(&self, modified: Timestamp) -> bool {
1900        self.timestamp == modified
1901    }
1902
1903    /// Return the [`Archive`] from the pointer.
1904    pub fn into_archive(self) -> Archive {
1905        self.archive
1906    }
1907
1908    /// Return the [`CacheInfo`] from the pointer.
1909    pub fn to_cache_info(&self) -> CacheInfo {
1910        CacheInfo::from_timestamp(self.timestamp)
1911    }
1912
1913    /// Return the [`BuildInfo`] from the pointer.
1914    pub fn to_build_info(&self) -> Option<BuildInfo> {
1915        None
1916    }
1917}