From f99e3560e84a8af855bcd7aa3d9cb6ae036ec2c9 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Tue, 5 Dec 2023 03:53:45 -0500
Subject: [PATCH] Avoid returning zipped wheels from registry and URL indexes
 (#558)

## Summary

This is hard to reproduce, but if you run a long installation process
that errors part-way through, you can end up with zipped wheels in the
`Wheels` cache, which is intended to contain only unzipped wheels. This
PR avoids returning those entries from the registry, which will then
lead to errors downstream when we treat them as directories.
---
 crates/puffin-distribution/src/distribution_database.rs | 4 ++--
 crates/puffin-installer/src/plan.rs                     | 9 ++++-----
 crates/puffin-installer/src/registry_index.rs           | 5 +++++
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/crates/puffin-distribution/src/distribution_database.rs b/crates/puffin-distribution/src/distribution_database.rs
index 2dafa0892..0ba21df49 100644
--- a/crates/puffin-distribution/src/distribution_database.rs
+++ b/crates/puffin-distribution/src/distribution_database.rs
@@ -103,7 +103,7 @@ impl<'a, Context: BuildContext + Send + Sync> DistributionDatabase<'a, Context>
         }
     }
 
-    /// In parallel, either fetch the wheel or fetch and built source distributions.
+    /// In parallel, either fetch each wheel or build each source distribution.
     pub async fn get_wheels(
         &self,
         dists: Vec<Dist>,
@@ -191,7 +191,7 @@ impl<'a, Context: BuildContext + Send + Sync> DistributionDatabase<'a, Context>
                         small_size.map_or("unknown size".to_string(), |size| size.to_string());
                     debug!("Fetching disk-based wheel from registry: {dist} ({size})");
 
-                    // Download the wheel to a temporary file.
+                    // Download the wheel into the cache.
                     fs::create_dir_all(&cache_entry.dir).await?;
                     let mut writer = fs::File::create(cache_entry.path()).await?;
                     tokio::io::copy(&mut reader.compat(), &mut writer).await?;
diff --git a/crates/puffin-installer/src/plan.rs b/crates/puffin-installer/src/plan.rs
index f93ffb2a9..a8b9a41e3 100644
--- a/crates/puffin-installer/src/plan.rs
+++ b/crates/puffin-installer/src/plan.rs
@@ -131,17 +131,16 @@ impl InstallPlan {
                             WheelCache::Url(url).wheel_dir(),
                             disk_filename.to_string(),
                         );
-                        if cache_entry.path().exists() {
-                            // TODO(charlie): This takes advantage of the fact that for URL dependencies, the package ID
-                            // and distribution ID are identical. We should either change the cache layout to use
-                            // distribution IDs, or implement package ID for URL.
+
+                        // Ignore zipped wheels, which represent intermediary cached artifacts.
+                        if cache_entry.path().is_dir() {
                             let cached_dist = CachedDirectUrlDist::from_url(
                                 filename,
                                 url.clone(),
                                 cache_entry.path(),
                             );
 
-                            debug!("Url wheel requirement already cached: {cached_dist}");
+                            debug!("URL wheel requirement already cached: {cached_dist}");
                             local.push(CachedDist::Url(cached_dist.clone()));
                             continue;
                         }
diff --git a/crates/puffin-installer/src/registry_index.rs b/crates/puffin-installer/src/registry_index.rs
index fbe7e41c9..17af2e546 100644
--- a/crates/puffin-installer/src/registry_index.rs
+++ b/crates/puffin-installer/src/registry_index.rs
@@ -41,6 +41,11 @@ impl RegistryIndex {
                     }
                 };
 
+                // Ignore zipped wheels, which represent intermediary cached artifacts.
+                if !path.is_dir() {
+                    continue;
+                }
+
                 match CachedRegistryDist::try_from_path(&path) {
                     Ok(None) => {}
                     Ok(Some(dist_info)) => {