Address review: update offsets only after moving each entry

Restoring every offset after a failed repack() is wrong once an entry
has been moved over its old position: the restored offset then points
at overwritten data.  Instead, set each entry's header_offset only
after its bytes have been copied, so entries that finished moving or
were never touched stay readable.

This also drops the list(removed) copy, which was only needed because
the rollback read removed a second time.

Author: fedonman

Lib/test/test_zipfile/test_core.py

@@ -2664,28 +2664,10 @@ def test_overlap_with_archive_comment(self):
         fz.seek(0)
         self.assertEqual(fz.read(), expected)
 
-    def test_repack_removed_iterator(self):
-        # A one-shot iterable of removed entries must still be repacked.
-        buf = io.BytesIO()
-        with zipfile.ZipFile(buf, 'w') as zh:
-            for name in ['a.txt', 'b.txt', 'c.txt']:
-                zh.writestr(name, name[0].upper().encode() * 5000)
-        full_size = len(buf.getvalue())
-
-        fz = io.BytesIO(buf.getvalue())
-        with zipfile.ZipFile(fz, 'a') as zh:
-            zh.repack(iter([zh.remove('b.txt')]))
-
-        self.assertLess(len(fz.getvalue()), full_size - 5000)
-        fz.seek(0)
-        with zipfile.ZipFile(fz) as zh:
-            self.assertIsNone(zh.testzip())
-            self.assertEqual(zh.namelist(), ['a.txt', 'c.txt'])
-
-    def test_repack_failed_write_restores_header_offsets(self):
-        # A failed repack() must leave the in-memory offsets describing the
-        # file that is actually on disk, so that a later close() does not
-        # commit a central directory pointing at data that was never written.
+    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
 
@@ -2696,34 +2678,34 @@ def write(self, data):
                         raise OSError(errno.ENOSPC, 'No space left on device')
                 return super().write(data)
 
-        names = ['a.txt', 'b.txt', 'c.txt', 'd.txt']
-        buf = io.BytesIO()
-        with zipfile.ZipFile(buf, 'w') as zh:
-            for name in names:
-                zh.writestr(name, name[0].upper().encode() * 5000)
-
-        fz = FlakyBytesIO(buf.getvalue())
-        with zipfile.ZipFile(fz, 'a') as zh:
-            zi = zh.remove('b.txt')
-            expected = [z.header_offset for z in zh.infolist()]
-            expected_removed = zi.header_offset
-
-            fz.countdown = 1
-            with self.assertRaises(OSError):
-                zh.repack([zi], chunk_size=4096)
-
-            self.assertEqual([z.header_offset for z in zh.infolist()], expected)
-            self.assertEqual(zi.header_offset, expected_removed)
-
-            # The archive is still usable once the write error clears.
-            fz.countdown = None
+        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:
-            self.assertIsNone(zh.testzip())
-            self.assertEqual(zh.namelist(), ['a.txt', 'c.txt', 'd.txt'])
-            for name in zh.namelist():
-                self.assertEqual(zh.read(name), name[0].upper().encode() * 5000)
+                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,

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
@@ -2441,23 +2443,12 @@ def repack(self, removed=None, *, strict_descriptor=True,
                 )
 
             self._writing = True
-            # *removed* is read twice below, so a one-shot iterable must not
-            # be consumed by the first read.  None (scan the archive) is not
-            # the same as an empty sequence, so keep it as is.
-            if removed is not None:
-                removed = list(removed)
-            header_offsets = [(zinfo, zinfo.header_offset)
-                              for zinfo in (*self.filelist, *(removed or ()))]
             try:
                 repacker = _ZipRepacker(
                     strict_descriptor=strict_descriptor,
                     chunk_size=chunk_size,
                 )
                 repacker.repack(self, removed)
-            except BaseException:
-                for zinfo, header_offset in header_offsets:
-                    zinfo.header_offset = header_offset
-                raise
             finally:
                 self._writing = False