Add ability to unzip artifacts zipped with zstd

Author: NicoHinderling

requirements.txt

@@ -20,3 +20,4 @@ sentry-kafka-schemas==2.1.2
 sentry-sdk>=2.36.0
 sortedcontainers>=2.4.0
 typing-extensions>=4.15.0
+zipfile-zstd==0.0.4

src/launchpad/artifacts/providers/zip_provider.py

@@ -1,7 +1,7 @@
 import zipfile
 
 from pathlib import Path
-from typing import List
+from typing import List, Set
 
 from launchpad.utils.file_utils import cleanup_directory, create_temp_directory
 from launchpad.utils.logging import get_logger
@@ -11,6 +11,9 @@
 DEFAULT_MAX_FILE_COUNT = 100000
 DEFAULT_MAX_UNCOMPRESSED_SIZE = 10 * 1024 * 1024 * 1024
 
+# Compression method constants
+COMPRESSION_ZSTD = 93  # Zstandard compression method
+
 
 class UnreasonableZipError(ValueError):
     """Raised when a zip file exceeds reasonable limits."""
@@ -92,13 +95,43 @@ def extract_to_temp_directory(self) -> Path:
         self._temp_dirs.append(temp_dir)
 
         self._safe_extract(str(self.path), str(temp_dir))
-        logger.debug(f"Extracted zip contents to {temp_dir} using system unzip")
+        logger.debug(f"Extracted zip contents to {temp_dir}")
 
         return temp_dir
 
+    def _detect_compression_methods(self, zip_path: str) -> Set[int]:
+        """Detect compression methods used in the zip file.
+
+        Args:
+            zip_path: Path to the zip file
+
+        Returns:
+            Set of compression method integers used in the zip file
+        """
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            return {info.compress_type for info in zf.infolist()}
+
     def _safe_extract(self, zip_path: str, extract_path: str):
-        """Extract the zip contents to a temporary directory, ensuring that the paths are safe from path traversal attacks."""
+        """Extract the zip contents to a temporary directory, ensuring that the paths are safe from path traversal attacks.
+
+        Supports both standard compression methods and Zstandard compression.
+        """
         base_dir = Path(extract_path)
+
+        # Detect if zstandard compression is used
+        compression_methods = self._detect_compression_methods(zip_path)
+        uses_zstd = COMPRESSION_ZSTD in compression_methods
+
+        if uses_zstd:
+            logger.debug("Detected Zstandard compression in zip file")
+            try:
+                import zipfile_zstd  # noqa: F401
+            except ImportError:
+                raise RuntimeError(
+                    "Zstandard-compressed zip file detected, but zipfile-zstd package is not installed. "
+                    "Install it with: pip install zipfile-zstd"
+                )
+
         with zipfile.ZipFile(zip_path, "r") as zip_ref:
             check_reasonable_zip(zip_ref)
             for member in zip_ref.namelist():

tests/unit/artifacts/providers/test_zip_provider.py

@@ -6,6 +6,7 @@
 import pytest
 
 from launchpad.artifacts.providers.zip_provider import (
+    COMPRESSION_ZSTD,
     UnreasonableZipError,
     UnsafePathError,
     ZipProvider,
@@ -88,6 +89,37 @@ def test_invalid_zip_file(self) -> None:
             with pytest.raises(zipfile.BadZipFile):
                 provider.extract_to_temp_directory()
 
+    def test_detect_compression_methods(self, hackernews_xcarchive: Path) -> None:
+        provider = ZipProvider(hackernews_xcarchive)
+        methods = provider._detect_compression_methods(str(hackernews_xcarchive))
+
+        # Verify detection works and standard zips use deflate compression, not zstd
+        assert zipfile.ZIP_DEFLATED in methods
+        assert COMPRESSION_ZSTD not in methods
+
+    def test_extract_zstd_zip(self) -> None:
+        """Test that zstd-compressed zips can be extracted when zipfile-zstd is available."""
+        try:
+            import zipfile_zstd  # noqa: F401
+        except ImportError:
+            pytest.skip("zipfile-zstd not installed")
+
+        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as temp_file:
+            temp_path = Path(temp_file.name)
+
+            with zipfile.ZipFile(temp_path, "w") as zf:
+                zf.writestr("test.txt", "content", compress_type=COMPRESSION_ZSTD)
+
+            try:
+                provider = ZipProvider(temp_path)
+                temp_dir = provider.extract_to_temp_directory()
+
+                assert temp_dir.exists()
+                assert (temp_dir / "test.txt").exists()
+                assert (temp_dir / "test.txt").read_text() == "content"
+            finally:
+                temp_path.unlink(missing_ok=True)
+
 
 class TestIsSafePath:
     def test_valid_paths(self) -> None: