Enforce extension validity at parse time (#5888)

This PR adds a `DistExtension` field to some of our distribution types,
which requires that we validate that the file type is known and
supported when parsing (rather than when attempting to unzip). It
removes a bunch of extension parsing from the code too, in favor of
doing it once upfront.

Closes #5858.

crates/distribution-filename/src/extension.rs

@@ -0,0 +1,99 @@
+use std::fmt::{Display, Formatter};
+use std::path::Path;
+
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum DistExtension {
+    Wheel,
+    Source(SourceDistExtension),
+}
+
+#[derive(
+    Clone,
+    Copy,
+    Debug,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    Hash,
+    Serialize,
+    Deserialize,
+    rkyv::Archive,
+    rkyv::Deserialize,
+    rkyv::Serialize,
+)]
+#[archive(check_bytes)]
+#[archive_attr(derive(Debug))]
+pub enum SourceDistExtension {
+    Zip,
+    TarGz,
+    TarBz2,
+    TarXz,
+    TarZst,
+}
+
+impl DistExtension {
+    /// Extract the [`DistExtension`] from a path.
+    pub fn from_path(path: impl AsRef<Path>) -> Result<Self, ExtensionError> {
+        let Some(extension) = path.as_ref().extension().and_then(|ext| ext.to_str()) else {
+            return Err(ExtensionError::Dist);
+        };
+
+        match extension {
+            "whl" => Ok(Self::Wheel),
+            _ => SourceDistExtension::from_path(path)
+                .map(Self::Source)
+                .map_err(|_| ExtensionError::Dist),
+        }
+    }
+}
+
+impl SourceDistExtension {
+    /// Extract the [`SourceDistExtension`] from a path.
+    pub fn from_path(path: impl AsRef<Path>) -> Result<Self, ExtensionError> {
+        /// Returns true if the path is a tar file (e.g., `.tar.gz`).
+        fn is_tar(path: &Path) -> bool {
+            path.file_stem().is_some_and(|stem| {
+                Path::new(stem)
+                    .extension()
+                    .is_some_and(|ext| ext.eq_ignore_ascii_case("tar"))
+            })
+        }
+
+        let Some(extension) = path.as_ref().extension().and_then(|ext| ext.to_str()) else {
+            return Err(ExtensionError::SourceDist);
+        };
+
+        match extension {
+            "zip" => Ok(Self::Zip),
+            "gz" if is_tar(path.as_ref()) => Ok(Self::TarGz),
+            "bz2" if is_tar(path.as_ref()) => Ok(Self::TarBz2),
+            "xz" if is_tar(path.as_ref()) => Ok(Self::TarXz),
+            "zst" if is_tar(path.as_ref()) => Ok(Self::TarZst),
+            _ => Err(ExtensionError::SourceDist),
+        }
+    }
+}
+
+impl Display for SourceDistExtension {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Zip => f.write_str("zip"),
+            Self::TarGz => f.write_str("tar.gz"),
+            Self::TarBz2 => f.write_str("tar.bz2"),
+            Self::TarXz => f.write_str("tar.xz"),
+            Self::TarZst => f.write_str("tar.zst"),
+        }
+    }
+}
+
+#[derive(Error, Debug)]
+pub enum ExtensionError {
+    #[error("`.whl`, `.zip`, `.tar.gz`, `.tar.bz2`, `.tar.xz`, or `.tar.zst`")]
+    Dist,
+    #[error("`.zip`, `.tar.gz`, `.tar.bz2`, `.tar.xz`, or `.tar.zst`")]
+    SourceDist,
+}

crates/distribution-filename/src/lib.rs

@@ -5,11 +5,13 @@ use uv_normalize::PackageName;
 
 pub use build_tag::{BuildTag, BuildTagError};
 pub use egg::{EggInfoFilename, EggInfoFilenameError};
-pub use source_dist::{SourceDistExtension, SourceDistFilename, SourceDistFilenameError};
+pub use extension::{DistExtension, ExtensionError, SourceDistExtension};
+pub use source_dist::SourceDistFilename;
 pub use wheel::{WheelFilename, WheelFilenameError};
 
 mod build_tag;
 mod egg;
+mod extension;
 mod source_dist;
 mod wheel;
 
@@ -22,13 +24,20 @@ pub enum DistFilename {
 impl DistFilename {
     /// Parse a filename as wheel or source dist name.
     pub fn try_from_filename(filename: &str, package_name: &PackageName) -> Option<Self> {
-        if let Ok(filename) = WheelFilename::from_str(filename) {
-            Some(Self::WheelFilename(filename))
-        } else if let Ok(filename) = SourceDistFilename::parse(filename, package_name) {
-            Some(Self::SourceDistFilename(filename))
-        } else {
-            None
+        match DistExtension::from_path(filename) {
+            Ok(DistExtension::Wheel) => {
+                if let Ok(filename) = WheelFilename::from_str(filename) {
+                    return Some(Self::WheelFilename(filename));
+                }
+            }
+            Ok(DistExtension::Source(extension)) => {
+                if let Ok(filename) = SourceDistFilename::parse(filename, extension, package_name) {
+                    return Some(Self::SourceDistFilename(filename));
+                }
+            }
+            Err(_) => {}
         }
+        None
     }
 
     /// Like [`DistFilename::try_from_normalized_filename`], but without knowing the package name.

crates/distribution-filename/src/source_dist.rs

@@ -1,69 +1,12 @@
 use std::fmt::{Display, Formatter};
 use std::str::FromStr;
 
+use crate::SourceDistExtension;
+use pep440_rs::{Version, VersionParseError};
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
-
-use pep440_rs::{Version, VersionParseError};
 use uv_normalize::{InvalidNameError, PackageName};
 
-#[derive(
-    Clone,
-    Debug,
-    PartialEq,
-    Eq,
-    Serialize,
-    Deserialize,
-    rkyv::Archive,
-    rkyv::Deserialize,
-    rkyv::Serialize,
-)]
-#[archive(check_bytes)]
-#[archive_attr(derive(Debug))]
-pub enum SourceDistExtension {
-    Zip,
-    TarGz,
-    TarBz2,
-}
-
-impl FromStr for SourceDistExtension {
-    type Err = String;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        Ok(match s {
-            "zip" => Self::Zip,
-            "tar.gz" => Self::TarGz,
-            "tar.bz2" => Self::TarBz2,
-            other => return Err(other.to_string()),
-        })
-    }
-}
-
-impl Display for SourceDistExtension {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Zip => f.write_str("zip"),
-            Self::TarGz => f.write_str("tar.gz"),
-            Self::TarBz2 => f.write_str("tar.bz2"),
-        }
-    }
-}
-
-impl SourceDistExtension {
-    pub fn from_filename(filename: &str) -> Option<(&str, Self)> {
-        if let Some(stem) = filename.strip_suffix(".zip") {
-            return Some((stem, Self::Zip));
-        }
-        if let Some(stem) = filename.strip_suffix(".tar.gz") {
-            return Some((stem, Self::TarGz));
-        }
-        if let Some(stem) = filename.strip_suffix(".tar.bz2") {
-            return Some((stem, Self::TarBz2));
-        }
-        None
-    }
-}
-
 /// Note that this is a normalized and not an exact representation, keep the original string if you
 /// need the latter.
 #[derive(
@@ -90,14 +33,18 @@ impl SourceDistFilename {
     /// these (consider e.g. `a-1-1.zip`)
     pub fn parse(
         filename: &str,
+        extension: SourceDistExtension,
         package_name: &PackageName,
     ) -> Result<Self, SourceDistFilenameError> {
-        let Some((stem, extension)) = SourceDistExtension::from_filename(filename) else {
+        // Drop the extension (e.g., given `tar.gz`, drop `.tar.gz`).
+        if filename.len() <= extension.to_string().len() + 1 {
             return Err(SourceDistFilenameError {
                 filename: filename.to_string(),
                 kind: SourceDistFilenameErrorKind::Extension,
             });
-        };
+        }
+
+        let stem = &filename[..(filename.len() - (extension.to_string().len() + 1))];
 
         if stem.len() <= package_name.as_ref().len() + "-".len() {
             return Err(SourceDistFilenameError {
@@ -138,13 +85,23 @@ impl SourceDistFilename {
     /// Source dist filenames can be ambiguous, e.g. `a-1-1.tar.gz`. Without knowing the package name, we assume that
     /// source dist filename version doesn't contain minus (the version is normalized).
     pub fn parsed_normalized_filename(filename: &str) -> Result<Self, SourceDistFilenameError> {
-        let Some((stem, extension)) = SourceDistExtension::from_filename(filename) else {
+        let Ok(extension) = SourceDistExtension::from_path(filename) else {
             return Err(SourceDistFilenameError {
                 filename: filename.to_string(),
                 kind: SourceDistFilenameErrorKind::Extension,
             });
         };
 
+        // Drop the extension (e.g., given `tar.gz`, drop `.tar.gz`).
+        if filename.len() <= extension.to_string().len() + 1 {
+            return Err(SourceDistFilenameError {
+                filename: filename.to_string(),
+                kind: SourceDistFilenameErrorKind::Extension,
+            });
+        }
+
+        let stem = &filename[..(filename.len() - (extension.to_string().len() + 1))];
+
         let Some((package_name, version)) = stem.rsplit_once('-') else {
             return Err(SourceDistFilenameError {
                 filename: filename.to_string(),
@@ -197,7 +154,7 @@ impl Display for SourceDistFilenameError {
 enum SourceDistFilenameErrorKind {
     #[error("Name doesn't start with package name {0}")]
     Filename(PackageName),
-    #[error("Source distributions filenames must end with .zip, .tar.gz, or .tar.bz2")]
+    #[error("File extension is invalid")]
     Extension,
     #[error("Version section is invalid")]
     Version(#[from] VersionParseError),
@@ -213,7 +170,7 @@ mod tests {
 
     use uv_normalize::PackageName;
 
-    use crate::SourceDistFilename;
+    use crate::{SourceDistExtension, SourceDistFilename};
 
     /// Only test already normalized names since the parsing is lossy
     #[test]
@@ -223,30 +180,40 @@ mod tests {
             "foo-lib-1.2.3a3.zip",
             "foo-lib-1.2.3.tar.gz",
             "foo-lib-1.2.3.tar.bz2",
+            "foo-lib-1.2.3.tar.zst",
         ] {
+            let ext = SourceDistExtension::from_path(normalized).unwrap();
             assert_eq!(
-                SourceDistFilename::parse(normalized, &PackageName::from_str("foo_lib").unwrap())
-                    .unwrap()
-                    .to_string(),
+                SourceDistFilename::parse(
+                    normalized,
+                    ext,
+                    &PackageName::from_str("foo_lib").unwrap()
+                )
+                .unwrap()
+                .to_string(),
                 normalized
             );
         }
     }
 
     #[test]
     fn errors() {
-        for invalid in ["b-1.2.3.zip", "a-1.2.3-gamma.3.zip", "a-1.2.3.tar.zstd"] {
+        for invalid in ["b-1.2.3.zip", "a-1.2.3-gamma.3.zip"] {
+            let ext = SourceDistExtension::from_path(invalid).unwrap();
             assert!(
-                SourceDistFilename::parse(invalid, &PackageName::from_str("a").unwrap()).is_err()
+                SourceDistFilename::parse(invalid, ext, &PackageName::from_str("a").unwrap())
+                    .is_err()
             );
         }
     }
 
     #[test]
-    fn name_to_long() {
-        assert!(
-            SourceDistFilename::parse("foo.zip", &PackageName::from_str("foo-lib").unwrap())
-                .is_err()
-        );
+    fn name_too_long() {
+        assert!(SourceDistFilename::parse(
+            "foo.zip",
+            SourceDistExtension::Zip,
+            &PackageName::from_str("foo-lib").unwrap()
+        )
+        .is_err());
     }
 }

crates/distribution-types/src/buildable.rs

@@ -1,6 +1,7 @@
 use std::borrow::Cow;
 use std::path::Path;
 
+use distribution_filename::SourceDistExtension;
 use pep440_rs::Version;
 use pep508_rs::VerbatimUrl;
 use url::Url;
@@ -109,6 +110,8 @@ impl std::fmt::Display for SourceUrl<'_> {
 #[derive(Debug, Clone)]
 pub struct DirectSourceUrl<'a> {
     pub url: &'a Url,
+    pub subdirectory: Option<&'a Path>,
+    pub ext: SourceDistExtension,
 }
 
 impl std::fmt::Display for DirectSourceUrl<'_> {
@@ -146,6 +149,7 @@ impl<'a> From<&'a GitSourceDist> for GitSourceUrl<'a> {
 pub struct PathSourceUrl<'a> {
     pub url: &'a Url,
     pub path: Cow<'a, Path>,
+    pub ext: SourceDistExtension,
 }
 
 impl std::fmt::Display for PathSourceUrl<'_> {
@@ -159,6 +163,7 @@ impl<'a> From<&'a PathSourceDist> for PathSourceUrl<'a> {
         Self {
             url: &dist.url,
             path: Cow::Borrowed(&dist.install_path),
+            ext: dist.ext,
         }
     }
 }