diff --git a/xtest/rangeserve.py b/xtest/rangeserve.py index 385c25b9..666e4287 100644 --- a/xtest/rangeserve.py +++ b/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(" 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() diff --git a/xtest/zipinspect.py b/xtest/zipinspect.py index de3faf31..379c26e1 100644 --- a/xtest/zipinspect.py +++ b/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("= 0 and tail[locator_at : locator_at + 4] == _EOCD64_LOCATOR_SIG: + (eocd64_offset,) = struct.unpack_from(" 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("= 0 and tail[locator_at : locator_at + 4] == _EOCD64_LOCATOR_SIG: - (eocd64_offset,) = struct.unpack_from(" 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(" 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(" 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