Skip to content

Commit f897dbf

Browse files
encukoutonghuaroot
andauthored
gh-156002: Bound zipfile decompression for bzip2/LZMA/Zstandard (GH-156003)
Patch by @tonghuaroot. zipfile.ZipExtFile._read1() bounds the output of each decompress() call for DEFLATE members by passing a max_length to zlib, but for bzip2, LZMA, and Zstandard members it called decompress() with no bound. A whole compressed chunk was therefore expanded into a single allocation before the data[:self._left] clip ran, so a consumer that deliberately reads in small chunks to limit memory (for example zf.open(name).read(8192)) was silently unprotected for non-DEFLATE members. A small, spec-conformant archive member declaring a large uncompressed size could drive multi-GB peak memory. _read1() now passes a per-call bound to the non-DEFLATE decompress() (mirroring the DEFLATE branch) and drains the decompressor's internal buffer across calls by checking needs_input before reading more compressed input. zipfile's LZMADecompressor wrapper forwards max_length and exposes needs_input so the bound also holds for LZMA members. Co-authored-by: tonghuaroot <tonghuaroot@gmail.com>
1 parent 9cbd578 commit f897dbf

3 files changed

Lines changed: 79 additions & 5 deletions

File tree

Lib/test/test_zipfile/test_core.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4874,6 +4874,48 @@ def tearDown(self):
48744874
unlink(TESTFN2)
48754875

48764876

4877+
class AbstractBoundedDecompressTests:
4878+
# ZipExtFile._read1() bounds the output of each decompress() call so that a
4879+
# small member declaring a large uncompressed size cannot expand into one
4880+
# unbounded read.
4881+
def test_read1_output_is_bounded(self):
4882+
buf = io.BytesIO()
4883+
with zipfile.ZipFile(buf, "w", compression=self.compression) as zf:
4884+
zf.writestr("big", b"\0" * (4 * 1024 * 1024))
4885+
with zipfile.ZipFile(io.BytesIO(buf.getvalue())) as zf:
4886+
with zf.open("big") as f:
4887+
self.assertLessEqual(len(f._read1(100)), f.MIN_READ_SIZE)
4888+
4889+
4890+
class StoredBoundedDecompressTests(AbstractBoundedDecompressTests,
4891+
unittest.TestCase):
4892+
compression = zipfile.ZIP_STORED
4893+
4894+
4895+
@requires_zlib()
4896+
class DeflateBoundedDecompressTests(AbstractBoundedDecompressTests,
4897+
unittest.TestCase):
4898+
compression = zipfile.ZIP_DEFLATED
4899+
4900+
4901+
@requires_bz2()
4902+
class Bzip2BoundedDecompressTests(AbstractBoundedDecompressTests,
4903+
unittest.TestCase):
4904+
compression = zipfile.ZIP_BZIP2
4905+
4906+
4907+
@requires_lzma()
4908+
class LzmaBoundedDecompressTests(AbstractBoundedDecompressTests,
4909+
unittest.TestCase):
4910+
compression = zipfile.ZIP_LZMA
4911+
4912+
4913+
@requires_zstd()
4914+
class ZstdBoundedDecompressTests(AbstractBoundedDecompressTests,
4915+
unittest.TestCase):
4916+
compression = zipfile.ZIP_ZSTANDARD
4917+
4918+
48774919
class AbstractBadCrcTests:
48784920
def test_testzip_with_bad_crc(self):
48794921
"""Tests that files with bad CRCs return their name from testzip."""

Lib/zipfile/__init__.py

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -801,7 +801,16 @@ def unused_data(self):
801801
except AttributeError:
802802
return b''
803803

804-
def decompress(self, data):
804+
@property
805+
def _needs_input(self):
806+
# While the LZMA properties header is still being buffered, more input
807+
# is required; afterwards defer to the wrapped decompressor so a bounded
808+
# decompress() call can be drained across reads.
809+
if self._decomp is None:
810+
return True
811+
return self._decomp.needs_input
812+
813+
def decompress(self, data, max_length=-1):
805814
if self._decomp is None:
806815
self._unconsumed += data
807816
if len(self._unconsumed) <= 4:
@@ -817,7 +826,7 @@ def decompress(self, data):
817826
data = self._unconsumed[4 + psize:]
818827
del self._unconsumed
819828

820-
result = self._decomp.decompress(data)
829+
result = self._decomp.decompress(data, max_length)
821830
self.eof = self._decomp.eof
822831
return result
823832

@@ -884,6 +893,13 @@ def _get_compressor(compress_type, compresslevel=None):
884893
return None
885894

886895

896+
def _decompressor_needs_input(decompressor):
897+
# bz2/zstd expose the stdlib decompressor's public needs_input; the LZMA
898+
# wrapper keeps it private (_needs_input) to avoid adding public API.
899+
needs_input = getattr(decompressor, "needs_input", None)
900+
return decompressor._needs_input if needs_input is None else needs_input
901+
902+
887903
def _get_decompressor(compress_type):
888904
_check_compression(compress_type)
889905
if compress_type == ZIP_STORED:
@@ -1186,8 +1202,15 @@ def _read1(self, n):
11861202
data = self._decompressor.unconsumed_tail
11871203
if n > len(data):
11881204
data += self._read2(n - len(data))
1189-
else:
1205+
elif self._compress_type == ZIP_STORED:
11901206
data = self._read2(n)
1207+
else:
1208+
# bzip2/lzma/zstd: a bounded decompress() call may leave input
1209+
# buffered inside the decompressor; drain that before reading more.
1210+
if _decompressor_needs_input(self._decompressor):
1211+
data = self._read2(n)
1212+
else:
1213+
data = b''
11911214

11921215
if self._compress_type == ZIP_STORED:
11931216
self._eof = self._compress_left <= 0
@@ -1200,8 +1223,13 @@ def _read1(self, n):
12001223
if self._eof:
12011224
data += self._decompressor.flush()
12021225
else:
1203-
data = self._decompressor.decompress(data)
1204-
self._eof = self._decompressor.eof or self._compress_left <= 0
1226+
# Bound the output of a single decompress() call (mirroring the
1227+
# DEFLATE path above) so that a small compressed member cannot
1228+
# expand into one unbounded read.
1229+
data = self._decompressor.decompress(data, max(n, self.MIN_READ_SIZE))
1230+
self._eof = (self._decompressor.eof or
1231+
self._compress_left <= 0 and
1232+
_decompressor_needs_input(self._decompressor))
12051233

12061234
data = data[:self._left]
12071235
self._left -= len(data)
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Bound the amount of data :mod:`zipfile` decompresses per read for members
2+
compressed with bzip2, LZMA, or Zstandard, matching the existing limit for
3+
deflate. A small archive member could previously expand into an unbounded
4+
allocation even when read in small chunks.

0 commit comments

Comments
 (0)