From bdeab5519358735b2a4f403c25eec80cbe91daa9 Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Mon, 8 Apr 2024 15:34:08 -0500
Subject: [PATCH] Add extract support for zstd (#2861)

We need this to extract toolchain downloads
---
 Cargo.lock                      | 30 ++++++++++++++++++++++++++++
 crates/uv-extract/Cargo.toml    |  2 +-
 crates/uv-extract/src/seek.rs   | 17 +++++++++++++++-
 crates/uv-extract/src/stream.rs | 35 +++++++++++++++++++++++++++++++--
 4 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0cd3b6ce6..c8bc8a192 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -211,6 +211,8 @@ dependencies = [
  "memchr",
  "pin-project-lite",
  "tokio",
+ "zstd",
+ "zstd-safe",
 ]
 
 [[package]]
@@ -5382,3 +5384,31 @@ dependencies = [
  "crossbeam-utils",
  "flate2",
 ]
+
+[[package]]
+name = "zstd"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2d789b1514203a1120ad2429eae43a7bd32b90976a7bb8a05f7ec02fa88cc23a"
+dependencies = [
+ "zstd-safe",
+]
+
+[[package]]
+name = "zstd-safe"
+version = "7.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1cd99b45c6bc03a018c8b8a86025678c87e55526064e38f9df301989dce7ec0a"
+dependencies = [
+ "zstd-sys",
+]
+
+[[package]]
+name = "zstd-sys"
+version = "2.0.10+zstd.1.5.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c253a4914af5bafc8fa8c86ee400827e83cf6ec01195ec1f1ed8441bf00d65aa"
+dependencies = [
+ "cc",
+ "pkg-config",
+]
diff --git a/crates/uv-extract/Cargo.toml b/crates/uv-extract/Cargo.toml
index ca0704c4c..f40dee266 100644
--- a/crates/uv-extract/Cargo.toml
+++ b/crates/uv-extract/Cargo.toml
@@ -13,7 +13,7 @@ license = { workspace = true }
 workspace = true
 
 [dependencies]
-async-compression = { workspace = true, features = ["gzip"] }
+async-compression = { workspace = true, features = ["gzip", "zstd"] }
 async_zip = { workspace = true, features = ["tokio"] }
 fs-err = { workspace = true, features = ["tokio"] }
 futures = { workspace = true }
diff --git a/crates/uv-extract/src/seek.rs b/crates/uv-extract/src/seek.rs
index 636bb6318..57499fc86 100644
--- a/crates/uv-extract/src/seek.rs
+++ b/crates/uv-extract/src/seek.rs
@@ -107,7 +107,22 @@ pub async fn archive<R: tokio::io::AsyncRead + tokio::io::AsyncSeek + Unpin>(
                 .is_some_and(|ext| ext.eq_ignore_ascii_case("tar"))
         })
     {
-        crate::stream::untar(reader, target).await?;
+        crate::stream::untar_gz(reader, target).await?;
+        return Ok(());
+    }
+
+    // `.tar.zst`
+    if source
+        .as_ref()
+        .extension()
+        .is_some_and(|ext| ext.eq_ignore_ascii_case("zst"))
+        && source.as_ref().file_stem().is_some_and(|stem| {
+            Path::new(stem)
+                .extension()
+                .is_some_and(|ext| ext.eq_ignore_ascii_case("tar"))
+        })
+    {
+        crate::stream::untar_zst(reader, target).await?;
         return Ok(());
     }
 
diff --git a/crates/uv-extract/src/stream.rs b/crates/uv-extract/src/stream.rs
index 7fb4707fc..cff2825dc 100644
--- a/crates/uv-extract/src/stream.rs
+++ b/crates/uv-extract/src/stream.rs
@@ -151,7 +151,7 @@ async fn untar_in<R: tokio::io::AsyncRead + Unpin, P: AsRef<Path>>(
 /// Unzip a `.tar.gz` archive into the target directory, without requiring `Seek`.
 ///
 /// This is useful for unpacking files as they're being downloaded.
-pub async fn untar<R: tokio::io::AsyncRead + Unpin>(
+pub async fn untar_gz<R: tokio::io::AsyncRead + Unpin>(
     reader: R,
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
@@ -163,6 +163,22 @@ pub async fn untar<R: tokio::io::AsyncRead + Unpin>(
     Ok(untar_in(&mut archive, target.as_ref()).await?)
 }
 
+/// Unzip a `.tar.zst` archive into the target directory, without requiring `Seek`.
+///
+/// This is useful for unpacking files as they're being downloaded.
+pub async fn untar_zst<R: tokio::io::AsyncRead + Unpin>(
+    reader: R,
+    target: impl AsRef<Path>,
+) -> Result<(), Error> {
+    let reader = tokio::io::BufReader::new(reader);
+    let decompressed_bytes = async_compression::tokio::bufread::ZstdDecoder::new(reader);
+
+    let mut archive = tokio_tar::ArchiveBuilder::new(decompressed_bytes)
+        .set_preserve_mtime(false)
+        .build();
+    Ok(untar_in(&mut archive, target.as_ref()).await?)
+}
+
 /// Unzip a `.zip` or `.tar.gz` archive into the target directory, without requiring `Seek`.
 pub async fn archive<R: tokio::io::AsyncRead + Unpin>(
     reader: R,
@@ -190,7 +206,22 @@ pub async fn archive<R: tokio::io::AsyncRead + Unpin>(
                 .is_some_and(|ext| ext.eq_ignore_ascii_case("tar"))
         })
     {
-        untar(reader, target).await?;
+        untar_gz(reader, target).await?;
+        return Ok(());
+    }
+
+    // `.tar.zst`
+    if source
+        .as_ref()
+        .extension()
+        .is_some_and(|ext| ext.eq_ignore_ascii_case("zst"))
+        && source.as_ref().file_stem().is_some_and(|stem| {
+            Path::new(stem)
+                .extension()
+                .is_some_and(|ext| ext.eq_ignore_ascii_case("tar"))
+        })
+    {
+        untar_zst(reader, target).await?;
         return Ok(());
     }