From f70b7867b0cf5abba1389c3963e634732f7373a5 Mon Sep 17 00:00:00 2001
From: lipengyu <lipengyu@kylinos.cn>
Date: Fri, 31 Jul 2026 17:40:50 +0800
Subject: [PATCH] gh-155152: Fix quadratic processing of concatenated zstd
 frames in zipimport

---
 Lib/test/test_zipimport.py                    | 56 +++++++++++++++++++
 Lib/zipimport.py                              | 30 +++++++---
 ...-08-04-16-38-59.gh-issue-155152.qN1EEN.rst |  2 +
 3 files changed, 79 insertions(+), 9 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst

diff --git a/Lib/test/test_zipimport.py b/Lib/test/test_zipimport.py
index 76cd85709a63af8..cdc68d401b5c807 100644
--- a/Lib/test/test_zipimport.py
+++ b/Lib/test/test_zipimport.py
@@ -1026,6 +1026,62 @@ class DeflateCompressedZipImportTestCase(UncompressedZipImportTestCase):
 class ZStdCompressedZipImportTestCase(UncompressedZipImportTestCase):
     compression = ZIP_ZSTANDARD
 
+    def test_concatenated_frames(self):
+        from compression import zstd
+        import zipfile
+
+        class ConcatenatedFrameCompressor:
+            def compress(self, data):
+                return b"".join(zstd.compress(bytes([byte])) for byte in data)
+
+            def flush(self):
+                return b""
+
+        with unittest.mock.patch.object(
+            zipfile,
+            "_get_compressor",
+            return_value=ConcatenatedFrameCompressor(),
+        ):
+            self.doTest(".py", {TESTMOD + ".py": test_src}, TESTMOD)
+
+    def test_concatenated_frames_linear_input(self):
+        from compression import zstd
+
+        frame_count = 100
+        frame = zstd.compress(b"x")
+        compressed = frame * frame_count
+        input_sizes = []
+        decompressor_class = zipimport._get_zstd_decompressor_class()
+
+        class RecordingDecompressor:
+            def __init__(self):
+                self._decompressor = decompressor_class()
+
+            def __getattr__(self, name):
+                return getattr(self._decompressor, name)
+
+            def decompress(self, data, max_length=-1):
+                input_sizes.append(memoryview(data).nbytes)
+                return self._decompressor.decompress(data, max_length)
+
+        with unittest.mock.patch.object(
+            zipimport, "_zstd_decompressor_class", RecordingDecompressor
+        ):
+            self.assertEqual(
+                zipimport._zstd_decompress(compressed), b"x" * frame_count
+            )
+
+        self.assertEqual(sum(input_sizes), len(compressed))
+
+    def test_truncated_frame(self):
+        from compression import zstd
+
+        compressed = zstd.compress(b"x")
+        for data in (b"", compressed[:-1]):
+            with self.subTest(data=data):
+                with self.assertRaises(zipimport.ZipImportError):
+                    zipimport._zstd_decompress(data)
+
 
 class BadFileZipImportTestCase(unittest.TestCase):
     def assertZipFailure(self, filename):
diff --git a/Lib/zipimport.py b/Lib/zipimport.py
index 19279d1c2bea368..626db6cd6e470f0 100644
--- a/Lib/zipimport.py
+++ b/Lib/zipimport.py
@@ -584,12 +584,13 @@ def _get_zlib_decompress_func():
 
 _importing_zstd = False
 _zstd_decompressor_class = None
+_zstd_get_frame_size = None
 
 # Return the _zstd.ZstdDecompressor function object, or NULL if _zstd couldn't
 # be imported. The result is cached when found.
 def _get_zstd_decompressor_class():
-    global _zstd_decompressor_class
-    if _zstd_decompressor_class:
+    global _zstd_decompressor_class, _zstd_get_frame_size
+    if _zstd_decompressor_class and _zstd_get_frame_size:
         return _zstd_decompressor_class
 
     global _importing_zstd
@@ -601,7 +602,10 @@ def _get_zstd_decompressor_class():
 
     _importing_zstd = True
     try:
-        from _zstd import ZstdDecompressor as _zstd_decompressor_class
+        from _zstd import (
+            ZstdDecompressor as _zstd_decompressor_class,
+            get_frame_size as _zstd_get_frame_size,
+        )
     except Exception:
         _bootstrap._verbose_message("zipimport: zstd UNAVAILABLE")
         raise ZipImportError("can't decompress data; zstd not available")
@@ -615,16 +619,24 @@ def _get_zstd_decompressor_class():
 def _zstd_decompress(data):
     # A simple version of compression.zstd.decompress() as we cannot import
     # that here as the stdlib itself could be being zipimported.
+    decomp_class = _get_zstd_decompressor_class()
+    data = memoryview(data)
+    if not data:
+        raise ZipImportError("zipimport: zstd compressed data ended before "
+                             "the end-of-stream marker")
     results = []
-    while True:
-        decomp = _get_zstd_decompressor_class()()
-        results.append(decomp.decompress(data))
+    while data:
+        try:
+            frame_size = _zstd_get_frame_size(data)
+        except Exception:
+            raise ZipImportError("zipimport: zstd compressed data ended before "
+                                 "the end-of-stream marker") from None
+        decomp = decomp_class()
+        results.append(decomp.decompress(data[:frame_size]))
         if not decomp.eof:
             raise ZipImportError("zipimport: zstd compressed data ended before "
                                  "the end-of-stream marker")
-        data = decomp.unused_data
-        if not data:
-            break
+        data = data[frame_size:]
     return b"".join(results)
 
 
diff --git a/Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst b/Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst
new file mode 100644
index 000000000000000..5eefb8bc93f4a72
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst
@@ -0,0 +1,2 @@
+Prevent unexpectedly slow imports and excessive CPU usage for some ZIP
+archives that use Zstandard compression.