Skip to content

Commit c4e4413

Browse files
committed
gh-88300: Support .xz stream padding in lzma.decompress() and LZMAFile
Null bytes in multiples of four after an .xz stream are now skipped, so that a stream following them is decoded instead of ignored.
1 parent 3c23ebf commit c4e4413

4 files changed

Lines changed: 121 additions & 4 deletions

File tree

‎Lib/compression/_common/_streams.py‎

Lines changed: 61 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ def readable(self):
3737
return True
3838

3939
def __init__(self, fp, decomp_factory, trailing_error=(), *,
40-
stream_magic=(), **decomp_args):
40+
stream_magic=(), stream_padding=False, **decomp_args):
4141
self._fp = fp
4242
self._eof = False
4343
self._pos = 0 # Current offset in decompressed stream
@@ -62,6 +62,12 @@ def __init__(self, fp, decomp_factory, trailing_error=(), *,
6262
# stream rather than trailing data, so its error is not ignored.
6363
self._stream_magic = stream_magic
6464

65+
# Whether a stream that starts with stream_magic (bytes) may be
66+
# followed by Stream Padding, as in the .xz format. _head holds the
67+
# first bytes of the current stream.
68+
self._stream_padding = stream_padding and bool(stream_magic)
69+
self._head = b""
70+
6571
def close(self):
6672
self._decompressor = None
6773
return super().close()
@@ -88,11 +94,15 @@ def read(self, size=-1):
8894
if self._decompressor.eof:
8995
rawblock = (self._decompressor.unused_data or
9096
self._fp.read(BUFFER_SIZE))
97+
if self._stream_padding and self._head == self._stream_magic:
98+
rawblock = self._skip_padding(rawblock)
9199
if not rawblock:
92100
break
93101
# Continue to next stream.
94102
self._decompressor = self._decomp_factory(
95103
**self._decomp_args)
104+
if self._stream_padding:
105+
self._head = rawblock[:len(self._stream_magic)]
96106
try:
97107
data = self._decompressor.decompress(rawblock, size)
98108
except self._trailing_error:
@@ -109,6 +119,10 @@ def read(self, size=-1):
109119
"end-of-stream marker was reached")
110120
else:
111121
rawblock = b""
122+
if (self._stream_padding
123+
and len(self._head) < len(self._stream_magic)):
124+
self._head += rawblock[:len(self._stream_magic)
125+
- len(self._head)]
112126
data = self._decompressor.decompress(rawblock, size)
113127
if data:
114128
break
@@ -129,12 +143,30 @@ def readall(self):
129143

130144
return b"".join(chunks)
131145

146+
def _skip_padding(self, rawblock):
147+
# Skip the null bytes at the start of the data following a stream
148+
# if they are Stream Padding, whose size is a multiple of four.
149+
size = 0
150+
while True:
151+
rest = rawblock.lstrip(b"\0")
152+
size += len(rawblock) - len(rest)
153+
if rest:
154+
break
155+
rawblock = self._fp.read(BUFFER_SIZE)
156+
if not rawblock:
157+
break
158+
if size % 4:
159+
# Not Stream Padding: leave the null bytes to the decompressor.
160+
return bytes(size) + rest
161+
return rest
162+
132163
# Rewind the file to the beginning of the data stream.
133164
def _rewind(self):
134165
self._fp.seek(0)
135166
self._eof = False
136167
self._pos = 0
137168
self._decompressor = self._decomp_factory(**self._decomp_args)
169+
self._head = b""
138170

139171
def seek(self, offset, whence=io.SEEK_SET):
140172
# Recalculate offset as an absolute file position.
@@ -171,15 +203,37 @@ def tell(self):
171203
return self._pos
172204

173205

206+
def _padding_size(data, pos):
207+
"""Return the size of the Stream Padding at *data[pos:]*, or 0.
208+
209+
Stream Padding is a run of null bytes whose size is a multiple of four.
210+
"""
211+
end = pos
212+
chunk_size = 4
213+
while end < len(data):
214+
with data[end:end + chunk_size] as chunk:
215+
block = chunk.tobytes()
216+
rest = len(block.lstrip(b"\0"))
217+
end += len(block) - rest
218+
if rest:
219+
break
220+
chunk_size *= 2
221+
size = end - pos
222+
return 0 if size % 4 else size
223+
224+
174225
def _decompress_all(data, decomp_factory, trailing_error=(), *,
175-
stream_magic=(), truncated_error, **decomp_args):
226+
stream_magic=(), stream_padding=False, truncated_error,
227+
**decomp_args):
176228
"""Decompress all of the concatenated streams in *data*.
177229
178230
Each stream is decompressed by a new decompressor, made by calling
179231
*decomp_factory* with *decomp_args*. As in DecompressReader, an error
180232
of class *trailing_error* in a stream after the first one ends the
181233
data, unless that stream starts with *stream_magic*. A stream without
182-
an end-of-stream marker raises *truncated_error*.
234+
an end-of-stream marker raises *truncated_error*. If *stream_padding*
235+
is true, a stream that starts with *stream_magic* (bytes) may be
236+
followed by Stream Padding, as in the .xz format.
183237
"""
184238
results = []
185239
# The views are released on exit, so that the caller can resize a
@@ -220,6 +274,10 @@ def _decompress_all(data, decomp_factory, trailing_error=(), *,
220274
raise truncated_error("Compressed data ended before the "
221275
"end-of-stream marker was reached")
222276
pos -= len(decomp.unused_data)
277+
if stream_padding and stream_magic:
278+
with data[start:start + len(stream_magic)] as head:
279+
if head == stream_magic:
280+
pos += _padding_size(data, pos)
223281
if pos == len(data):
224282
break
225283
# Later streams start small: a larger first chunk reads further

‎Lib/lzma.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ def __init__(self, filename=None, mode="r", *,
138138
if self._mode == _MODE_READ:
139139
raw = _streams.DecompressReader(self._fp, LZMADecompressor,
140140
trailing_error=LZMAError, stream_magic=_stream_magic(format),
141-
format=format, filters=filters)
141+
stream_padding=True, format=format, filters=filters)
142142
self._buffer = io.BufferedReader(raw)
143143

144144
def close(self):
@@ -356,6 +356,7 @@ def decompress(data, format=FORMAT_AUTO, memlimit=None, filters=None):
356356
"""
357357
return _streams._decompress_all(data, LZMADecompressor, LZMAError,
358358
stream_magic=_stream_magic(format),
359+
stream_padding=True,
359360
truncated_error=LZMAError,
360361
format=format, memlimit=memlimit,
361362
filters=filters)

‎Lib/test/test_lzma.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,28 @@ def test_decompress_multistream_damaged_stream(self):
585585
self.assertRaises(LZMAError, lzma.decompress,
586586
COMPRESSED_XZ + damaged, format)
587587

588+
def test_decompress_stream_padding(self):
589+
# Null bytes in multiples of four can follow each .xz stream.
590+
cdata = (COMPRESSED_XZ + bytes(4) + COMPRESSED_XZ + bytes(20) +
591+
COMPRESSED_XZ + bytes(100_000) + COMPRESSED_XZ + bytes(8))
592+
for format in lzma.FORMAT_AUTO, lzma.FORMAT_XZ:
593+
with self.subTest(format=format):
594+
self.assertEqual(lzma.decompress(COMPRESSED_XZ + bytes(4),
595+
format), INPUT)
596+
self.assertEqual(lzma.decompress(cdata, format), INPUT * 4)
597+
# A damaged stream after padding is not trailing junk.
598+
damaged = bytearray(COMPRESSED_XZ)
599+
damaged[-20] ^= 1
600+
self.assertRaises(LZMAError, lzma.decompress,
601+
COMPRESSED_XZ + bytes(4) + damaged, format)
602+
self.assertRaises(LZMAError, lzma.decompress,
603+
bytes(4) + COMPRESSED_XZ, format)
604+
# Not after an .lzma stream.
605+
for format in lzma.FORMAT_AUTO, lzma.FORMAT_ALONE:
606+
with self.subTest(format=format):
607+
self.assertRaises(LZMAError, lzma.decompress,
608+
COMPRESSED_ALONE + bytes(4), format)
609+
588610

589611
class TempFile:
590612
"""Context manager - creates a file, and deletes it on __exit__."""
@@ -965,6 +987,39 @@ def test_read_multistream_damaged_stream(self):
965987
with LZMAFile(BytesIO(cdata), format=format) as f:
966988
self.assertRaises(LZMAError, f.seek, 0, 2)
967989

990+
def test_read_stream_padding(self):
991+
# Null bytes in multiples of four can follow each .xz stream,
992+
# whatever the size of the reads.
993+
cdata = (COMPRESSED_XZ + bytes(4) + COMPRESSED_XZ + bytes(20) +
994+
COMPRESSED_XZ + bytes(10_000) + COMPRESSED_XZ + bytes(8))
995+
damaged = bytearray(COMPRESSED_XZ)
996+
damaged[-20] ^= 1
997+
damaged = COMPRESSED_XZ + bytes(4) + damaged
998+
for buffer_size in _streams.BUFFER_SIZE, len(COMPRESSED_XZ) + 2, 3:
999+
for format in lzma.FORMAT_AUTO, lzma.FORMAT_XZ:
1000+
with (self.subTest(buffer_size=buffer_size, format=format),
1001+
support.swap_attr(_streams, "BUFFER_SIZE",
1002+
buffer_size)):
1003+
with LZMAFile(BytesIO(COMPRESSED_XZ + bytes(4)),
1004+
format=format) as f:
1005+
self.assertEqual(f.read(), INPUT)
1006+
with LZMAFile(BytesIO(cdata), format=format) as f:
1007+
self.assertEqual(f.read(), INPUT * 4)
1008+
self.assertEqual(f.seek(0), 0)
1009+
chunks = []
1010+
while result := f.read(10):
1011+
chunks.append(result)
1012+
self.assertEqual(b"".join(chunks), INPUT * 4)
1013+
# A damaged stream after padding is not trailing junk.
1014+
with LZMAFile(BytesIO(damaged), format=format) as f:
1015+
self.assertRaises(LZMAError, f.read)
1016+
# Not after an .lzma stream.
1017+
for format in lzma.FORMAT_AUTO, lzma.FORMAT_ALONE:
1018+
with self.subTest(format=format):
1019+
with LZMAFile(BytesIO(COMPRESSED_ALONE + bytes(4)),
1020+
format=format) as f:
1021+
self.assertRaises(EOFError, f.read)
1022+
9681023
def test_read_from_file(self):
9691024
with TempFile(TESTFN, COMPRESSED_XZ):
9701025
with LZMAFile(TESTFN) as f:
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix :func:`lzma.decompress` and :class:`lzma.LZMAFile` raising an error for
2+
``.xz`` stream padding (null bytes in multiples of four) at the end of the
3+
data, and ignoring any ``.xz`` streams that follow such padding.

0 commit comments

Comments
 (0)