Key hash policy on version, rather than package

astral-sh · Jul 17, 2024 · fa0bb41 · fa0bb41
1 parent 82d9483
commit fa0bb41
Show file tree

Hide file tree

Showing 7 changed files with 97 additions and 62 deletions.
diff --git a/crates/uv-distribution/src/index/registry_wheel_index.rs b/crates/uv-distribution/src/index/registry_wheel_index.rs
@@ -114,7 +114,10 @@ impl<'a> RegistryWheelIndex<'a> {
                                 CachedWheel::from_http_pointer(wheel_dir.join(file), cache)
                             {
                                 // Enforce hash-checking based on the built distribution.
-                                if wheel.satisfies(hasher.get_package(package)) {
+                                if wheel.satisfies(
+                                    hasher
+                                        .get_package(&wheel.filename.name, &wheel.filename.version),
+                                ) {
                                     Self::add_wheel(wheel, tags, &mut versions);
                                 }
                             }
@@ -130,7 +133,10 @@ impl<'a> RegistryWheelIndex<'a> {
                                 CachedWheel::from_local_pointer(wheel_dir.join(file), cache)
                             {
                                 // Enforce hash-checking based on the built distribution.
-                                if wheel.satisfies(hasher.get_package(package)) {
+                                if wheel.satisfies(
+                                    hasher
+                                        .get_package(&wheel.filename.name, &wheel.filename.version),
+                                ) {
                                     Self::add_wheel(wheel, tags, &mut versions);
                                 }
                             }
@@ -174,10 +180,12 @@ impl<'a> RegistryWheelIndex<'a> {
                 };
 
                 if let Some(revision) = revision {
-                    // Enforce hash-checking based on the source distribution.
-                    if revision.satisfies(hasher.get_package(package)) {
-                        for wheel_dir in symlinks(cache_shard.join(revision.id())) {
-                            if let Some(wheel) = CachedWheel::from_built_source(wheel_dir) {
+                    for wheel_dir in symlinks(cache_shard.join(revision.id())) {
+                        if let Some(wheel) = CachedWheel::from_built_source(wheel_dir) {
+                            // Enforce hash-checking based on the source distribution.
+                            if revision.satisfies(
+                                hasher.get_package(&wheel.filename.name, &wheel.filename.version),
+                            ) {
                                 Self::add_wheel(wheel, tags, &mut versions);
                             }
                         }

diff --git a/crates/uv-resolver/src/flat_index.rs b/crates/uv-resolver/src/flat_index.rs
@@ -124,7 +124,9 @@ impl FlatIndex {
         }
 
         // Check if hashes line up
-        let hash = if let HashPolicy::Validate(required) = hasher.get_package(&filename.name) {
+        let hash = if let HashPolicy::Validate(required) =
+            hasher.get_package(&filename.name, &filename.version)
+        {
             if hashes.is_empty() {
                 HashComparison::Missing
             } else if required.iter().any(|hash| hashes.contains(hash)) {
@@ -163,7 +165,9 @@ impl FlatIndex {
         };
 
         // Check if hashes line up.
-        let hash = if let HashPolicy::Validate(required) = hasher.get_package(&filename.name) {
+        let hash = if let HashPolicy::Validate(required) =
+            hasher.get_package(&filename.name, &filename.version)
+        {
             if hashes.is_empty() {
                 HashComparison::Missing
             } else if required.iter().any(|hash| hashes.contains(hash)) {

diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
@@ -722,11 +722,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                 request_sink.blocking_send(Request::Dist(dist))?;
             }
         } else {
-            // Verify that the package is allowed under the hash-checking policy.
-            if !self.hasher.allows_package(name) {
-                return Err(ResolveError::UnhashedPackage(name.clone()));
-            }
-
             // Emit a request to fetch the metadata for this package.
             if self.index.packages().register(name.clone()) {
                 request_sink.blocking_send(Request::Package(name.clone()))?;
@@ -1077,6 +1072,14 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
         // Emit a request to fetch the metadata for this version.
         if matches!(&**package, PubGrubPackageInner::Package { .. }) {
             if self.index.distributions().register(candidate.version_id()) {
+                // Verify that the package is allowed under the hash-checking policy.
+                if !self
+                    .hasher
+                    .allows_package(candidate.name(), candidate.version())
+                {
+                    return Err(ResolveError::UnhashedPackage(candidate.name().clone()));
+                }
+
                 let request = Request::from(dist.for_resolution());
                 request_sink.blocking_send(request)?;
             }
@@ -1801,6 +1804,14 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
 
                 // Emit a request to fetch the metadata for this version.
                 if self.index.distributions().register(candidate.version_id()) {
+                    // Verify that the package is allowed under the hash-checking policy.
+                    if !self
+                        .hasher
+                        .allows_package(candidate.name(), candidate.version())
+                    {
+                        return Err(ResolveError::UnhashedPackage(candidate.name().clone()));
+                    }
+
                     let dist = dist.for_resolution().to_owned();
 
                     let response = match dist {

diff --git a/crates/uv-resolver/src/version_map.rs b/crates/uv-resolver/src/version_map.rs
@@ -2,7 +2,6 @@ use std::collections::btree_map::{BTreeMap, Entry};
 use std::sync::OnceLock;
 
 use rkyv::{de::deserializers::SharedDeserializeMap, Deserialize};
-use rustc_hash::FxHashSet;
 use tracing::instrument;
 
 use distribution_filename::{DistFilename, WheelFilename};
@@ -94,11 +93,6 @@ impl VersionMap {
                 },
             }
         }
-        let allowed_yanks = allowed_yanks
-            .allowed_versions(package_name)
-            .cloned()
-            .unwrap_or_default();
-        let required_hashes = hasher.get_package(package_name).digests().to_vec();
         Self {
             inner: VersionMapInner::Lazy(VersionMapLazy {
                 map,
@@ -107,10 +101,10 @@ impl VersionMap {
                 no_build: build_options.no_build_package(package_name),
                 index: index.clone(),
                 tags: tags.cloned(),
+                allowed_yanks: allowed_yanks.clone(),
+                hasher: hasher.clone(),
                 requires_python: requires_python.cloned(),
                 exclude_newer: exclude_newer.copied(),
-                allowed_yanks,
-                required_hashes,
             }),
         }
     }
@@ -288,9 +282,9 @@ struct VersionMapLazy {
     /// Whether files newer than this timestamp should be excluded or not.
     exclude_newer: Option<ExcludeNewer>,
     /// Which yanked versions are allowed
-    allowed_yanks: FxHashSet<Version>,
+    allowed_yanks: AllowedYanks,
     /// The hashes of allowed distributions.
-    required_hashes: Vec<HashDigest>,
+    hasher: HashStrategy,
     /// The `requires-python` constraint for the resolution.
     requires_python: Option<RequiresPython>,
 }
@@ -366,14 +360,14 @@ impl VersionMapLazy {
                 };
 
                 // Prioritize amongst all available files.
-                let version = filename.version().clone();
                 let yanked = file.yanked.clone();
                 let hashes = file.hashes.clone();
                 match filename {
                     DistFilename::WheelFilename(filename) => {
                         let compatibility = self.wheel_compatibility(
                             &filename,
-                            &version,
+                            &filename.name,
+                            &filename.version,
                             &hashes,
                             yanked,
                             excluded,
@@ -388,7 +382,8 @@ impl VersionMapLazy {
                     }
                     DistFilename::SourceDistFilename(filename) => {
                         let compatibility = self.source_dist_compatibility(
-                            &version,
+                            &filename.name,
+                            &filename.version,
                             &hashes,
                             yanked,
                             excluded,
@@ -416,6 +411,7 @@ impl VersionMapLazy {
 
     fn source_dist_compatibility(
         &self,
+        name: &PackageName,
         version: &Version,
         hashes: &[HashDigest],
         yanked: Option<Yanked>,
@@ -436,21 +432,20 @@ impl VersionMapLazy {
 
         // Check if yanked
         if let Some(yanked) = yanked {
-            if yanked.is_yanked() && !self.allowed_yanks.contains(version) {
+            if yanked.is_yanked() && !self.allowed_yanks.contains(name, version) {
                 return SourceDistCompatibility::Incompatible(IncompatibleSource::Yanked(yanked));
             }
         }
 
         // Check if hashes line up. If hashes aren't required, they're considered matching.
-        let hash = if self.required_hashes.is_empty() {
+        let hash_policy = self.hasher.get_package(name, version);
+        let required_hashes = hash_policy.digests();
+        let hash = if required_hashes.is_empty() {
             HashComparison::Matched
         } else {
             if hashes.is_empty() {
                 HashComparison::Missing
-            } else if hashes
-                .iter()
-                .any(|hash| self.required_hashes.contains(hash))
-            {
+            } else if hashes.iter().any(|hash| required_hashes.contains(hash)) {
                 HashComparison::Matched
             } else {
                 HashComparison::Mismatched
@@ -463,6 +458,7 @@ impl VersionMapLazy {
     fn wheel_compatibility(
         &self,
         filename: &WheelFilename,
+        name: &PackageName,
         version: &Version,
         hashes: &[HashDigest],
         yanked: Option<Yanked>,
@@ -481,7 +477,7 @@ impl VersionMapLazy {
 
         // Check if yanked
         if let Some(yanked) = yanked {
-            if yanked.is_yanked() && !self.allowed_yanks.contains(version) {
+            if yanked.is_yanked() && !self.allowed_yanks.contains(name, version) {
                 return WheelCompatibility::Incompatible(IncompatibleWheel::Yanked(yanked));
             }
         }
@@ -498,15 +494,14 @@ impl VersionMapLazy {
         };
 
         // Check if hashes line up. If hashes aren't required, they're considered matching.
-        let hash = if self.required_hashes.is_empty() {
+        let hash_policy = self.hasher.get_package(name, version);
+        let required_hashes = hash_policy.digests();
+        let hash = if required_hashes.is_empty() {
             HashComparison::Matched
         } else {
             if hashes.is_empty() {
                 HashComparison::Missing
-            } else if hashes
-                .iter()
-                .any(|hash| self.required_hashes.contains(hash))
-            {
+            } else if hashes.iter().any(|hash| required_hashes.contains(hash)) {
                 HashComparison::Matched
             } else {
                 HashComparison::Mismatched

diff --git a/crates/uv-resolver/src/yanks.rs b/crates/uv-resolver/src/yanks.rs
@@ -1,5 +1,6 @@
 use pypi_types::RequirementSource;
 use rustc_hash::{FxHashMap, FxHashSet};
+use std::sync::Arc;
 
 use pep440_rs::Version;
 use pep508_rs::MarkerEnvironment;
@@ -10,7 +11,7 @@ use crate::{DependencyMode, Manifest};
 /// A set of package versions that are permitted, even if they're marked as yanked by the
 /// relevant index.
 #[derive(Debug, Default, Clone)]
-pub struct AllowedYanks(FxHashMap<PackageName, FxHashSet<Version>>);
+pub struct AllowedYanks(Arc<FxHashMap<PackageName, FxHashSet<Version>>>);
 
 impl AllowedYanks {
     pub fn from_manifest(
@@ -47,7 +48,14 @@ impl AllowedYanks {
                 .insert(version.clone());
         }
 
-        Self(allowed_yanks)
+        Self(Arc::new(allowed_yanks))
+    }
+
+    /// Returns `true` if the package-version is allowed, even if it's marked as yanked.
+    pub fn contains(&self, package_name: &PackageName, version: &Version) -> bool {
+        self.0
+            .get(package_name)
+            .map_or(false, |versions| versions.contains(version))
     }
 
     /// Returns versions for the given package which are allowed even if marked as yanked by the