diff --git a/Lib/bz2.py b/Lib/bz2.py index f4f3cda5a66cbd8..7e41866004f7518 100644 --- a/Lib/bz2.py +++ b/Lib/bz2.py @@ -338,20 +338,8 @@ def decompress(data): For incremental decompression, use a BZ2Decompressor object instead. """ - results = [] - while data: - decomp = BZ2Decompressor() - try: - res = decomp.decompress(data) - except OSError: - if not results: - raise # Error on the first iteration; bail out. - if data.startswith(_STREAM_MAGIC): - raise # Leftover data is a damaged bzip2 stream; bail out. - break # Leftover data is not a valid bzip2 stream; ignore it. - results.append(res) - if not decomp.eof: - raise ValueError("Compressed data ended before the " - "end-of-stream marker was reached") - data = decomp.unused_data - return b"".join(results) + if not data: + return b"" + return _streams._decompress_all(data, BZ2Decompressor, OSError, + stream_magic=_STREAM_MAGIC, + truncated_error=ValueError) diff --git a/Lib/compression/_common/_streams.py b/Lib/compression/_common/_streams.py index 09c65507297886b..0e3308c0c6e5fdc 100644 --- a/Lib/compression/_common/_streams.py +++ b/Lib/compression/_common/_streams.py @@ -169,3 +169,60 @@ def seek(self, offset, whence=io.SEEK_SET): def tell(self): """Return the current file position.""" return self._pos + + +def _decompress_all(data, decomp_factory, trailing_error=(), *, + stream_magic=(), truncated_error, **decomp_args): + """Decompress all of the concatenated streams in *data*. + + Each stream is decompressed by a new decompressor, made by calling + *decomp_factory* with *decomp_args*. As in DecompressReader, an error + of class *trailing_error* in a stream after the first one ends the + data, unless that stream starts with *stream_magic*. A stream without + an end-of-stream marker raises *truncated_error*. + """ + results = [] + # 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 = decomp_factory(**decomp_args) + start = pos + 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 trailing_error: + if not results: + raise # Error in the first stream; bail out. + if stream_magic: + if isinstance(stream_magic, bytes): + stream_magic = (stream_magic,) + size = max(map(len, stream_magic)) + with data[start:start + size] as head: + if head.tobytes().startswith(stream_magic): + raise # Damaged stream, not trailing data. + # Trailing data isn't a valid compressed stream; ignore it. + break + results.extend(res) + if not decomp.eof: + raise truncated_error("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 5326cf9b19cf883..b1b43cee15111f1 100644 --- a/Lib/compression/zstd/__init__.py +++ b/Lib/compression/zstd/__init__.py @@ -30,6 +30,7 @@ import enum from _zstd import (ZstdCompressor, ZstdDecompressor, ZstdDict, ZstdError, get_frame_size, zstd_version) +from compression._common import _streams from compression.zstd._zstdfile import ZstdFile, open, _nbytes # zstd_version_number is (MAJOR * 100 * 100 + MINOR * 100 + RELEASE) @@ -161,17 +162,9 @@ 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 - return b''.join(results) + return _streams._decompress_all(data, ZstdDecompressor, + truncated_error=ZstdError, + options=options, zstd_dict=zstd_dict) class CompressionParameter(enum.IntEnum): diff --git a/Lib/gzip.py b/Lib/gzip.py index b50089c31fadc86..a478329aef77f7a 100644 --- a/Lib/gzip.py +++ b/Lib/gzip.py @@ -681,24 +681,47 @@ 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 = len(data) 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("