Backtrack on distributions with invalid metadata (#2834)

charliermarsh · web-flow · commit 00934044aaa6 · 2024-04-05T18:00:48.000-04:00
## Summary Closes #2821.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/uv-client/src/error.rs b/crates/uv-client/src/error.rs
@@ -132,7 +132,7 @@ pub enum ErrorKind {
 
     /// Dist-info error
     #[error(transparent)]
-    InstallWheel(#[from] install_wheel_rs::Error),
+    DistInfo(#[from] install_wheel_rs::Error),
 
     #[error("{0} isn't available locally, but making network requests to registries was banned.")]
     NoIndex(String),
diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs
@@ -632,7 +632,7 @@ async fn read_metadata_async_seek(
             .enumerate()
             .filter_map(|(index, entry)| Some((index, entry.filename().as_str().ok()?))),
     )
-    .map_err(ErrorKind::InstallWheel)?;
+    .map_err(ErrorKind::DistInfo)?;
 
     // Read the contents of the `METADATA` file.
     let mut contents = Vec::new();
diff --git a/crates/uv-client/src/remote_metadata.rs b/crates/uv-client/src/remote_metadata.rs
@@ -75,7 +75,7 @@ pub(crate) async fn wheel_metadata_from_remote_zip(
             .enumerate()
             .filter_map(|(idx, e)| Some(((idx, e), e.filename().as_str().ok()?))),
     )
-    .map_err(ErrorKind::InstallWheel)?;
+    .map_err(ErrorKind::DistInfo)?;
 
     let offset = metadata_entry.header_offset();
     let size = metadata_entry.compressed_size()
diff --git a/crates/uv-requirements/src/lookahead.rs b/crates/uv-requirements/src/lookahead.rs
@@ -11,7 +11,7 @@ use pep508_rs::{MarkerEnvironment, Requirement, VersionOrUrl};
 use pypi_types::Metadata23;
 use uv_client::RegistryClient;
 use uv_distribution::{DistributionDatabase, Reporter};
-use uv_resolver::InMemoryIndex;
+use uv_resolver::{InMemoryIndex, MetadataResponse};
 use uv_types::{BuildContext, Constraints, Overrides, RequestedRequirements};
 
 /// A resolver for resolving lookahead requirements from direct URLs.
@@ -134,7 +134,18 @@ impl<'a, Context: BuildContext + Send + Sync> LookaheadResolver<'a, Context> {
         // Fetch the metadata for the distribution.
         let requires_dist = {
             let id = dist.package_id();
-            if let Some(metadata) = self.index.get_metadata(&id) {
+            if let Some(metadata) = self
+                .index
+                .get_metadata(&id)
+                .as_deref()
+                .and_then(|response| {
+                    if let MetadataResponse::Found(metadata) = response {
+                        Some(metadata)
+                    } else {
+                        None
+                    }
+                })
+            {
                 // If the metadata is already in the index, return it.
                 metadata.requires_dist.clone()
             } else {
@@ -151,7 +162,8 @@ impl<'a, Context: BuildContext + Send + Sync> LookaheadResolver<'a, Context> {
                 let requires_dist = metadata.requires_dist.clone();
 
                 // Insert the metadata into the index.
-                self.index.insert_metadata(id, metadata);
+                self.index
+                    .insert_metadata(id, MetadataResponse::Found(metadata));
 
                 requires_dist
             }
diff --git a/crates/uv-requirements/src/source_tree.rs b/crates/uv-requirements/src/source_tree.rs
@@ -1,5 +1,5 @@
 use std::borrow::Cow;
-use std::ops::Deref;
+
 use std::path::{Path, PathBuf};
 
 use anyhow::{Context, Result};
@@ -11,7 +11,7 @@ use pep508_rs::Requirement;
 use uv_client::RegistryClient;
 use uv_distribution::{DistributionDatabase, Reporter};
 use uv_fs::Simplified;
-use uv_resolver::InMemoryIndex;
+use uv_resolver::{InMemoryIndex, MetadataResponse};
 use uv_types::BuildContext;
 
 use crate::ExtrasSpecification;
@@ -87,16 +87,28 @@ impl<'a, Context: BuildContext + Send + Sync> SourceTreeResolver<'a, Context> {
         // Fetch the metadata for the distribution.
         let metadata = {
             let id = PackageId::from_url(source.url());
-            if let Some(metadata) = self.index.get_metadata(&id) {
+            if let Some(metadata) = self
+                .index
+                .get_metadata(&id)
+                .as_deref()
+                .and_then(|response| {
+                    if let MetadataResponse::Found(metadata) = response {
+                        Some(metadata)
+                    } else {
+                        None
+                    }
+                })
+            {
                 // If the metadata is already in the index, return it.
-                metadata.deref().clone()
+                metadata.clone()
             } else {
                 // Run the PEP 517 build process to extract metadata from the source distribution.
                 let source = BuildableSource::Url(source);
                 let metadata = self.database.build_wheel_metadata(&source).await?;
 
                 // Insert the metadata into the index.
-                self.index.insert_metadata(id, metadata.clone());
+                self.index
+                    .insert_metadata(id, MetadataResponse::Found(metadata.clone()));
 
                 metadata
             }
diff --git a/crates/uv-requirements/src/unnamed.rs b/crates/uv-requirements/src/unnamed.rs
@@ -20,7 +20,7 @@ use pypi_types::Metadata10;
 use uv_client::RegistryClient;
 use uv_distribution::{DistributionDatabase, Reporter};
 use uv_normalize::PackageName;
-use uv_resolver::InMemoryIndex;
+use uv_resolver::{InMemoryIndex, MetadataResponse};
 use uv_types::BuildContext;
 
 /// Like [`RequirementsSpecification`], but with concrete names for all requirements.
@@ -236,7 +236,13 @@ impl<'a, Context: BuildContext + Send + Sync> NamedRequirementsResolver<'a, Cont
         // Fetch the metadata for the distribution.
         let name = {
             let id = PackageId::from_url(source.url());
-            if let Some(metadata) = index.get_metadata(&id) {
+            if let Some(metadata) = index.get_metadata(&id).as_deref().and_then(|response| {
+                if let MetadataResponse::Found(metadata) = response {
+                    Some(metadata)
+                } else {
+                    None
+                }
+            }) {
                 // If the metadata is already in the index, return it.
                 metadata.name.clone()
             } else {
@@ -247,7 +253,7 @@ impl<'a, Context: BuildContext + Send + Sync> NamedRequirementsResolver<'a, Cont
                 let name = metadata.name.clone();
 
                 // Insert the metadata into the index.
-                index.insert_metadata(id, metadata);
+                index.insert_metadata(id, MetadataResponse::Found(metadata));
 
                 name
             }
diff --git a/crates/uv-resolver/Cargo.toml b/crates/uv-resolver/Cargo.toml
@@ -16,6 +16,7 @@ workspace = true
 cache-key = { workspace = true }
 distribution-filename = { workspace = true, features = ["serde"] }
 distribution-types = { workspace = true }
+install-wheel-rs = { workspace = true }
 once-map = { workspace = true }
 pep440_rs = { workspace = true }
 pep508_rs = { workspace = true }
diff --git a/crates/uv-resolver/src/lib.rs b/crates/uv-resolver/src/lib.rs
@@ -9,7 +9,7 @@ pub use python_requirement::PythonRequirement;
 pub use resolution::{AnnotationStyle, Diagnostic, DisplayResolutionGraph, ResolutionGraph};
 pub use resolution_mode::ResolutionMode;
 pub use resolver::{
-    BuildId, DefaultResolverProvider, InMemoryIndex, PackageVersionsResult,
+    BuildId, DefaultResolverProvider, InMemoryIndex, MetadataResponse, PackageVersionsResult,
     Reporter as ResolverReporter, Resolver, ResolverProvider, VersionsResponse,
     WheelMetadataResult,
 };
diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs
@@ -18,7 +18,7 @@ use distribution_types::{
 use once_map::OnceMap;
 use pep440_rs::Version;
 use pep508_rs::MarkerEnvironment;
-use pypi_types::{Hashes, Metadata23};
+use pypi_types::Hashes;
 use uv_distribution::to_precise;
 use uv_normalize::{ExtraName, PackageName};
 
@@ -28,7 +28,7 @@ use crate::pins::FilePins;
 use crate::preferences::Preferences;
 use crate::pubgrub::{PubGrubDistribution, PubGrubPackage};
 use crate::redirect::apply_redirect;
-use crate::resolver::{InMemoryIndex, VersionsResponse};
+use crate::resolver::{InMemoryIndex, MetadataResponse, VersionsResponse};
 use crate::{Manifest, ResolveError};
 
 /// Indicate the style of annotation comments, used to indicate the dependencies that requested each
@@ -66,7 +66,7 @@ impl ResolutionGraph {
         selection: &SelectedDependencies<UvDependencyProvider>,
         pins: &FilePins,
         packages: &OnceMap<PackageName, VersionsResponse>,
-        distributions: &OnceMap<PackageId, Metadata23>,
+        distributions: &OnceMap<PackageId, MetadataResponse>,
         state: &State<UvDependencyProvider>,
         preferences: &Preferences,
         editables: Editables,
@@ -164,13 +164,20 @@ impl ResolutionGraph {
                             });
                         }
                     } else {
-                        let metadata = distributions.get(&dist.package_id()).unwrap_or_else(|| {
+                        let response = distributions.get(&dist.package_id()).unwrap_or_else(|| {
                             panic!(
                                 "Every package should have metadata: {:?}",
                                 dist.package_id()
                             )
                         });
 
+                        let MetadataResponse::Found(metadata) = &*response else {
+                            panic!(
+                                "Every package should have metadata: {:?}",
+                                dist.package_id()
+                            )
+                        };
+
                         if metadata.provides_extras.contains(extra) {
                             extras
                                 .entry(package_name.clone())
@@ -211,13 +218,20 @@ impl ResolutionGraph {
                             });
                         }
                     } else {
-                        let metadata = distributions.get(&dist.package_id()).unwrap_or_else(|| {
+                        let response = distributions.get(&dist.package_id()).unwrap_or_else(|| {
                             panic!(
                                 "Every package should have metadata: {:?}",
                                 dist.package_id()
                             )
                         });
 
+                        let MetadataResponse::Found(metadata) = &*response else {
+                            panic!(
+                                "Every package should have metadata: {:?}",
+                                dist.package_id()
+                            )
+                        };
+
                         if metadata.provides_extras.contains(extra) {
                             extras
                                 .entry(package_name.clone())
@@ -417,10 +431,16 @@ impl ResolutionGraph {
                 }
                 VersionOrUrl::Url(verbatim_url) => PackageId::from_url(verbatim_url.raw()),
             };
-            let md = index
+            let res = index
                 .distributions
                 .get(&package_id)
                 .expect("every package in resolution graph has metadata");
+            let MetadataResponse::Found(md) = &*res else {
+                panic!(
+                    "Every package should have metadata: {:?}",
+                    dist.package_id()
+                )
+            };
             for req in manifest.apply(&md.requires_dist) {
                 let Some(ref marker_tree) = req.marker else {
                     continue;
diff --git a/crates/uv-resolver/src/resolver/index.rs b/crates/uv-resolver/src/resolver/index.rs
@@ -2,10 +2,9 @@ use std::sync::Arc;
 
 use distribution_types::PackageId;
 use once_map::OnceMap;
-use pypi_types::Metadata23;
 use uv_normalize::PackageName;
 
-use crate::resolver::provider::VersionsResponse;
+use crate::resolver::provider::{MetadataResponse, VersionsResponse};
 
 /// In-memory index of package metadata.
 #[derive(Default)]
@@ -15,27 +14,27 @@ pub struct InMemoryIndex {
     pub(crate) packages: OnceMap<PackageName, VersionsResponse>,
 
     /// A map from package ID to metadata for that distribution.
-    pub(crate) distributions: OnceMap<PackageId, Metadata23>,
+    pub(crate) distributions: OnceMap<PackageId, MetadataResponse>,
 }
 
 impl InMemoryIndex {
     /// Insert a [`VersionsResponse`] into the index.
-    pub fn insert_package(&self, package_name: PackageName, metadata: VersionsResponse) {
-        self.packages.done(package_name, metadata);
+    pub fn insert_package(&self, package_name: PackageName, response: VersionsResponse) {
+        self.packages.done(package_name, response);
     }
 
     /// Insert a [`Metadata23`] into the index.
-    pub fn insert_metadata(&self, package_id: PackageId, metadata: Metadata23) {
-        self.distributions.done(package_id, metadata);
+    pub fn insert_metadata(&self, package_id: PackageId, response: MetadataResponse) {
+        self.distributions.done(package_id, response);
     }
 
     /// Get the [`VersionsResponse`] for a given package name, without waiting.
     pub fn get_package(&self, package_name: &PackageName) -> Option<Arc<VersionsResponse>> {
         self.packages.get(package_name)
     }
 
-    /// Get the [`Metadata23`] for a given package ID, without waiting.
-    pub fn get_metadata(&self, package_id: &PackageId) -> Option<Arc<Metadata23>> {
+    /// Get the [`MetadataResponse`] for a given package ID, without waiting.
+    pub fn get_metadata(&self, package_id: &PackageId) -> Option<Arc<MetadataResponse>> {
         self.distributions.get(package_id)
     }
 }
diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
diff --git a/crates/uv-resolver/src/resolver/provider.rs b/crates/uv-resolver/src/resolver/provider.rs
diff --git a/crates/uv/tests/pip_compile.rs b/crates/uv/tests/pip_compile.rs
diff --git a/crates/uv/tests/pip_sync.rs b/crates/uv/tests/pip_sync.rs

Original file line number	Diff line number	Diff line change
`@@ -632,7 +632,7 @@ async fn read_metadata_async_seek(`
`632`	`632`	`.enumerate()`
`633`	`633`	`.filter_map(\|(index, entry)\| Some((index, entry.filename().as_str().ok()?))),`
`634`	`634`	`)`
`635`		`- .map_err(ErrorKind::InstallWheel)?;`
	`635`	`+ .map_err(ErrorKind::DistInfo)?;`
`636`	`636`
`637`	`637`	// Read the contents of the `METADATA` file.
`638`	`638`	`let mut contents = Vec::new();`
Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ pub(crate) async fn wheel_metadata_from_remote_zip(`
`75`	`75`	`.enumerate()`
`76`	`76`	`.filter_map(\|(idx, e)\| Some(((idx, e), e.filename().as_str().ok()?))),`
`77`	`77`	`)`
`78`		`- .map_err(ErrorKind::InstallWheel)?;`
	`78`	`+ .map_err(ErrorKind::DistInfo)?;`
`79`	`79`
`80`	`80`	`let offset = metadata_entry.header_offset();`
`81`	`81`	`let size = metadata_entry.compressed_size()`