feat(xtest): sparse sink that keeps only planned TDF extents

xtest/rangeserve.py

@@ -198,6 +198,101 @@ def read(self, start: int, end: int) -> bytes:
         return b"".join(os.pread(fd, n, p) for p, n in self._pieces(start, end))
 
 
+class ManifestTooLarge(Exception):
+    """The TDF's tail did not fit in the sink's rolling buffer."""
+
+
+class _TailView:
+    """The rolling tail buffer, addressed by offsets into the whole stream."""
+
+    def __init__(self, buf: bytearray, base: int):
+        self._buf, self._base, self._pos = buf, base, base
+
+    def seek(self, offset: int, /) -> int:
+        if offset < self._base:
+            raise ManifestTooLarge(
+                f"the TDF tail reaches back to {offset}, but only bytes from "
+                f"{self._base} were buffered; raise tail_ceiling"
+            )
+        self._pos = offset
+        return offset
+
+    def read(self, n: int, /) -> bytes:
+        start = self._pos - self._base
+        self._pos += n
+        return bytes(self._buf[start : start + n])
+
+
+class SparseSink:
+    """Keep only the planned extents and the tail of a TDF as it streams past.
+
+    ``write`` the TDF in order, in chunks of any size, then ``close`` to get
+    the :class:`ExtentStore`. The tail -- everything after the payload data:
+    manifest entry, central directory, end records -- is not known until the
+    stream ends, so the last ``tail_ceiling`` bytes are held in memory. A
+    tail longer than that raises :class:`ManifestTooLarge` rather than
+    silently keeping part of it; at 50 TiB with 2 MiB segments the manifest
+    alone is about 2.6 GB, so size the ceiling to the run.
+    """
+
+    def __init__(
+        self, directory: Path, plan: Sequence[tuple[int, int]], *, tail_ceiling: int
+    ):
+        # The EOCD search reads up to this much from the end, whatever the ceiling.
+        self._ceiling = max(tail_ceiling, zipinspect.MAX_EOCD_SEARCH)
+        self._plan = merge(plan)
+        self._next = 0  # first plan extent not yet fully written
+        self._pos = 0
+        self._tail = bytearray()
+        # Logical length is unknown until close; ExtentStore checks it on add.
+        self._store = ExtentStore(directory, 2**63)
+
+    def write(self, data: bytes) -> None:
+        start, end = self._pos, self._pos + len(data)
+        i = self._next
+        while i < len(self._plan) and self._plan[i][0] < end:
+            lo, hi = max(self._plan[i][0], start), min(self._plan[i][1], end)
+            if lo < hi:
+                self._store.add(lo, data[lo - start : hi - start])
+            if self._plan[i][1] > end:
+                break
+            i += 1
+        self._next = i
+        self._pos = end
+        self._tail += data
+        del self._tail[: -self._ceiling]
+
+    def close(self) -> ExtentStore:
+        total = self._pos
+        if self._plan and self._plan[-1][1] > total:
+            raise ValueError(
+                f"stream ended at {total}, before planned extent {self._plan[-1]}"
+            )
+        self._store.length = total
+        base = total - len(self._tail)
+        cd = zipinspect.read_central_directory(_TailView(self._tail, base), total)
+        payload = next((e for e in cd.entries if e.name == "0.payload"), None)
+        if payload is None:
+            raise zipinspect.MalformedZipError("no 0.payload entry")
+        # The payload's local header is in the head, which every plan keeps.
+        header = self._store.read(
+            payload.local_header_offset,
+            payload.local_header_offset + _LOCAL_HEADER_FIXED_SIZE,
+        )
+        name_len, extra_len = struct.unpack_from("<HH", header, 26)
+        tail_start = (
+            payload.local_header_offset
+            + _LOCAL_HEADER_FIXED_SIZE
+            + name_len
+            + extra_len
+            + payload.compressed_size
+        )
+        _TailView(self._tail, base).seek(tail_start)
+        self._store.add(tail_start, bytes(self._tail[tail_start - base :]))
+        self._tail = bytearray()
+        return self._store
+
+
 class Request(NamedTuple):
     method: str
     range: str | None

xtest/test_rangeserve_units.py

@@ -206,3 +206,58 @@ def test_fixture_serves_several(range_server, tdf: Path):
     b, _ = range_server(FullFile(tdf))
     assert a != b
     assert get(a, "bytes=0-3") == get(b, "bytes=0-3")
+
+
+# --- SparseSink ---------------------------------------------------------------
+
+
+def make_tdf(path: Path, manifest_len: int, *, force_zip64: bool = False) -> Path:
+    with zipfile.ZipFile(path, "w", zipfile.ZIP_STORED) as z:
+        with z.open("0.payload", "w", force_zip64=force_zip64) as f:
+            f.write(random.Random(5111).randbytes(PT_LEN + 5 * GCM))
+        z.writestr("0.manifest.json", "{}" + " " * manifest_len)
+    return path
+
+
+def stream(sink: rangeserve.SparseSink, data: bytes, chunk: int) -> None:
+    for i in range(0, len(data), chunk):
+        sink.write(data[i : i + chunk])
+
+
+@pytest.mark.parametrize("chunk", [1, 7, 4096, 1 << 20])
+@pytest.mark.parametrize("force_zip64", [False, True], ids=["zip32", "zip64"])
+def test_sink_keeps_plan_and_tail(tmp_path: Path, chunk: int, force_zip64: bool):
+    tdf = make_tdf(tmp_path / "t.tdf", 3000, force_zip64=force_zip64)
+    data = tdf.read_bytes()
+    start = rangeserve.payload_start(tdf)
+    kept = plan((SEG - 1, SEG + 1), (3 * SEG + 1, 3 * SEG + 2), start=start)
+    (tmp_path / "store").mkdir()
+    sink = rangeserve.SparseSink(tmp_path / "store", kept, tail_ceiling=1 << 16)
+    stream(sink, data, chunk)
+    store = sink.close()
+
+    payload_end = start + PT_LEN + 5 * GCM
+    assert store.length == len(data)
+    assert store.read(payload_end, len(data)) == data[payload_end:]
+    for s, e in kept:
+        assert store.read(s, e) == data[s:e]
+    assert store.physical_size == sum(e - s for s, e in kept) + len(data) - payload_end
+    with serve(store) as (url, _):
+        assert get(url, "bytes=-22") == (206, data[-22:])
+        assert get(url, f"bytes={kept[0][1]}-{kept[1][0] - 1}")[0] == 410
+    store.close()
+
+
+def test_sink_refuses_tail_over_ceiling(tmp_path: Path):
+    data = make_tdf(tmp_path / "t.tdf", 200_000).read_bytes()
+    sink = rangeserve.SparseSink(tmp_path, [(0, 100)], tail_ceiling=1 << 16)
+    stream(sink, data, 65536)
+    with pytest.raises(rangeserve.ManifestTooLarge):
+        sink.close()
+
+
+def test_sink_refuses_short_stream(tmp_path: Path):
+    sink = rangeserve.SparseSink(tmp_path, [(0, 100), (500, 600)], tail_ceiling=0)
+    sink.write(b"x" * 550)
+    with pytest.raises(ValueError, match="before planned extent"):
+        sink.close()

xtest/zipinspect.py

@@ -19,7 +19,7 @@
 import struct
 from dataclasses import dataclass
 from pathlib import Path
-from typing import NamedTuple
+from typing import NamedTuple, Protocol
 
 from sizes import ZIP64_WINDOW_HIGH, in_zip64_window
 
@@ -43,7 +43,7 @@
 _CEN_FIXED_SIZE = 46
 #: A ZIP comment is a 16-bit length, so the EOCD cannot start further back
 #: than this from the end of the file.
-_MAX_EOCD_SEARCH = _EOCD_SIZE + 0xFFFF
+MAX_EOCD_SEARCH = _EOCD_SIZE + 0xFFFF
 
 
 class MalformedZipError(Exception):
@@ -221,55 +221,68 @@ def central_directory(path: Path) -> CentralDirectory:
     :class:`MalformedZipError` naming the problem, not as a bare
     ``struct.error`` from somewhere in the middle of the parse.
     """
-    file_size = path.stat().st_size
     with path.open("rb") as f:
-        tail_len = min(file_size, _MAX_EOCD_SEARCH)
-        f.seek(file_size - tail_len)
-        tail = f.read(tail_len)
+        return read_central_directory(f, path.stat().st_size)
 
-        eocd_at = _find_eocd(tail)
-        if eocd_at + _EOCD_SIZE > len(tail):
+
+class Seekable(Protocol):
+    def seek(self, offset: int, /) -> object: ...
+    def read(self, n: int, /) -> bytes: ...
+
+
+def read_central_directory(f: Seekable, file_size: int) -> CentralDirectory:
+    """:func:`central_directory` over any seekable source of ``file_size`` bytes.
+
+    Only the tail and the directory are read, which is what lets a streaming
+    writer's rolling tail buffer stand in for the whole file.
+    """
+    tail_len = min(file_size, MAX_EOCD_SEARCH)
+    f.seek(file_size - tail_len)
+    tail = f.read(tail_len)
+
+    eocd_at = _find_eocd(tail)
+    if eocd_at + _EOCD_SIZE > len(tail):
+        raise MalformedZipError(
+            f"end-of-central-directory record at {eocd_at} is truncated: "
+            f"{len(tail) - eocd_at} of {_EOCD_SIZE} bytes"
+        )
+    (
+        cd_entries_this_disk,
+        cd_entries_total,
+        cd_size,
+        cd_offset,
+    ) = struct.unpack_from("<HHII", tail, eocd_at + 8)
+    del cd_entries_this_disk
+
+    entry_count = cd_entries_total
+    # ZIP64: the 32-bit EOCD holds sentinels and the real values live in
+    # the ZIP64 EOCD record, found via the locator that precedes the EOCD.
+    locator_at = eocd_at - _EOCD64_LOCATOR_SIZE
+    if locator_at >= 0 and tail[locator_at : locator_at + 4] == _EOCD64_LOCATOR_SIG:
+        (eocd64_offset,) = struct.unpack_from("<Q", tail, locator_at + 8)
+        if eocd64_offset + _EOCD64_SIZE > file_size:
             raise MalformedZipError(
-                f"end-of-central-directory record at {eocd_at} is truncated: "
-                f"{len(tail) - eocd_at} of {_EOCD_SIZE} bytes"
+                f"zip64 locator points at {eocd64_offset}, past the end of "
+                f"a {file_size}-byte file"
             )
-        (
-            cd_entries_this_disk,
-            cd_entries_total,
-            cd_size,
-            cd_offset,
-        ) = struct.unpack_from("<HHII", tail, eocd_at + 8)
-        del cd_entries_this_disk
-
-        entry_count = cd_entries_total
-        # ZIP64: the 32-bit EOCD holds sentinels and the real values live in
-        # the ZIP64 EOCD record, found via the locator that precedes the EOCD.
-        locator_at = eocd_at - _EOCD64_LOCATOR_SIZE
-        if locator_at >= 0 and tail[locator_at : locator_at + 4] == _EOCD64_LOCATOR_SIG:
-            (eocd64_offset,) = struct.unpack_from("<Q", tail, locator_at + 8)
-            if eocd64_offset + _EOCD64_SIZE > file_size:
-                raise MalformedZipError(
-                    f"zip64 locator points at {eocd64_offset}, past the end of "
-                    f"a {file_size}-byte file"
-                )
-            f.seek(eocd64_offset)
-            eocd64 = f.read(_EOCD64_SIZE)
-            if eocd64[:4] != _EOCD64_SIG:
-                raise MalformedZipError(
-                    f"zip64 locator points at {eocd64_offset}, which is not a "
-                    "zip64 end-of-central-directory record"
-                )
-            entry_count, cd_size, cd_offset = struct.unpack_from("<QQQ", eocd64, 32)
-
-        # Before the seek, not after: cd_size comes straight off the wire and
-        # a garbage value here would otherwise drive an unbounded read.
-        if cd_offset + cd_size > file_size:
+        f.seek(eocd64_offset)
+        eocd64 = f.read(_EOCD64_SIZE)
+        if eocd64[:4] != _EOCD64_SIG:
             raise MalformedZipError(
-                f"central directory claims {cd_size} bytes at offset "
-                f"{cd_offset}, past the end of a {file_size}-byte file"
+                f"zip64 locator points at {eocd64_offset}, which is not a "
+                "zip64 end-of-central-directory record"
             )
-        f.seek(cd_offset)
-        cd = f.read(cd_size)
+        entry_count, cd_size, cd_offset = struct.unpack_from("<QQQ", eocd64, 32)
+
+    # Before the seek, not after: cd_size comes straight off the wire and
+    # a garbage value here would otherwise drive an unbounded read.
+    if cd_offset + cd_size > file_size:
+        raise MalformedZipError(
+            f"central directory claims {cd_size} bytes at offset "
+            f"{cd_offset}, past the end of a {file_size}-byte file"
+        )
+    f.seek(cd_offset)
+    cd = f.read(cd_size)
 
     entries: list[CentralDirectoryEntry] = []
     pos = 0