Skip to content

Commit 4be7eee

Browse files
committed
update
1 parent 1a80f6d commit 4be7eee

3 files changed

Lines changed: 54 additions & 3 deletions

File tree

‎Lib/test/test_zipfile/test_core.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -719,6 +719,42 @@ class ZstdTestsWithSourceFile(AbstractTestsWithSourceFile,
719719
unittest.TestCase):
720720
compression = zipfile.ZIP_ZSTANDARD
721721

722+
def test_read_multiframe(self):
723+
first = b'first frame' * 400
724+
second = b'second frame' * 400
725+
data = first + second
726+
727+
class MultiframeCompressor:
728+
def compress(self, value):
729+
value = bytes(value)
730+
return (zipfile.zstd.compress(value[:len(first)]) +
731+
zipfile.zstd.compress(value[len(first):]))
732+
733+
def flush(self):
734+
return b''
735+
736+
archive = io.BytesIO()
737+
with mock.patch.object(zipfile, '_get_compressor',
738+
return_value=MultiframeCompressor()):
739+
with zipfile.ZipFile(archive, 'w', self.compression) as zipfp:
740+
zipfp.writestr('multiframe', data)
741+
742+
with zipfile.ZipFile(archive) as zipfp:
743+
self.assertEqual(zipfp.read('multiframe'), data)
744+
745+
with zipfile.ZipFile(archive) as zipfp:
746+
with zipfp.open('multiframe') as fp:
747+
chunk = fp._read1(1)
748+
self.assertLessEqual(len(chunk), fp.MIN_READ_SIZE)
749+
self.assertEqual(chunk + fp.read(), data)
750+
751+
# Exercise a frame boundary between two reads from the archive.
752+
with zipfile.ZipFile(archive) as zipfp:
753+
with zipfp.open('multiframe') as fp:
754+
fp.MIN_READ_SIZE = 1
755+
self.assertEqual(b''.join(iter(lambda: fp.read(1), b'')),
756+
data)
757+
722758
class AbstractTestZip64InSmallFiles:
723759
# These tests test the ZIP64 functionality without using large files,
724760
# see test_zipfile64 for proper tests.

‎Lib/zipfile/__init__.py‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1195,6 +1195,13 @@ def _read1(self, n):
11951195
data = self._decompressor.unconsumed_tail
11961196
if n > len(data):
11971197
data += self._read2(n - len(data))
1198+
elif (self._compress_type == ZIP_ZSTANDARD and
1199+
self._decompressor.eof):
1200+
# Continue with the next Zstandard frame.
1201+
data = self._decompressor.unused_data
1202+
self._decompressor = _get_decompressor(self._compress_type)
1203+
if not data:
1204+
data = self._read2(n)
11981205
elif self._compress_type == ZIP_STORED:
11991206
data = self._read2(n)
12001207
else:
@@ -1233,9 +1240,15 @@ def _read1(self, n):
12331240
remove=(3, 21),
12341241
stacklevel=4)
12351242
data = self._decompressor.decompress(data)
1236-
self._eof = (self._decompressor.eof or
1237-
self._compress_left <= 0 and
1238-
getattr(self._decompressor, "needs_input", True))
1243+
if self._compress_type == ZIP_ZSTANDARD:
1244+
self._eof = (self._compress_left <= 0 and
1245+
not self._decompressor.unused_data and
1246+
(self._decompressor.eof or
1247+
self._decompressor.needs_input))
1248+
else:
1249+
self._eof = (self._decompressor.eof or
1250+
self._compress_left <= 0 and
1251+
getattr(self._decompressor, "needs_input", True))
12391252

12401253
data = data[:self._left]
12411254
self._left -= len(data)
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix :mod:`zipfile` to correctly read Zstandard-compressed archive members
2+
that could previously be truncated or incorrectly rejected.

0 commit comments

Comments
 (0)