gh-160310: Fix quadratic one-shot decompress() with many members

gzip.decompress(), bz2.decompress(), lzma.decompress() and
compression.zstd.decompress() passed all of the remaining input to a new
decompressor for every member and copied the rest back via unused_data.
Feed each member in chunks of increasing size instead.

Author: gpshead

Lib/bz2.py

@@ -337,18 +337,42 @@ def decompress(data):
     For incremental decompression, use a BZ2Decompressor object instead.
     """
     results = []
-    while data:
-        decomp = BZ2Decompressor()
-        try:
-            res = decomp.decompress(data)
-        except OSError:
-            if results:
-                break  # Leftover data is not a valid bzip2 stream; ignore it.
-            else:
-                raise  # Error on the first iteration; bail out.
-        results.append(res)
-        if not decomp.eof:
-            raise ValueError("Compressed data ended before the "
-                             "end-of-stream marker was reached")
-        data = decomp.unused_data
+    if not data:
+        return b""
+    # The views are released on exit, so that the caller can resize a
+    # bytearray argument again, also while handling an exception.
+    with memoryview(data) as view, view.cast("B") as data:
+        pos = 0
+        # There is usually a single stream: the first chunk is all of the
+        # data, to decompress it in one go.
+        chunk_size = len(data)
+        while True:
+            decomp = BZ2Decompressor()
+            res = []
+            try:
+                # Feed the stream in chunks of increasing size, so that what
+                # is read past its end (and copied into unused_data) is
+                # bounded by the size of the stream, not by the size of all
+                # the streams that follow it.
+                while not decomp.eof and pos < len(data):
+                    with data[pos:pos + chunk_size] as chunk:
+                        res.append(decomp.decompress(chunk))
+                        pos += len(chunk)
+                    chunk_size *= 2
+            except OSError:
+                if results:
+                    # Leftover data is not a valid bzip2 stream; ignore it.
+                    break
+                else:
+                    raise  # Error on the first iteration; bail out.
+            results.extend(res)
+            if not decomp.eof:
+                raise ValueError("Compressed data ended before the "
+                                 "end-of-stream marker was reached")
+            pos -= len(decomp.unused_data)
+            if pos == len(data):
+                break
+            # Later streams start small: a larger first chunk reads further
+            # past the end of a short stream, a smaller one needs more calls.
+            chunk_size = 1024
     return b"".join(results)

Lib/compression/zstd/__init__.py

@@ -161,15 +161,33 @@ def decompress(data, zstd_dict=None, options=None):
     For incremental decompression, use a ZstdDecompressor instead.
     """
     results = []
-    while True:
-        decomp = ZstdDecompressor(options=options, zstd_dict=zstd_dict)
-        results.append(decomp.decompress(data))
-        if not decomp.eof:
-            raise ZstdError('Compressed data ended before the '
-                            'end-of-stream marker was reached')
-        data = decomp.unused_data
-        if not data:
-            break
+    # The views are released on exit, so that the caller can resize a
+    # bytearray argument again, also while handling an exception.
+    with memoryview(data) as view, view.cast('B') as data:
+        pos = 0
+        # There is usually a single frame: the first chunk is all of the
+        # data, to decompress it in one go.
+        chunk_size = len(data)
+        while True:
+            decomp = ZstdDecompressor(options=options, zstd_dict=zstd_dict)
+            # Feed the frame in chunks of increasing size, so that what is
+            # read past its end (and copied into unused_data) is bounded by
+            # the size of the frame, not by the size of all the frames that
+            # follow it.
+            while not decomp.eof and pos < len(data):
+                with data[pos:pos + chunk_size] as chunk:
+                    results.append(decomp.decompress(chunk))
+                    pos += len(chunk)
+                chunk_size *= 2
+            if not decomp.eof:
+                raise ZstdError('Compressed data ended before the '
+                                'end-of-stream marker was reached')
+            pos -= len(decomp.unused_data)
+            if pos == len(data):
+                break
+            # Later frames start small: a larger first chunk reads further
+            # past the end of a short frame, a smaller one needs more calls.
+            chunk_size = 1024
     return b''.join(results)
 
 

Lib/gzip.py

@@ -682,24 +682,48 @@ def decompress(data):
     Return the decompressed string.
     """
     decompressed_members = []
+    fp = io.BytesIO(data)
+    # There is usually a single member: the first chunk is all of the data,
+    # to decompress it in one go.
+    chunk_size = fp.seek(0, io.SEEK_END)
+    fp.seek(0)
     while True:
-        fp = io.BytesIO(data)
         if _read_gzip_header(fp) is None:
             return b"".join(decompressed_members)
         # Use a zlib raw deflate compressor
         do = zlib.decompressobj(wbits=-zlib.MAX_WBITS)
-        # Read all the data except the header
-        decompressed = do.decompress(data[fp.tell():])
-        if not do.eof or len(do.unused_data) < 8:
+        crc = zlib.crc32(b"")
+        length = 0
+        # Feed the member in chunks of increasing size, so that what is read
+        # past its end (and copied into unused_data) is bounded by the size
+        # of the member, not by the size of all the members that follow it.
+        while not do.eof and (chunk := fp.read(chunk_size)):
+            decompressed = do.decompress(chunk)
+            crc = zlib.crc32(decompressed, crc)
+            length += len(decompressed)
+            decompressed_members.append(decompressed)
+            chunk_size *= 2
+        # Later members start small: a larger first chunk reads further past
+        # the end of a short member, a smaller one needs more calls.
+        chunk_size = 1024
+        # Rewind to the end of the deflate stream and read the trailer
+        fp.seek(-len(do.unused_data), io.SEEK_CUR)
+        trailer = fp.read(8)
+        if not do.eof or len(trailer) < 8:
             raise EOFError("Compressed file ended before the end-of-stream "
                            "marker was reached")
-        crc, length = struct.unpack("<II", do.unused_data[:8])
-        if crc != zlib.crc32(decompressed):
+        crc32, isize = struct.unpack("<II", trailer)
+        if crc32 != crc:
             raise BadGzipFile("CRC check failed")
-        if length != (len(decompressed) & 0xffffffff):
+        if isize != (length & 0xffffffff):
             raise BadGzipFile("Incorrect length of data produced")
-        decompressed_members.append(decompressed)
-        data = do.unused_data[8:].lstrip(b"\x00")
+        # Gzip files can be padded with zeroes and still have archives.
+        # Consume all zero bytes and set the file position to the first
+        # non-zero byte.
+        while chunk := fp.read(chunk_size):
+            if chunk := chunk.lstrip(b"\x00"):
+                fp.seek(-len(chunk), io.SEEK_CUR)
+                break
 
 
 def main():

Lib/lzma.py

@@ -347,20 +347,40 @@ def decompress(data, format=FORMAT_AUTO, memlimit=None, filters=None):
     For incremental decompression, use an LZMADecompressor instead.
     """
     results = []
-    while True:
-        decomp = LZMADecompressor(format, memlimit, filters)
-        try:
-            res = decomp.decompress(data)
-        except LZMAError:
-            if results:
-                break  # Leftover data is not a valid LZMA/XZ stream; ignore it.
-            else:
-                raise  # Error on the first iteration; bail out.
-        results.append(res)
-        if not decomp.eof:
-            raise LZMAError("Compressed data ended before the "
-                            "end-of-stream marker was reached")
-        data = decomp.unused_data
-        if not data:
-            break
+    # The views are released on exit, so that the caller can resize a
+    # bytearray argument again, also while handling an exception.
+    with memoryview(data) as view, view.cast("B") as data:
+        pos = 0
+        # There is usually a single stream: the first chunk is all of the
+        # data, to decompress it in one go.
+        chunk_size = len(data)
+        while True:
+            decomp = LZMADecompressor(format, memlimit, filters)
+            res = []
+            try:
+                # Feed the stream in chunks of increasing size, so that what
+                # is read past its end (and copied into unused_data) is
+                # bounded by the size of the stream, not by the size of all
+                # the streams that follow it.
+                while not decomp.eof and pos < len(data):
+                    with data[pos:pos + chunk_size] as chunk:
+                        res.append(decomp.decompress(chunk))
+                        pos += len(chunk)
+                    chunk_size *= 2
+            except LZMAError:
+                if results:
+                    # Leftover data is not a valid LZMA/XZ stream; ignore it.
+                    break
+                else:
+                    raise  # Error on the first iteration; bail out.
+            results.extend(res)
+            if not decomp.eof:
+                raise LZMAError("Compressed data ended before the "
+                                "end-of-stream marker was reached")
+            pos -= len(decomp.unused_data)
+            if pos == len(data):
+                break
+            # Later streams start small: a larger first chunk reads further
+            # past the end of a short stream, a smaller one needs more calls.
+            chunk_size = 1024
     return b"".join(results)

Lib/test/test_bz2.py

@@ -1100,6 +1100,49 @@ def testDecompressMultiStreamTrailingJunk(self):
         text = bz2.decompress(self.DATA * 5 + self.BAD_DATA)
         self.assertEqual(text, self.TEXT * 5)
 
+    def testDecompressMultiStreamBig(self):
+        big = random.randbytes(100_000)
+        streams = [self.TEXT, b'', big, self.TEXT, big]
+        expected = b''.join(streams)
+        data = b''.join(map(bz2.compress, streams))
+        self.assertEqual(bz2.decompress(data), expected)
+        self.assertEqual(bz2.decompress(bytearray(data)), expected)
+        self.assertEqual(bz2.decompress(memoryview(data)), expected)
+        self.assertEqual(bz2.decompress(data + self.BAD_DATA), expected)
+        self.assertRaises(ValueError, bz2.decompress, data[:-50_000])
+
+    def testDecompressErrorReleasesBuffer(self):
+        # The argument can be resized while the exception is handled.
+        for data in self.DATA[:-10], self.BAD_DATA:
+            with self.subTest(data=data[:10]):
+                data = bytearray(data)
+                try:
+                    bz2.decompress(data)
+                except (ValueError, OSError):
+                    data += b'more'
+                else:
+                    self.fail('no exception')
+
+    def testDecompressManyStreams(self):
+        # The amount of data passed to the decompressors should grow
+        # linearly with the number of streams.
+        class Decompressor:
+            def __init__(self, *args, **kwargs):
+                self.decomp = BZ2Decompressor(*args, **kwargs)
+            def __getattr__(self, name):
+                return getattr(self.decomp, name)
+            def decompress(self, data):
+                nonlocal fed
+                fed += len(data)
+                return self.decomp.decompress(data)
+        with support.swap_attr(bz2, 'BZ2Decompressor', Decompressor):
+            fed = 0
+            self.assertEqual(bz2.decompress(self.EMPTY_DATA * 1000), b'')
+            fed1 = fed
+            fed = 0
+            self.assertEqual(bz2.decompress(self.EMPTY_DATA * 2000), b'')
+        self.assertLess(fed, fed1 * 3)
+
 
 class OpenTest(BaseTest):
     "Test the open function."

Lib/test/test_gzip.py

@@ -827,6 +827,50 @@ def test_decompress_missing_trailer(self):
         compressed_data = gzip.compress(data1)
         self.assertRaises(EOFError, gzip.decompress, compressed_data[:-8])
 
+    def test_decompress_multi_member(self):
+        big = os.urandom(100_000)
+        members = [data1, b'', big, data2, big]
+        expected = b''.join(members)
+        for padding in b'', b'\0', b'\0' * 5000:
+            with self.subTest(padding=len(padding)):
+                datac = b''.join(gzip.compress(data) + padding
+                                 for data in members)
+                self.assertEqual(gzip.decompress(datac), expected)
+        datac = b''.join(map(gzip.compress, members))
+        self.assertEqual(gzip.decompress(bytearray(datac)), expected)
+        self.assertEqual(gzip.decompress(memoryview(datac)), expected)
+        self.assertRaises(EOFError, gzip.decompress, datac[:-4])
+        self.assertRaises(EOFError, gzip.decompress, datac[:-50_000])
+        self.assertRaises(gzip.BadGzipFile, gzip.decompress, datac + b'junk')
+        # Damage the CRC in the trailer of the third member.
+        parts = [gzip.compress(data) for data in members]
+        corrupted = bytearray(b''.join(parts))
+        corrupted[len(b''.join(parts[:3])) - 8] ^= 1
+        with self.assertRaisesRegex(gzip.BadGzipFile, 'CRC check failed'):
+            gzip.decompress(corrupted)
+
+    def test_decompress_many_members(self):
+        # The amount of data passed to the decompressors should grow
+        # linearly with the number of members.
+        decompressobj = zlib.decompressobj
+        class Decompressor:
+            def __init__(self, *args, **kwargs):
+                self.decomp = decompressobj(*args, **kwargs)
+            def __getattr__(self, name):
+                return getattr(self.decomp, name)
+            def decompress(self, data):
+                nonlocal fed
+                fed += len(data)
+                return self.decomp.decompress(data)
+        datac = gzip.compress(b'')
+        with mock.patch.object(zlib, 'decompressobj', Decompressor):
+            fed = 0
+            self.assertEqual(gzip.decompress(datac * 1000), b'')
+            fed1 = fed
+            fed = 0
+            self.assertEqual(gzip.decompress(datac * 2000), b'')
+        self.assertLess(fed, fed1 * 3)
+
     def test_truncated_header(self):
         truncated_headers = [
             b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00",             # Missing OS byte

Lib/test/test_lzma.py

@@ -514,6 +514,50 @@ def test_decompress_multistream(self):
         ddata = lzma.decompress(COMPRESSED_XZ + COMPRESSED_ALONE)
         self.assertEqual(ddata, INPUT * 2)
 
+    def test_decompress_multistream_big(self):
+        big = random.randbytes(100_000)
+        streams = [INPUT, b"", big, INPUT, big]
+        expected = b"".join(streams)
+        cdata = b"".join(map(lzma.compress, streams))
+        self.assertEqual(lzma.decompress(cdata), expected)
+        self.assertEqual(lzma.decompress(bytearray(cdata)), expected)
+        self.assertEqual(lzma.decompress(memoryview(cdata)), expected)
+        self.assertEqual(lzma.decompress(cdata + COMPRESSED_BOGUS), expected)
+        self.assertRaises(LZMAError, lzma.decompress, cdata[:-50_000])
+
+    def test_decompress_error_releases_buffer(self):
+        # The argument can be resized while the exception is handled.
+        for cdata in COMPRESSED_XZ[:-10], COMPRESSED_BOGUS:
+            with self.subTest(cdata=cdata[:10]):
+                cdata = bytearray(cdata)
+                try:
+                    lzma.decompress(cdata)
+                except LZMAError:
+                    cdata += b"more"
+                else:
+                    self.fail("no exception")
+
+    def test_decompress_many_streams(self):
+        # The amount of data passed to the decompressors should grow
+        # linearly with the number of streams.
+        class Decompressor:
+            def __init__(self, *args, **kwargs):
+                self.decomp = LZMADecompressor(*args, **kwargs)
+            def __getattr__(self, name):
+                return getattr(self.decomp, name)
+            def decompress(self, data):
+                nonlocal fed
+                fed += len(data)
+                return self.decomp.decompress(data)
+        cdata = lzma.compress(b"")
+        with support.swap_attr(lzma, "LZMADecompressor", Decompressor):
+            fed = 0
+            self.assertEqual(lzma.decompress(cdata * 1000), b"")
+            fed1 = fed
+            fed = 0
+            self.assertEqual(lzma.decompress(cdata * 2000), b"")
+        self.assertLess(fed, fed1 * 3)
+
     # Test robust handling of non-LZMA data following the compressed stream(s).
 
     def test_decompress_trailing_junk(self):