diff --git a/Lib/bz2.py b/Lib/bz2.py index 79305669efc38c9..f6f6223b65f4186 100644 --- a/Lib/bz2.py +++ b/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) diff --git a/Lib/compression/zstd/__init__.py b/Lib/compression/zstd/__init__.py index 22ce028c435caa8..757c11a82833b83 100644 --- a/Lib/compression/zstd/__init__.py +++ b/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) diff --git a/Lib/gzip.py b/Lib/gzip.py index 5b6ebcfef45c3ad..ca2834827cffaea 100644 --- a/Lib/gzip.py +++ b/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("