Cap zstd frames at 200 MiB to fix JS-decoder overflow on large .eval files

src/inspect_ai/_util/zipfile.py

@@ -1,17 +1,183 @@
+"""ZIP file helpers and monkey-patches for zstd support.
+
+On Python < 3.14 we import ``zipfile_zstd`` to monkey-patch zstandard
+compression into the stdlib ``zipfile`` module. Python 3.14+ handles zstd
+natively via stdlib.
+
+Additionally, we install a second monkey-patch on ``zipfile._get_compressor``
+that caps each emitted zstd frame at ``_MAX_INPUT_PER_FRAME`` bytes of input.
+Large single zstd frames (>= 256 MiB compressed) trigger an overflow bug in
+the pure-JS ``fzstd`` decoder used by our viewers; multi-framing keeps every
+frame well under that threshold.
+"""
+
+from __future__ import annotations
+
 import logging
 import sys
 import zipfile
-from typing import Any
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import zstandard
 
 logger = logging.getLogger(__name__)
 
 # On Python < 3.14, monkey-patch zipfile to support zstandard compression.
 if sys.version_info < (3, 14):
     import zipfile_zstd  # type: ignore[import-not-found, import-untyped]  # noqa: F401
 
+# Resolve once so the rest of the module can use a plain int; also fails loudly
+# at import time rather than at first call if the attribute is ever missing.
+_ZIP_ZSTANDARD: int = zipfile.ZIP_ZSTANDARD  # type: ignore[attr-defined]
+
 zipfile_compress_kwargs: dict[str, Any] = {
-    "compression": zipfile.ZIP_ZSTANDARD,  # type: ignore[attr-defined]
+    "compression": _ZIP_ZSTANDARD,
     "compresslevel": None,
 }
 
+
+# 200 MiB. Well under fzstd's 256 MiB (2^28) compressed-frame overflow
+# threshold, applied to *input* bytes which compress smaller.
+_MAX_INPUT_PER_FRAME = 200 * 1024 * 1024
+
+# Matches zipfile_zstd's hardcoded default so multi-frame and single-frame zip
+# writes use the same thread count. If zipfile_zstd ever changes its default,
+# update here too.
+_ZSTD_THREADS = 12
+
+
+class _MultiFrameZstdCompressObj:
+    """A zstd compressobj that chunks its output into multiple frames.
+
+    Wraps a ``zstandard`` compressobj and flushes it (finalizing the current
+    frame) and replaces it (starting a new frame) every
+    ``_MAX_INPUT_PER_FRAME`` bytes of input. Multi-frame zstd streams are
+    valid per spec -- any compliant decoder reads them transparently.
+    """
+
+    def __init__(self, factory: Callable[[], zstandard.ZstdCompressionObj]) -> None:
+        self._factory = factory
+        self._obj = factory()
+        self._input_bytes = 0
+
+    def compress(self, data: bytes) -> bytes:
+        view = memoryview(data)
+        pieces: list[bytes] = []
+        offset = 0
+        n = len(view)
+        while offset < n:
+            remaining_cap = _MAX_INPUT_PER_FRAME - self._input_bytes
+            end = min(offset + remaining_cap, n)
+            chunk = view[offset:end]
+            pieces.append(self._obj.compress(chunk))
+            self._input_bytes += end - offset
+            offset = end
+            if self._input_bytes >= _MAX_INPUT_PER_FRAME:
+                pieces.append(self._obj.flush())
+                self._obj = self._factory()
+                self._input_bytes = 0
+        return b"".join(pieces)
+
+    def flush(self) -> bytes:
+        # If the last ``compress()`` call landed exactly on the frame boundary,
+        # ``self._obj`` was replaced with a fresh compressobj that has received
+        # no bytes. Flushing it would append an empty 9-byte trailing frame.
+        if self._input_bytes == 0:
+            return b""
+        return self._obj.flush()
+
+
+class _MultiFrameZstdDecompressObj:
+    """A zstd decompressobj that transparently spans multiple frames.
+
+    ``zstandard.ZstdDecompressor().decompressobj()`` stops at the first frame
+    boundary and marks ``eof=True``.  When the compressor splits large entries
+    into multiple frames (see ``_MultiFrameZstdCompressObj``), the reader must
+    recognise the frame boundary, start a fresh inner decompressobj, and
+    continue until all compressed bytes have been consumed.
+
+    The stdlib ``zipfile._read1`` path for non-deflate compression does:
+
+        data = self._decompressor.decompress(data)
+        self._eof = self._decompressor.eof or self._compress_left <= 0
+
+    We return ``eof=False`` always so that ``self._eof`` is driven purely by
+    ``self._compress_left <= 0`` (all compressed bytes fed).  Meanwhile,
+    ``decompress`` buffers leftover bytes from a completed frame and feeds them
+    into the next inner decompressobj.
+    """
+
+    def __init__(self) -> None:
+        import zstandard as zstd  # local import -- already a hard dep
+
+        self._dctx: zstandard.ZstdDecompressor = zstd.ZstdDecompressor()
+        self._obj: zstandard.ZstdDecompressionObj = self._dctx.decompressobj()
+        self._pending: bytes = b""
+
+    def decompress(self, data: bytes) -> bytes:
+        self._pending += data
+        out = b""
+        while self._pending:
+            result = self._obj.decompress(self._pending)
+            out += result
+            if self._obj.eof:
+                # Inner frame complete; unused_data holds the start of the next.
+                self._pending = self._obj.unused_data
+                self._obj = self._dctx.decompressobj()
+            else:
+                # All pending input consumed; wait for more.
+                self._pending = b""
+                break
+        return out
+
+    def flush(self) -> bytes:
+        return b""
+
+    @property
+    def eof(self) -> bool:
+        # Always False: let compress_left drive the outer EOF check.
+        return False
+
+
+def _install_multiframe_patches() -> None:
+    """Install multi-frame zstd compressor and decompressor patches.
+
+    Idempotent. Wraps whatever ``_get_compressor`` / ``_get_decompressor`` were
+    installed before us (stdlib on Py >= 3.14; ``zipfile_zstd``'s versions on
+    Py < 3.14), so compression level and thread count are preserved.
+    """
+    if getattr(zipfile, "_inspect_ai_multiframe_installed", False):
+        return
+
+    original_compressor = zipfile._get_compressor  # type: ignore[attr-defined]
+    original_decompressor = zipfile._get_decompressor  # type: ignore[attr-defined]
+
+    def patched_compressor(compress_type: int, compresslevel: int | None = None) -> Any:
+        if compress_type == _ZIP_ZSTANDARD:
+            # Share one ``ZstdCompressor`` across all frames of the entry.
+            # Delegating to ``original_compressor`` would instead create a
+            # fresh ``ZstdCompressor(threads=N)`` per frame, re-initialising
+            # its thread pool every 200 MiB.
+            import zstandard
+
+            level = 3 if compresslevel is None else compresslevel
+            compressor = zstandard.ZstdCompressor(level=level, threads=_ZSTD_THREADS)
+            return _MultiFrameZstdCompressObj(compressor.compressobj)
+        return original_compressor(compress_type, compresslevel)
+
+    def patched_decompressor(compress_type: int) -> Any:
+        if compress_type == _ZIP_ZSTANDARD:
+            return _MultiFrameZstdDecompressObj()
+        return original_decompressor(compress_type)
+
+    zipfile._get_compressor = patched_compressor  # type: ignore[attr-defined]
+    zipfile._get_decompressor = patched_decompressor  # type: ignore[attr-defined]
+    zipfile._inspect_ai_multiframe_installed = True  # type: ignore[attr-defined]
+
+
+_install_multiframe_patches()
+
+
 __all__ = ["zipfile_compress_kwargs"]

tests/util/test_zipfile_multiframe.py

@@ -0,0 +1,135 @@
+"""Tests for multi-frame zstd compression of ZIP entries.
+
+Producer-side fix: every ZIP_ZSTANDARD entry should be a multi-frame zstd
+stream, capped at 200 MiB of input per frame, so JS decoders with per-frame
+size limits (fzstd @ 256 MiB) can decode large inspect_ai .eval files.
+"""
+
+from __future__ import annotations
+
+import zipfile
+from pathlib import Path
+
+import pytest
+import zstandard
+
+# Importing this module installs the zstd compression patches (both the
+# zipfile_zstd delegation on Python < 3.14 and our multi-frame wrapper).
+import inspect_ai._util.zipfile  # noqa: F401
+
+MAX_INPUT_PER_FRAME = 200 * 1024 * 1024  # must match the value in _util/zipfile.py
+ZSTD_MAGIC = b"\x28\xb5\x2f\xfd"
+
+# zipfile_zstd adds ZIP_ZSTANDARD (93) to zipfile at import time; use getattr
+# so type checkers on Python < 3.14 (where stdlib zipfile doesn't declare it)
+# stay happy.
+ZIP_ZSTANDARD: int = getattr(zipfile, "ZIP_ZSTANDARD", 93)
+
+
+def _moderately_compressible_payload(total_bytes: int) -> bytes:
+    """Build a deterministic payload that compresses at a realistic ratio.
+
+    Pure random data bypasses zstd's match-finder (compressor emits raw
+    blocks); pure-constant data collapses to tiny output. We want something
+    in between so the compressed output is a non-trivial fraction of input —
+    close to the JSON-like eval workload we care about.
+    """
+    fragments = [
+        b'{"role": "assistant", "content": "The quick brown fox jumps over the lazy dog."}\n',
+        b'{"role": "user", "content": "What is the capital of France? Please answer in one sentence."}\n',
+        b'{"tool_calls": [{"name": "search", "arguments": {"query": "weather in paris"}}]}\n',
+        b'{"metadata": {"timestamp": "2026-04-23T10:00:00Z", "model": "claude-opus-4-7"}}\n',
+    ]
+    chunk = b"".join(fragments)
+    n = (total_bytes // len(chunk)) + 1
+    return (chunk * n)[:total_bytes]
+
+
+@pytest.fixture(scope="module")
+def large_payload() -> bytes:
+    """~450 MiB of moderately compressible JSON-like bytes."""
+    return _moderately_compressible_payload(450 * 1024 * 1024)
+
+
+def _read_raw_compressed_entry(zip_path: Path, entry_name: str) -> bytes:
+    """Return the raw compressed bytes stored for a zip entry (no decompression)."""
+    with zipfile.ZipFile(zip_path) as zf:
+        info = zf.getinfo(entry_name)
+        with open(zip_path, "rb") as f:
+            f.seek(info.header_offset)
+            # Skip local file header to reach the data. Using ZipFile's private
+            # helper is fragile; instead, decode the local header manually.
+            import struct
+
+            sig_version_etc = f.read(30)
+            assert sig_version_etc[:4] == b"PK\x03\x04", "not a local file header"
+            name_len = struct.unpack("<H", sig_version_etc[26:28])[0]
+            extra_len = struct.unpack("<H", sig_version_etc[28:30])[0]
+            f.read(name_len + extra_len)
+            return f.read(info.compress_size)
+
+
+def _iter_frames(compressed: bytes):
+    """Yield (frame_index, decompressed_bytes) for each frame in a zstd stream."""
+    dctx = zstandard.ZstdDecompressor()
+    idx = 0
+    remaining = compressed
+    while remaining:
+        obj = dctx.decompressobj()
+        out = obj.decompress(remaining)
+        # decompressobj consumes one frame; remaining bytes are in unused_data
+        yield idx, out
+        idx += 1
+        remaining = obj.unused_data
+
+
+def test_large_entry_emits_multiple_capped_frames(
+    tmp_path: Path, large_payload: bytes
+) -> None:
+    """A ~450 MiB payload should produce ≥3 zstd frames, each ≤ 200 MiB uncompressed."""
+    zip_path = tmp_path / "large.zip"
+    with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf:
+        zf.writestr("big.json", large_payload)
+
+    raw = _read_raw_compressed_entry(zip_path, "big.json")
+
+    magic_count = raw.count(ZSTD_MAGIC)
+    assert magic_count >= 3, (
+        f"expected ≥3 zstd frames in a 450 MiB entry, found {magic_count}"
+    )
+
+    frame_sizes = [len(out) for _, out in _iter_frames(raw)]
+    assert len(frame_sizes) == magic_count, (
+        f"magic count {magic_count} != frames walked {len(frame_sizes)}"
+    )
+    for i, size in enumerate(frame_sizes):
+        assert size <= MAX_INPUT_PER_FRAME, (
+            f"frame {i} decompressed to {size} bytes, exceeds cap {MAX_INPUT_PER_FRAME}"
+        )
+
+
+def test_small_entry_emits_single_frame(tmp_path: Path) -> None:
+    """A 1 MiB entry should still be exactly one zstd frame."""
+    zip_path = tmp_path / "small.zip"
+    payload = _moderately_compressible_payload(1 * 1024 * 1024)
+    with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf:
+        zf.writestr("small.json", payload)
+
+    raw = _read_raw_compressed_entry(zip_path, "small.json")
+    assert raw.count(ZSTD_MAGIC) == 1, (
+        f"expected exactly 1 frame for small entry, got {raw.count(ZSTD_MAGIC)}"
+    )
+
+
+def test_large_entry_round_trip(tmp_path: Path, large_payload: bytes) -> None:
+    """Writing and reading back a multi-frame entry must preserve bytes exactly."""
+    zip_path = tmp_path / "rt.zip"
+    with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf:
+        zf.writestr("big.json", large_payload)
+
+    with zipfile.ZipFile(zip_path) as zf:
+        got = zf.read("big.json")
+
+    assert got == large_payload, (
+        f"round-trip mismatch: input {len(large_payload)} bytes, got {len(got)} bytes"
+    )