diff --git a/scripts/generate_pack.py b/scripts/generate_pack.py index 733d4bd8..a88958c6 100644 --- a/scripts/generate_pack.py +++ b/scripts/generate_pack.py @@ -4207,41 +4207,45 @@ def _manifest_core_index( def inject_manifest(zip_path: str, manifest: dict) -> None: - """Inject manifest.json into an existing ZIP pack.""" + """Inject manifest.json into an existing ZIP pack. + + The pack is rewritten beside its final name and swapped in, as + _writing_pack builds it: appending in place rewrites the central + directory last, so a process killed in between left an unreadable ZIP + where a valid pack had been. + """ manifest_json = json.dumps(manifest, indent=2, ensure_ascii=False) - # Check if manifest already exists with zipfile.ZipFile(zip_path, "r") as zf: has_manifest = "manifest.json" in zf.namelist() - if not has_manifest: - # Fast path: append directly. Append mode defaults to ZIP_STORED, - # which would leave the manifest uncompressed and give the same pack - # two different byte layouts depending on the path taken here. - with zipfile.ZipFile(zip_path, "a", zipfile.ZIP_DEFLATED) as zf: - _write_generated_member(zf, "manifest.json", manifest_json) - else: - # Rebuild to replace existing manifest - import tempfile as _tempfile - - tmp_fd, tmp_path = _tempfile.mkstemp( - suffix=".zip", dir=os.path.dirname(zip_path) - ) - os.close(tmp_fd) - try: + directory = os.path.dirname(zip_path) or "." + fd, scratch = tempfile.mkstemp( + prefix=f".{os.path.basename(zip_path)}.", suffix=".partial", dir=directory + ) + os.close(fd) + try: + if not has_manifest: + # Append to a copy. Append mode defaults to ZIP_STORED, which + # would leave the manifest uncompressed and give the same pack + # two different byte layouts depending on the path taken here. + shutil.copyfile(zip_path, scratch) + with zipfile.ZipFile(scratch, "a", zipfile.ZIP_DEFLATED) as zf: + _write_generated_member(zf, "manifest.json", manifest_json) + else: with ( zipfile.ZipFile(zip_path, "r") as src, - zipfile.ZipFile(tmp_path, "w", zipfile.ZIP_DEFLATED) as dst, + zipfile.ZipFile(scratch, "w", zipfile.ZIP_DEFLATED) as dst, ): for item in src.infolist(): if item.filename == "manifest.json": continue dst.writestr(item, src.read(item.filename)) _write_generated_member(dst, "manifest.json", manifest_json) - os.replace(tmp_path, zip_path) - except (OSError, zipfile.BadZipFile): - os.unlink(tmp_path) - raise + os.replace(scratch, zip_path) + except BaseException: + Path(scratch).unlink(missing_ok=True) + raise diff --git a/tests/test_deterministic_zip.py b/tests/test_deterministic_zip.py index a2eb91a6..78653f49 100644 --- a/tests/test_deterministic_zip.py +++ b/tests/test_deterministic_zip.py @@ -3,6 +3,7 @@ from __future__ import annotations import hashlib +import os import tempfile import time import unittest @@ -377,5 +378,40 @@ class TestPackDeterminism(unittest.TestCase): self.assertEqual(info.date_time, _FIXED_DATE_TIME) + +class ManifestInjectionIsAtomic(unittest.TestCase): + """inject_manifest appended to the final pack in place: zipfile writes + the member over the old central directory and the new one only on + close, so a process killed in between left a BadZipFile where a valid + pack had been.""" + + def test_a_kill_mid_write_leaves_the_pack_readable(self): + import tempfile + from unittest import mock + + from scripts import generate_pack as gp + + with tempfile.TemporaryDirectory(dir=Path(__file__).resolve().parents[1] / "tmp") as d: + pack = Path(d) / "Pack.zip" + with zipfile.ZipFile(pack, "w", zipfile.ZIP_DEFLATED) as zf: + for i in range(20): + zf.writestr(f"bios/{i}.bin", os.urandom(4096)) + before = pack.read_bytes() + + def die(zf, name, text): + zf.writestr(name, text) + zf.fp.flush() + raise KeyboardInterrupt + + with mock.patch.object(gp, "_write_generated_member", die), \ + self.assertRaises(KeyboardInterrupt): + gp.inject_manifest(str(pack), {"files": ["x"] * 2000}) + self.assertEqual(pack.read_bytes(), before) + self.assertEqual( + [p.name for p in Path(d).iterdir()], ["Pack.zip"], + "the scratch copy is removed", + ) + + if __name__ == "__main__": unittest.main()