Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 97 additions & 5 deletions Lib/compression/_common/_streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ def readable(self):
return True

def __init__(self, fp, decomp_factory, trailing_error=(), *,
stream_magic=(), **decomp_args):
stream_magic=(), xz_stream_padding=False, **decomp_args):
self._fp = fp
self._eof = False
self._pos = 0 # Current offset in decompressed stream
Expand All @@ -62,6 +62,20 @@ def __init__(self, fp, decomp_factory, trailing_error=(), *,
# stream rather than trailing data, so its error is not ignored.
self._stream_magic = stream_magic

# Whether to skip the Stream Padding of the .xz format: null bytes,
# in multiples of four, after a stream that starts with stream_magic.
# No other format here has such padding, so only lzma sets this, and
# stream_magic must then be a single bytes object, not a tuple.
# _head holds the first bytes of the current stream.
self._xz_stream_padding = xz_stream_padding and bool(stream_magic)
self._head = b""

# Null bytes that _skip_xz_padding() read past but did not skip
# (only their number is kept), and the data that was read after
# them. _read_block() returns these before reading the file again.
self._unread_nulls = 0
self._unread_data = b""

def close(self):
self._decompressor = None
return super().close()
Expand All @@ -87,12 +101,17 @@ def read(self, size=-1):
while True:
if self._decompressor.eof:
rawblock = (self._decompressor.unused_data or
self._fp.read(BUFFER_SIZE))
self._read_block())
if (self._xz_stream_padding
and self._head == self._stream_magic):
rawblock = self._skip_xz_padding(rawblock)
if not rawblock:
break
# Continue to next stream.
self._decompressor = self._decomp_factory(
**self._decomp_args)
if self._xz_stream_padding:
self._head = rawblock[:len(self._stream_magic)]
try:
data = self._decompressor.decompress(rawblock, size)
except self._trailing_error:
Expand All @@ -103,12 +122,16 @@ def read(self, size=-1):
break
else:
if self._decompressor.needs_input:
rawblock = self._fp.read(BUFFER_SIZE)
rawblock = self._read_block()
if not rawblock:
raise EOFError("Compressed file ended before the "
"end-of-stream marker was reached")
else:
rawblock = b""
if (self._xz_stream_padding
and len(self._head) < len(self._stream_magic)):
self._head += rawblock[:len(self._stream_magic)
- len(self._head)]
data = self._decompressor.decompress(rawblock, size)
if data:
break
Expand All @@ -129,12 +152,53 @@ def readall(self):

return b"".join(chunks)

def _read_block(self):
# Return the next block of compressed data: what _skip_xz_padding()
# left unread, if anything, and otherwise a block from the file.
if self._unread_nulls:
size = min(self._unread_nulls, BUFFER_SIZE)
self._unread_nulls -= size
return bytes(size)
if self._unread_data:
rawblock, self._unread_data = self._unread_data, b""
return rawblock
return self._fp.read(BUFFER_SIZE)

def _skip_xz_padding(self, rawblock):
# Skip the null bytes at the start of the data following an .xz
# stream if they are Stream Padding, whose size is a multiple of
# four. Return the block of data to go on with.
first = rawblock
size = 0
while True:
rest = rawblock.lstrip(b"\0")
size += len(rawblock) - len(rest)
if rest:
break
rawblock = self._read_block()
if not rawblock:
break
if size % 4 == 0:
return rest
# Not Stream Padding: the null bytes are left to the decompressor,
# in the blocks they were read in, as if nothing had been read past
# the first one. There can be any number of null bytes, so the
# blocks of them between the first and the last are not kept.
if rawblock is not first:
self._unread_nulls = (size - len(first)
- (len(rawblock) - len(rest)))
self._unread_data = rawblock
return first

# Rewind the file to the beginning of the data stream.
def _rewind(self):
self._fp.seek(0)
self._eof = False
self._pos = 0
self._decompressor = self._decomp_factory(**self._decomp_args)
self._head = b""
self._unread_nulls = 0
self._unread_data = b""

def seek(self, offset, whence=io.SEEK_SET):
# Recalculate offset as an absolute file position.
Expand Down Expand Up @@ -171,15 +235,39 @@ def tell(self):
return self._pos


def _xz_padding_size(data, pos):
"""Return the size of the .xz Stream Padding at *data[pos:]*, or 0.

Stream Padding is a run of null bytes whose size is a multiple of four.
"""
end = pos
# There is usually no padding, or little of it: the chunks to look at
# start small. Each one is copied, so their size is limited too.
chunk_size = 4
while end < len(data):
with data[end:end + chunk_size] as chunk:
block = chunk.tobytes()
rest = len(block.lstrip(b"\0"))
end += len(block) - rest
if rest:
break
chunk_size = min(chunk_size * 2, BUFFER_SIZE)
size = end - pos
return 0 if size % 4 else size


def _decompress_all(data, decomp_factory, trailing_error=(), *,
stream_magic=(), truncated_error, **decomp_args):
stream_magic=(), xz_stream_padding=False,
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*.
an end-of-stream marker raises *truncated_error*. If *xz_stream_padding*
is true, the Stream Padding of the .xz format is skipped after a stream
that starts with *stream_magic*, which must then be a bytes object.
"""
results = []
# The views are released on exit, so that the caller can resize a
Expand Down Expand Up @@ -220,6 +308,10 @@ def _decompress_all(data, decomp_factory, trailing_error=(), *,
raise truncated_error("Compressed data ended before the "
"end-of-stream marker was reached")
pos -= len(decomp.unused_data)
if xz_stream_padding and stream_magic:
with data[start:start + len(stream_magic)] as head:
if head == stream_magic:
pos += _xz_padding_size(data, pos)
if pos == len(data):
break
# Later streams start small: a larger first chunk reads further
Expand Down
3 changes: 2 additions & 1 deletion Lib/lzma.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ def __init__(self, filename=None, mode="r", *,
if self._mode == _MODE_READ:
raw = _streams.DecompressReader(self._fp, LZMADecompressor,
trailing_error=LZMAError, stream_magic=_stream_magic(format),
format=format, filters=filters)
xz_stream_padding=True, format=format, filters=filters)
self._buffer = io.BufferedReader(raw)

def close(self):
Expand Down Expand Up @@ -356,6 +356,7 @@ def decompress(data, format=FORMAT_AUTO, memlimit=None, filters=None):
"""
return _streams._decompress_all(data, LZMADecompressor, LZMAError,
stream_magic=_stream_magic(format),
xz_stream_padding=True,
truncated_error=LZMAError,
format=format, memlimit=memlimit,
filters=filters)
134 changes: 134 additions & 0 deletions Lib/test/test_lzma.py
Original file line number Diff line number Diff line change
Expand Up @@ -585,6 +585,55 @@ def test_decompress_multistream_damaged_stream(self):
self.assertRaises(LZMAError, lzma.decompress,
COMPRESSED_XZ + damaged, format)

def test_decompress_stream_padding(self):
# Null bytes in multiples of four can follow each .xz stream.
cdata = (COMPRESSED_XZ + bytes(4) + COMPRESSED_XZ + bytes(20) +
COMPRESSED_XZ + bytes(100_000) + COMPRESSED_XZ + bytes(8))
for format in lzma.FORMAT_AUTO, lzma.FORMAT_XZ:
with self.subTest(format=format):
self.assertEqual(lzma.decompress(COMPRESSED_XZ + bytes(4),
format), INPUT)
self.assertEqual(lzma.decompress(cdata, format), INPUT * 4)
# A damaged stream after padding is not trailing junk.
damaged = bytearray(COMPRESSED_XZ)
damaged[-20] ^= 1
self.assertRaises(LZMAError, lzma.decompress,
COMPRESSED_XZ + bytes(4) + damaged, format)
self.assertRaises(LZMAError, lzma.decompress,
bytes(4) + COMPRESSED_XZ, format)
# Not after an .lzma stream.
for format in lzma.FORMAT_AUTO, lzma.FORMAT_ALONE:
with self.subTest(format=format):
self.assertRaises(LZMAError, lzma.decompress,
COMPRESSED_ALONE + bytes(4), format)

def test_decompress_null_bytes_not_stream_padding(self):
# Null bytes that are not in multiples of four are not stream
# padding: they are decompressed like any other data.
def decompress(data, format, **kwargs):
try:
return _streams._decompress_all(
data, LZMADecompressor, LZMAError,
stream_magic=lzma._stream_magic(format),
truncated_error=LZMAError, format=format, **kwargs)
except LZMAError:
return LZMAError

for format in lzma.FORMAT_AUTO, lzma.FORMAT_XZ:
for size in 1, 3, 5, 13, 18, 30, 1001, 300_001:
for tail in b"", COMPRESSED_XZ, b"junk" * 4:
with self.subTest(format=format, size=size, tail=tail):
cdata = COMPRESSED_XZ + bytes(size) + tail
self.assertEqual(
decompress(cdata, format, xz_stream_padding=True),
decompress(cdata, format))
# A stream after them is not decompressed.
self.assertRaises(LZMAError, lzma.decompress,
COMPRESSED_XZ + bytes(5), lzma.FORMAT_XZ)
self.assertEqual(lzma.decompress(COMPRESSED_XZ + bytes(13) +
COMPRESSED_XZ, lzma.FORMAT_XZ),
INPUT)


class TempFile:
"""Context manager - creates a file, and deletes it on __exit__."""
Expand Down Expand Up @@ -965,6 +1014,91 @@ def test_read_multistream_damaged_stream(self):
with LZMAFile(BytesIO(cdata), format=format) as f:
self.assertRaises(LZMAError, f.seek, 0, 2)

def test_read_stream_padding(self):
# Null bytes in multiples of four can follow each .xz stream,
# whatever the size of the reads.
cdata = (COMPRESSED_XZ + bytes(4) + COMPRESSED_XZ + bytes(20) +
COMPRESSED_XZ + bytes(10_000) + COMPRESSED_XZ + bytes(8))
damaged = bytearray(COMPRESSED_XZ)
damaged[-20] ^= 1
damaged = COMPRESSED_XZ + bytes(4) + damaged
for buffer_size in _streams.BUFFER_SIZE, len(COMPRESSED_XZ) + 2, 3:
for format in lzma.FORMAT_AUTO, lzma.FORMAT_XZ:
with (self.subTest(buffer_size=buffer_size, format=format),
support.swap_attr(_streams, "BUFFER_SIZE",
buffer_size)):
with LZMAFile(BytesIO(COMPRESSED_XZ + bytes(4)),
format=format) as f:
self.assertEqual(f.read(), INPUT)
with LZMAFile(BytesIO(cdata), format=format) as f:
self.assertEqual(f.read(), INPUT * 4)
self.assertEqual(f.seek(0), 0)
chunks = []
while result := f.read(10):
chunks.append(result)
self.assertEqual(b"".join(chunks), INPUT * 4)
# A damaged stream after padding is not trailing junk.
with LZMAFile(BytesIO(damaged), format=format) as f:
self.assertRaises(LZMAError, f.read)
# Not after an .lzma stream.
for format in lzma.FORMAT_AUTO, lzma.FORMAT_ALONE:
with self.subTest(format=format):
with LZMAFile(BytesIO(COMPRESSED_ALONE + bytes(4)),
format=format) as f:
self.assertRaises(EOFError, f.read)

def test_read_null_bytes_not_stream_padding(self):
# Null bytes that are not in multiples of four are not stream
# padding. However many there are, the decompressors are given
# the same blocks of data as without support for stream padding.
class Decompressor:
def __init__(self, **kwargs):
self.decomp = LZMADecompressor(**kwargs)
def __getattr__(self, name):
return getattr(self.decomp, name)
def decompress(self, data, max_length):
blocks.append(len(data))
return self.decomp.decompress(data, max_length)

def read(data, format, **kwargs):
blocks.clear()
raw = _streams.DecompressReader(
BytesIO(data), Decompressor, trailing_error=LZMAError,
stream_magic=lzma._stream_magic(format), format=format,
**kwargs)
try:
result = raw.readall()
except (LZMAError, EOFError) as exc:
result = type(exc)
return result, blocks[:]

blocks = []
for buffer_size in 3, 64, len(COMPRESSED_XZ) + 2, 4096:
for format in lzma.FORMAT_AUTO, lzma.FORMAT_XZ:
for size in 1, 3, 5, 13, 18, 30, 1001, 10_001:
for tail in b"", COMPRESSED_XZ, b"junk" * 4:
with (self.subTest(buffer_size=buffer_size,
format=format, size=size,
tail=tail),
support.swap_attr(_streams, "BUFFER_SIZE",
buffer_size)):
cdata = COMPRESSED_XZ + bytes(size) + tail
result = read(cdata, format,
xz_stream_padding=True)
self.assertEqual(result, read(cdata, format))
self.assertLessEqual(max(result[1]), buffer_size)

# A stream after them is not read, also after seeking back.
cdata = COMPRESSED_XZ + bytes(300_001) + COMPRESSED_XZ
self.assertGreater(len(cdata), 2 * _streams.BUFFER_SIZE)
with LZMAFile(BytesIO(cdata), format=lzma.FORMAT_XZ) as f:
self.assertEqual(f.read(), INPUT)
self.assertEqual(f.seek(0), 0)
self.assertEqual(f.read(), INPUT)
with LZMAFile(BytesIO(COMPRESSED_XZ + bytes(5)),
format=lzma.FORMAT_XZ) as f:
self.assertRaises(EOFError, f.read)

def test_read_from_file(self):
with TempFile(TESTFN, COMPRESSED_XZ):
with LZMAFile(TESTFN) as f:
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Fix :func:`lzma.decompress` and :class:`lzma.LZMAFile` raising an error for
``.xz`` stream padding (null bytes in multiples of four) at the end of the
data, and ignoring any ``.xz`` streams that follow such padding.
Loading