diff --git a/Lib/test/test_zipfile/test_core.py b/Lib/test/test_zipfile/test_core.py index 84000f472c4df9..0edfa138acb1f0 100644 --- a/Lib/test/test_zipfile/test_core.py +++ b/Lib/test/test_zipfile/test_core.py @@ -2664,6 +2664,49 @@ def test_overlap_with_archive_comment(self): fz.seek(0) self.assertEqual(fz.read(), expected) + def test_repack_write_error_keeps_entries_readable(self): + # close() after a failed repack() writes the in-memory offsets to + # the central directory, so an entry's offset must only change once + # its data has been moved. + class FlakyBytesIO(io.BytesIO): + countdown = None + + def write(self, data): + if self.countdown is not None: + self.countdown -= 1 + if self.countdown < 0: + raise OSError(errno.ENOSPC, 'No space left on device') + return super().write(data) + + for sizes, countdown, readable in [ + # The error hits b, the first entry to move. The gap left by a + # is larger than one chunk, so b's old data is still intact. + ([5000, 5000, 5000, 5000], 1, 'bcd'), + # b has been moved when the error hits c. c is not checked: + # the gap left by a is small, so its move overwrote its old data. + ([100, 10000, 10000, 100], 4, 'bd'), + ]: + with self.subTest(sizes=sizes): + contents = {name: randbytes(size) + for name, size in zip('abcd', sizes)} + buf = io.BytesIO() + with zipfile.ZipFile(buf, 'w') as zh: + for name, data in contents.items(): + zh.writestr(name, data) + + fz = FlakyBytesIO(buf.getvalue()) + with zipfile.ZipFile(fz, 'a') as zh: + zi = zh.remove('a') + fz.countdown = countdown + with self.assertRaises(OSError): + zh.repack([zi], chunk_size=4096) + fz.countdown = None + + fz.seek(0) + with zipfile.ZipFile(fz) as zh: + for name in readable: + self.assertEqual(zh.read(name), contents[name]) + class ZipRepackerTests(unittest.TestCase): def _generate_local_file_entry(self, arcname, raw_bytes, compression=zipfile.ZIP_STORED, diff --git a/Lib/zipfile/__init__.py b/Lib/zipfile/__init__.py index 3110759b29f271..151179fa6d15a3 100644 --- a/Lib/zipfile/__init__.py +++ b/Lib/zipfile/__init__.py @@ -1581,17 +1581,18 @@ def repack(self, zfile, removed=None): entry_size = entry_size_list[i] used_entry_size = used_entry_size_list[i] - # update the header and move entry data to the new position + # move entry data to the new position, then update the header old_header_offset = zinfo.header_offset - zinfo.header_offset -= entry_offset + new_header_offset = old_header_offset - entry_offset if zinfo in removed_zinfos: self._copy_bytes( fp, old_header_offset + used_entry_size, - zinfo.header_offset, + new_header_offset, entry_size - used_entry_size, ) + zinfo.header_offset = new_header_offset # update entry_offset for subsequent files to follow entry_offset += used_entry_size @@ -1601,7 +1602,7 @@ def repack(self, zfile, removed=None): self._copy_bytes( fp, old_header_offset, - zinfo.header_offset, + new_header_offset, used_entry_size, ) @@ -1615,12 +1616,13 @@ def repack(self, zfile, removed=None): self._copy_bytes( fp, old_header_offset + used_entry_size + stale_entry_size, - zinfo.header_offset + used_entry_size, + new_header_offset + used_entry_size, entry_size - used_entry_size - stale_entry_size, ) # update entry_offset for subsequent files to follow entry_offset += stale_entry_size + zinfo.header_offset = new_header_offset # update state zfile.start_dir -= entry_offset