From f4aec0fbeaabf324156ff4f1544585f8b3ae4bee Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Thu, 6 Aug 2026 14:25:15 -0600 Subject: [PATCH 1/4] =?UTF-8?q?pgw#973=20=C2=A74.24=20wave=202:=20six=20sa?= =?UTF-8?q?fetensors=20header=20caps=20become=20one,=20plus=20three=20dupl?= =?UTF-8?q?icate=20bounds=20and=20two=20false=20justifications?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ONE THREAT, ONE NUMBER (the headline). A safetensors file opens with an 8-byte declared header length read straight from the file, and every reader turns it into json.loads(f.read(n)) before anything has validated it. That threat was stated six times: w4a4, w8a8, svdq, loading and convert/ingest at 100 MiB, convert/writer at 512 MiB. The outlier was a live defect, not just noise: writer accepted headers loading refuses, so the re-shard path could emit a shard the serving path could not open. Same bytes, two verdicts. All six now read models/safetensors_header.py, where the threat, why nothing else prevents it (the length is read before any other structure exists), and why 100 MiB is a plausibility floor rather than a measurement are stated once — with what would change it. DELETED, dead: models/memory.py _DEFAULT_VAE_SLICE_THRESHOLD_GB = 10.0, zero references in src, tests, docs or scripts. SINGLE OWNER, duplicated value: _READ_CHUNK_BYTES (chunk_cas already exports the chunk vocabulary chunk_upload imports). CENSUS CORRECTED — input_assets vs url_fetch is NOT two caps on one path. open_guarded_stream deliberately caps nothing ("the caller owns the read and its byte cap"); fetch_bytes and _download_one are separate entry points each enforcing its own. Acting on the filed verdict literally would have left one path with no cap. Aliased to one owner instead of deleted. CENSUS CORRECTED — _PRESIGNED_PUT_BUDGET is not the authoritative cap and is not deletable either. Its justification cited _concurrent_upload.py, a module that NO LONGER EXISTS, and claimed file-level fan-out of 4; the only in-repo caller is sequential, so the file axis is 1 and min(total_parts, 4) is what binds. It is KEPT because it covers the one axis optimal_part_concurrency cannot see — an endpoint author calling ctx.save() from their own threads — and both docstrings now say that instead of citing each other. Proof: tests/test_one_safetensors_header_bound_pgw973.py — real files on disk through every real entry point, no mocks and no monkeypatching of the bound. The runaway (2**63-1), the over-cap and the absent (0) length are refused at all six; a legitimate file still parses at all six; writer and loader now agree on a 200 MiB-declared header. A structural pin fails the suite if any module re-declares a header cap. RED-verified: restoring 512 MiB in convert/writer.py fails test_writer_and_loader_agree_on_the_same_file — the writer accepts the length and only trips later on the short body, i.e. it would have attempted the 200 MiB read. Gates: tests/ 3297 passed, 37 skipped, 1 xfailed (the single -n 4 failure, test_procsplit_pgw763, passes in isolation — load-induced, untouched by this change); test_import_cycles_pgw981 green (the new leaf module adds no cycle). Python 3.12. --- changelog.d/pgw973-wave2.md | 36 ++++ src/gen_worker/_upload_transport.py | 13 +- src/gen_worker/convert/ingest.py | 4 +- src/gen_worker/convert/writer.py | 4 +- src/gen_worker/input_assets.py | 9 +- src/gen_worker/models/chunk_upload.py | 8 +- src/gen_worker/models/loading.py | 8 +- src/gen_worker/models/memory.py | 1 - src/gen_worker/models/safetensors_header.py | 45 +++++ src/gen_worker/models/svdq.py | 4 +- src/gen_worker/models/w4a4.py | 4 +- src/gen_worker/models/w8a8.py | 4 +- src/gen_worker/presigned_upload.py | 25 ++- ...est_one_safetensors_header_bound_pgw973.py | 155 ++++++++++++++++++ 14 files changed, 292 insertions(+), 28 deletions(-) create mode 100644 changelog.d/pgw973-wave2.md create mode 100644 src/gen_worker/models/safetensors_header.py create mode 100644 tests/test_one_safetensors_header_bound_pgw973.py diff --git a/changelog.d/pgw973-wave2.md b/changelog.d/pgw973-wave2.md new file mode 100644 index 00000000..2f7d1b99 --- /dev/null +++ b/changelog.d/pgw973-wave2.md @@ -0,0 +1,36 @@ +- **pgw#973 (§4.24 execution wave 2): the safetensors header cap is now one + number, not six.** A safetensors file opens with an 8-byte declared header + length taken straight from the file, and every reader turned it directly into + `json.loads(f.read(n))`. That threat was bounded six times — `models/w4a4.py`, + `models/w8a8.py`, `models/svdq.py`, `models/loading.py` and + `convert/ingest.py` at 100 MiB, and `convert/writer.py` at **512 MiB**. The + outlier was not harmless: the writer accepted headers the loader would refuse, + so the re-shard path could emit a shard the serving path could not open — same + bytes, two verdicts. All six now read + `gen_worker.models.safetensors_header.MAX_HEADER_BYTES` / `header_len_ok()`, + where the threat, the reason nothing else prevents it, and why 100 MiB is a + plausibility floor rather than a measurement are stated once. RED-verified: + restoring the 512 MiB cap makes the writer parse a 200 MiB-declared header the + loader rejects. +- **pgw#973: a dead offload threshold deleted.** + `models/memory.py:_DEFAULT_VAE_SLICE_THRESHOLD_GB = 10.0` had zero references + anywhere in the repo — src, tests, docs and scripts. +- **pgw#973: two verbatim-duplicated bounds given single owners.** + `_READ_CHUNK_BYTES` (4 MiB) was defined identically in `models/chunk_cas.py` + and `models/chunk_upload.py`, which already imports the rest of the chunk + vocabulary from `chunk_cas`; and `input_assets._DEFAULT_MAX_BYTES` / + `_CHUNK` restated `url_fetch.DEFAULT_MAX_BYTES` / `_CHUNK`. **Correcting the + census on the second one:** these are *not* two caps on one fetch path. + `url_fetch.open_guarded_stream` deliberately caps nothing ("the caller owns + the read and its byte cap"), and `url_fetch.fetch_bytes` and + `input_assets._download_one` are separate entry points that each enforce their + own. Deleting either would leave one path with no cap at all, so the value is + aliased to a single owner rather than removed. +- **pgw#973: two limit justifications that cited a deleted module are corrected.** + `presigned_upload._PRESIGNED_PUT_BUDGET = 8` called itself "the authoritative + cap" on the grounds that "file-level fan-out is fixed at 4" — but the module + that owned file-level fan-out (`_concurrent_upload.py`) no longer exists and + the only in-repo caller is sequential, so the file axis is 1 and the binding + bound is `optimal_part_concurrency`'s `min(total_parts, 4)`. The semaphore is + KEPT because it covers the one axis nothing else can see — an endpoint author + calling `ctx.save()` from their own threads — and both docstrings now say so. diff --git a/src/gen_worker/_upload_transport.py b/src/gen_worker/_upload_transport.py index 9a86f5d7..b5faf067 100644 --- a/src/gen_worker/_upload_transport.py +++ b/src/gen_worker/_upload_transport.py @@ -402,9 +402,16 @@ def optimal_part_concurrency(total_parts: int) -> int: """Fixed part-level concurrency for one file's multipart upload. A single file can saturate R2 with a small number of in-flight PUTs. - Keep this fixed so it cannot multiply with file-level fan-out into an - uncontrolled retry storm. The current Tensorhub presigned path also - has a process-wide PUT budget in ``presigned_upload.py``. + + pgw#973 (§4.24): this is the BINDING bound on the in-repo presigned path — + the caller is sequential, so 4 is the real ceiling on concurrent PUTs. + ``presigned_upload._PRESIGNED_PUT_BUDGET`` (8) is not a second cap on this + axis; it covers a different one (an endpoint author saving from their own + threads). The previous text here cited ``_concurrent_upload.py`` as the + file-level fan-out owner — that module no longer exists. + + Without this, one large file's part count IS the concurrency, and a + thousand-part upload opens a thousand PUTs. """ if total_parts <= 1: return 1 diff --git a/src/gen_worker/convert/ingest.py b/src/gen_worker/convert/ingest.py index a1884407..23e4d6c7 100644 --- a/src/gen_worker/convert/ingest.py +++ b/src/gen_worker/convert/ingest.py @@ -30,6 +30,7 @@ ) from ..net import hf, install_hf_http_timeouts +from ..models.safetensors_header import header_len_ok from .classifier import RepoClassification, apply_source_include, classify_repo from .layout import detect_huggingface_source_layout from huggingface_hub.errors import EntryNotFoundError, GatedRepoError, RepositoryNotFoundError, RevisionNotFoundError @@ -78,7 +79,6 @@ def _download_attempts() -> int: "F32": "fp32", "F16": "fp16", "BF16": "bf16", "F8_E4M3": "fp8", "F8_E5M2": "fp8:e5m2", } -_MAX_SAFETENSORS_HEADER_BYTES = 100 * 1024 * 1024 def _detect_snapshot_dtype(root: Path) -> str: @@ -95,7 +95,7 @@ def _detect_snapshot_dtype(root: Path) -> str: if len(raw) < 8: continue (n,) = struct.unpack(" _MAX_SAFETENSORS_HEADER_BYTES: + if not header_len_ok(n): continue header = json.loads(f.read(n)) for value in header.values(): diff --git a/src/gen_worker/convert/writer.py b/src/gen_worker/convert/writer.py index a47acad7..0cd82bc4 100644 --- a/src/gen_worker/convert/writer.py +++ b/src/gen_worker/convert/writer.py @@ -34,6 +34,7 @@ from fnmatch import fnmatch import random from gen_worker.models.w8a8 import detect_w8a8_artifact +from gen_worker.models.safetensors_header import header_len_ok if TYPE_CHECKING: import torch @@ -1391,7 +1392,6 @@ def _tensor_map(files: list[Path]) -> dict[str, Path]: # --------------------------------------------------------------------------- _HEADER_LEN_PREFIX = 8 -_MAX_HEADER_BYTES = 512 * 1024 * 1024 _RAW_COPY_CHUNK = 8 * 1024 * 1024 @@ -1402,7 +1402,7 @@ def _read_safetensors_header(fd: int) -> tuple[dict, int]: if len(prefix) != _HEADER_LEN_PREFIX: raise ValueError("safetensors: short read on header length prefix") header_len = int.from_bytes(prefix, "little") - if header_len <= 0 or header_len > _MAX_HEADER_BYTES: + if not header_len_ok(header_len): raise ValueError(f"safetensors: implausible header_length={header_len}") body = os.read(fd, header_len) if len(body) != header_len: diff --git a/src/gen_worker/input_assets.py b/src/gen_worker/input_assets.py index ee8d266a..9f25d21c 100644 --- a/src/gen_worker/input_assets.py +++ b/src/gen_worker/input_assets.py @@ -33,16 +33,19 @@ from .api.errors import CanceledError, RetryableError, ValidationError from .api.types import Asset, AudioAsset, ImageAsset, VideoAsset from .request_context._helpers import _infer_mime_type, _url_is_blocked -from .url_fetch import open_guarded_stream +from .url_fetch import DEFAULT_MAX_BYTES, _CHUNK, open_guarded_stream logger = logging.getLogger(__name__) -_DEFAULT_MAX_BYTES = 50 << 20 # matches tensorhub's default media cap +# pgw#973 (§4.24): one number for tensorhub's media cap, owned by url_fetch. +# open_guarded_stream deliberately caps nothing ("the caller owns the read +# and its byte cap"), so this path really does need its own enforcement — +# it just must not re-decide the value. +_DEFAULT_MAX_BYTES = DEFAULT_MAX_BYTES _DOWNLOAD_TIMEOUT_S = 120 _RESOLVE_TIMEOUT_S = 30 _MAX_RESOLVE_BODY = 8 << 20 _MAX_WALK_DEPTH = 32 -_CHUNK = 1 << 20 _INPUT_DIR_PREFIX = "gen-worker-inputs-" _RESOLVE_PATH = "/api/v1/worker/input-assets/resolve" _MANIFEST_KINDS = ("media", "image", "video", "audio") diff --git a/src/gen_worker/models/chunk_upload.py b/src/gen_worker/models/chunk_upload.py index d4227b23..cc8c8087 100644 --- a/src/gen_worker/models/chunk_upload.py +++ b/src/gen_worker/models/chunk_upload.py @@ -41,7 +41,12 @@ import requests -from .chunk_cas import CAS_CHUNK_SIZE_BYTES, chunk_count_for, chunk_len_at +from .chunk_cas import ( + CAS_CHUNK_SIZE_BYTES, + _READ_CHUNK_BYTES, + chunk_count_for, + chunk_len_at, +) __all__ = [ "ChunkPlan", @@ -54,7 +59,6 @@ _log = logging.getLogger(__name__) -_READ_CHUNK_BYTES = 4 * 1024 * 1024 _MAX_ATTEMPTS = 5 # Shared with the legacy presigned path's intent: total concurrent PUTs across diff --git a/src/gen_worker/models/loading.py b/src/gen_worker/models/loading.py index 07b6da93..2f1c3d29 100644 --- a/src/gen_worker/models/loading.py +++ b/src/gen_worker/models/loading.py @@ -28,6 +28,7 @@ from .artifact_contract import CONTRACT_PLAIN_BF16, implements_contract from .fp8_storage import restructure_fp8_storage from .memory import get_available_vram_gb, meta_tensors +from .safetensors_header import MAX_HEADER_BYTES, header_len_ok from .svdq import detect_svdq_artifact, load_svdq_pipeline from .w4a4 import ( detect_w4a4_artifact, @@ -94,7 +95,6 @@ def detect_diffusers_variant(model_path: Path) -> Optional[str]: _SAFETENSORS_DTYPE_NAMES = { "BF16": "bf16", "F16": "fp16", "F32": "fp32", "F8_E4M3": "fp8", } -_MAX_SAFETENSORS_HEADER_BYTES = 100 << 20 def safetensors_file_valid(path: Path) -> bool: @@ -111,7 +111,7 @@ def safetensors_file_valid(path: Path) -> bool: if len(raw) < 8: return False (n,) = struct.unpack(" _MAX_SAFETENSORS_HEADER_BYTES or 8 + n > size: + if not header_len_ok(n) or 8 + n > size: return False header = json.loads(f.read(n)) if not isinstance(header, dict): @@ -144,7 +144,7 @@ def detect_on_disk_dtype(model_path: Path) -> str: if len(raw) < 8: continue (n,) = struct.unpack(" _MAX_SAFETENSORS_HEADER_BYTES: + if not header_len_ok(n): continue header = json.loads(f.read(n)) for value in header.values(): @@ -1342,7 +1342,7 @@ def _safetensors_data_bytes(p: Path) -> int: if len(raw) < 8: return 0 (n,) = struct.unpack(" _MAX_SAFETENSORS_HEADER_BYTES: + if not header_len_ok(n): return 0 header = json.loads(f.read(n)) total = 0 diff --git a/src/gen_worker/models/memory.py b/src/gen_worker/models/memory.py index 59194ec6..139f7e11 100644 --- a/src/gen_worker/models/memory.py +++ b/src/gen_worker/models/memory.py @@ -41,7 +41,6 @@ "auto", "off", "vae_only", "model_offload", "group_offload", "sequential", ) -_DEFAULT_VAE_SLICE_THRESHOLD_GB = 10.0 _DEFAULT_MODEL_OFFLOAD_THRESHOLD_GB = 8.0 _DEFAULT_GROUP_OFFLOAD_THRESHOLD_GB = 6.0 # Safety margin below free VRAM reserved for activations. diff --git a/src/gen_worker/models/safetensors_header.py b/src/gen_worker/models/safetensors_header.py new file mode 100644 index 00000000..5ab3a207 --- /dev/null +++ b/src/gen_worker/models/safetensors_header.py @@ -0,0 +1,45 @@ +"""The one bound on a safetensors declared header length. + +pgw#973 (DESIGN-RULINGS §4.24). A safetensors file opens with an 8-byte +little-endian header length taken straight from the file. It is attacker- or +corruption-controlled, and every reader turns it directly into an allocation +(``json.loads(f.read(n))``). Unbounded, one crafted file declaring 2**63-1 +is an OOM in whichever process opened it — a serving worker, a conversion +pod, or the shard writer. + +THE THREAT, stated once: a declared header length that the file cannot back +sizes a read and a parse before anything has validated it. + +WHY NOTHING ELSE PREVENTS IT: the length is read before any other structure +exists, so there is nothing earlier to lean on. This bound is load-bearing. + +It was previously stated SIX times — ``models/w4a4.py``, ``models/w8a8.py``, +``models/svdq.py``, ``models/loading.py``, ``convert/ingest.py`` all at +100 MiB, and ``convert/writer.py`` at 512 MiB. The odd one out was not +harmless: writer accepted headers loading would refuse, so the re-shard path +could emit a shard the serving path could not open. Same bytes, two verdicts. + +WHY 100 MiB AND NOT A MEASUREMENT: real safetensors headers are tens of KB; +the largest sharded checkpoints in the fleet are a few MB of JSON. 100 MiB is +~20x above anything observed and exists only to make the number finite — it +is a plausibility floor, not a tuned capacity. What would change it: a +legitimate model whose header exceeds ~10 MiB, which would mean the tensor +count per shard grew by an order of magnitude. +""" + +from __future__ import annotations + +MAX_HEADER_BYTES: int = 100 << 20 + + +def header_len_ok(n: int) -> bool: + """Whether a declared safetensors header length is plausible. + + Zero and negative are refusals, not "no header": a file that declares + nothing is malformed, and treating it as an empty header would let a + truncated blob parse as a valid one (§4.24 item 4). + """ + return 0 < n <= MAX_HEADER_BYTES + + +__all__ = ["MAX_HEADER_BYTES", "header_len_ok"] diff --git a/src/gen_worker/models/svdq.py b/src/gen_worker/models/svdq.py index 72d16a33..4873c2f5 100644 --- a/src/gen_worker/models/svdq.py +++ b/src/gen_worker/models/svdq.py @@ -25,6 +25,7 @@ from dataclasses import dataclass from pathlib import Path from ..component_vocab import denoiser_components +from .safetensors_header import header_len_ok from typing import Any, Optional import importlib.metadata as md import os @@ -37,7 +38,6 @@ SVDQ_FP4_SMS = (120, 121) SVDQ_INT4_SMS = (75, 80, 86, 89) -_MAX_HEADER_BYTES = 100 << 20 class SvdqError(RuntimeError): @@ -291,7 +291,7 @@ def _read_safetensors_metadata(path: Path) -> dict: if len(raw) < 8: return {} (n,) = struct.unpack(" _MAX_HEADER_BYTES: + if not header_len_ok(n): return {} header = json.loads(f.read(n)) except (OSError, ValueError): diff --git a/src/gen_worker/models/w4a4.py b/src/gen_worker/models/w4a4.py index 56035348..a74dbe79 100644 --- a/src/gen_worker/models/w4a4.py +++ b/src/gen_worker/models/w4a4.py @@ -47,6 +47,7 @@ from pathlib import Path from .. import activity as activity_mod from ..component_vocab import denoiser_components +from .safetensors_header import header_len_ok from typing import Any, Dict, List, Optional import shutil @@ -72,7 +73,6 @@ # dims % 16 == 0 => in_features % 32 == 0, out_features % 16 == 0. _K_ALIGN = 32 _N_ALIGN = 16 -_MAX_HEADER_BYTES = 100 << 20 _AUX_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale", ".pre_quant_scale") @@ -102,7 +102,7 @@ def _read_header(path: Path) -> dict: if len(raw) < 8: return {} (n,) = struct.unpack(" _MAX_HEADER_BYTES: + if not header_len_ok(n): return {} header = json.loads(f.read(n)) except (OSError, ValueError): diff --git a/src/gen_worker/models/w8a8.py b/src/gen_worker/models/w8a8.py index 14993cfb..a22ef317 100644 --- a/src/gen_worker/models/w8a8.py +++ b/src/gen_worker/models/w8a8.py @@ -39,6 +39,7 @@ from pathlib import Path from .. import activity as activity_mod from ..component_vocab import denoiser_components +from .safetensors_header import header_len_ok from .artifact_contract import CONTRACT_COZY_FP8_ROWWISE, implements_contract from typing import Any, Dict, List, Optional import shutil @@ -53,7 +54,6 @@ W8A8_ROWWISE_MIN_SM = 90 _FP8_MAX = 448.0 _DIM_ALIGN = 16 -_MAX_HEADER_BYTES = 100 << 20 class W8a8Error(RuntimeError): @@ -82,7 +82,7 @@ def _read_header(path: Path) -> dict: if len(raw) < 8: return {} (n,) = struct.unpack(" _MAX_HEADER_BYTES: + if not header_len_ok(n): return {} header = json.loads(f.read(n)) except (OSError, ValueError): diff --git a/src/gen_worker/presigned_upload.py b/src/gen_worker/presigned_upload.py index 8c0879ce..db5e5a5d 100644 --- a/src/gen_worker/presigned_upload.py +++ b/src/gen_worker/presigned_upload.py @@ -104,11 +104,26 @@ # Default part size sent by server, but we read it from the response. _FALLBACK_PART_SIZE = 64 * 1024 * 1024 # 64 MiB -# Hard-coded internal safety budget for the current Tensorhub presigned -# upload path. File-level fan-out is fixed at 4 and per-file part fan-out is -# fixed at 4, so this semaphore is the authoritative cap that keeps the two -# axes from multiplying. Eight concurrent PUTs preserves useful parallelism -# while avoiding the 100+ PUT retry storm that broke R2 mirrors. +# pgw#973 (§4.24) — KEEP, but the old justification was false and is replaced. +# +# It read: "File-level fan-out is fixed at 4 and per-file part fan-out is fixed +# at 4, so this semaphore is the authoritative cap that keeps the two axes from +# multiplying." Both halves were wrong. The module that owned file-level +# fan-out (``_concurrent_upload.py``) NO LONGER EXISTS; the only in-repo caller +# of this path (``request_context/_stream.py:_finalize_presigned_upload``) is +# sequential, with no pool and no gather. So the file axis is 1, in-flight PUTs +# are capped at 4 by ``optimal_part_concurrency``, and this semaphore of 8 is +# never the binding constraint on any in-repo path. +# +# THE THREAT IT ACTUALLY COVERS, which nothing else does: an endpoint author +# calling ``ctx.save()`` from their own threads. That is the one axis +# ``optimal_part_concurrency`` cannot see, because it bounds ONE file's parts +# and knows nothing about how many files are in flight beside it. Without this, +# N author threads x 4 parts is unbounded and rebuilds the 100+ PUT retry storm +# that broke R2 mirrors. +# +# NOT DERIVED. 8 is round. What would change it: one measured author workload +# whose concurrent saves exceed 2 files. _PRESIGNED_PUT_BUDGET = 8 _presigned_put_slots = threading.BoundedSemaphore(_PRESIGNED_PUT_BUDGET) diff --git a/tests/test_one_safetensors_header_bound_pgw973.py b/tests/test_one_safetensors_header_bound_pgw973.py new file mode 100644 index 00000000..3ed8f361 --- /dev/null +++ b/tests/test_one_safetensors_header_bound_pgw973.py @@ -0,0 +1,155 @@ +"""pgw#973 (DESIGN-RULINGS §4.24): ONE bound on the safetensors header length. + +The threat is named in ``gen_worker/models/safetensors_header.py``: an 8-byte +declared header length, read from the file before anything has validated it, +sizes a read and a JSON parse in whichever process opened the file. + +It used to be stated six times — five readers at 100 MiB and +``convert/writer.py`` at 512 MiB. The outlier was not harmless: the writer +accepted headers the loader would refuse, so the re-shard path could emit a +shard the serving path could not open. Same bytes, two verdicts. + +These tests drive every real entry point against real files on disk — no +mocks, no monkeypatching of the bound. They prove the surviving bound still +refuses the runaway everywhere the deleted copies used to, and that writer and +loader now agree. + +RED-verified: restoring ``_MAX_HEADER_BYTES = 512 * 1024 * 1024`` in +convert/writer.py fails ``test_writer_and_loader_agree_on_the_same_file`` +(the writer parses a 200 MiB-declared header the loader rejects). +""" + +from __future__ import annotations + +import json +import os +import struct +from pathlib import Path + +import pytest + +from gen_worker.convert.ingest import _detect_snapshot_dtype +from gen_worker.convert.writer import _read_safetensors_header +from gen_worker.models.loading import safetensors_file_valid +from gen_worker.models.safetensors_header import MAX_HEADER_BYTES, header_len_ok +from gen_worker.models.svdq import _read_safetensors_metadata +from gen_worker.models.w4a4 import _read_header as w4a4_read_header +from gen_worker.models.w8a8 import _read_header as w8a8_read_header + +# One tensor of one F32 element, so the header describes something real. +_HEADER = {"t": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}} + + +def _write_safetensors(path: Path, *, declared_len: int | None = None) -> Path: + """Write a real safetensors file. ``declared_len`` overrides ONLY the + 8-byte prefix, which is exactly what a hostile or corrupt file does: it + promises a header it never has to back with bytes.""" + body = json.dumps(_HEADER).encode() + n = len(body) if declared_len is None else declared_len + path.write_bytes(struct.pack(" no dtype could be detected. + assert _detect_snapshot_dtype(snapshot) in (None, "") + + +def test_a_legitimate_file_still_parses_everywhere(tmp_path: Path): + """The deletion must not have made anything stricter than it was.""" + f = _write_safetensors(tmp_path / "ok.safetensors") + + assert safetensors_file_valid(str(f)) is True + assert w8a8_read_header(f) == _HEADER + assert w4a4_read_header(f) == _HEADER + fd = os.open(f, os.O_RDONLY) + try: + header, _ = _read_safetensors_header(fd) + finally: + os.close(fd) + assert header == _HEADER + + +def test_writer_and_loader_agree_on_the_same_file(tmp_path: Path): + """The defect the 512 MiB outlier caused, pinned. + + A 200 MiB declared header sits between the old writer cap (512 MiB) and + the reader cap (100 MiB). Before this change the writer parsed it and the + loader refused it — a shard the re-shard path would emit and the serving + path could not open. Now both refuse. + """ + between = 200 * 1024 * 1024 + assert MAX_HEADER_BYTES < between < 512 * 1024 * 1024 + f = _write_safetensors(tmp_path / "between.safetensors", declared_len=between) + + assert safetensors_file_valid(str(f)) is False, "premise: the loader refuses it" + fd = os.open(f, os.O_RDONLY) + try: + with pytest.raises(ValueError, match="implausible header_length"): + _read_safetensors_header(fd) + finally: + os.close(fd) From c3a3016e64e761958dee18dc5d9ffa0839820434 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Thu, 6 Aug 2026 14:34:57 -0600 Subject: [PATCH 2/4] =?UTF-8?q?pgw#973:=20keep=20input=5Fassets.=5FCHUNK?= =?UTF-8?q?=20local=20=E2=80=94=20it=20is=20a=20buffer,=20not=20a=20bound?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aliasing it to url_fetch's copy meant importing a private name across modules to de-duplicate something §4.24 does not reach: a 1 MiB streaming read buffer refuses nothing, so it is not a limit and sharing it buys no threat coherence. Only DEFAULT_MAX_BYTES — which does refuse — stays aliased to its single owner. Labelled in place so the next census does not re-file it. --- changelog.d/pgw973-wave2.md | 4 ++-- src/gen_worker/input_assets.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/changelog.d/pgw973-wave2.md b/changelog.d/pgw973-wave2.md index 2f7d1b99..f8b87a39 100644 --- a/changelog.d/pgw973-wave2.md +++ b/changelog.d/pgw973-wave2.md @@ -18,8 +18,8 @@ - **pgw#973: two verbatim-duplicated bounds given single owners.** `_READ_CHUNK_BYTES` (4 MiB) was defined identically in `models/chunk_cas.py` and `models/chunk_upload.py`, which already imports the rest of the chunk - vocabulary from `chunk_cas`; and `input_assets._DEFAULT_MAX_BYTES` / - `_CHUNK` restated `url_fetch.DEFAULT_MAX_BYTES` / `_CHUNK`. **Correcting the + vocabulary from `chunk_cas`; and `input_assets._DEFAULT_MAX_BYTES` + restated `url_fetch.DEFAULT_MAX_BYTES`. **Correcting the census on the second one:** these are *not* two caps on one fetch path. `url_fetch.open_guarded_stream` deliberately caps nothing ("the caller owns the read and its byte cap"), and `url_fetch.fetch_bytes` and diff --git a/src/gen_worker/input_assets.py b/src/gen_worker/input_assets.py index 9f25d21c..2c086ede 100644 --- a/src/gen_worker/input_assets.py +++ b/src/gen_worker/input_assets.py @@ -33,7 +33,7 @@ from .api.errors import CanceledError, RetryableError, ValidationError from .api.types import Asset, AudioAsset, ImageAsset, VideoAsset from .request_context._helpers import _infer_mime_type, _url_is_blocked -from .url_fetch import DEFAULT_MAX_BYTES, _CHUNK, open_guarded_stream +from .url_fetch import DEFAULT_MAX_BYTES, open_guarded_stream logger = logging.getLogger(__name__) @@ -46,6 +46,8 @@ _RESOLVE_TIMEOUT_S = 30 _MAX_RESOLVE_BODY = 8 << 20 _MAX_WALK_DEPTH = 32 +# Streaming read buffer, not a bound — it refuses nothing. +_CHUNK = 1 << 20 _INPUT_DIR_PREFIX = "gen-worker-inputs-" _RESOLVE_PATH = "/api/v1/worker/input-assets/resolve" _MANIFEST_KINDS = ("media", "image", "video", "audio") From cb10e310d82dcafdaec3eef46958b6ea10bf591b Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Thu, 6 Aug 2026 14:42:17 -0600 Subject: [PATCH 3/4] pgw#973: fold the wave-2 changelog bullets into changelog.d/pgw973.md MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit scripts/assemble_changelog.py requires .md — the number is what orders the release section — so 'pgw973-wave2.md' would have failed the cut, not the PR. Same issue, same file, no sibling lane touches it. (Pre-existing and untouched: changelog.d/pgw868-a4.md violates the same rule on master and belongs to the pgw#868 A1 lane.) --- changelog.d/{pgw973-wave2.md => pgw973.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changelog.d/{pgw973-wave2.md => pgw973.md} (100%) diff --git a/changelog.d/pgw973-wave2.md b/changelog.d/pgw973.md similarity index 100% rename from changelog.d/pgw973-wave2.md rename to changelog.d/pgw973.md From e0425c259ba21898b8676c045903ee68c85e430b Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Fri, 7 Aug 2026 01:37:47 -0600 Subject: [PATCH 4/4] pgw#973: two safetensors readers had NO header bound, and the census was structurally unable to see them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found while re-verifying this branch against a master that had moved 15 commits. Not new code: both have been unbounded since 2026-07-24 (6714ad8b). convert/writer.component_stored_tensor_names json.loads(fh.read(header_len)) models/loading (deshard path) json.loads(f.read(n)) Both read the attacker-controlled 8-byte prefix and allocate on it with no check whatsoever — the exact threat models/safetensors_header.py names, in two places the wave-2 collapse did not touch because there was no cap there to collapse. THE METHOD DEFECT, which is the part worth keeping: the §4.24 census enumerated every numeric bound and adjudicated it. That procedure can only find bounds that EXIST. A reader with no bound at all is invisible to an inventory of what is present — so "958 bounds counted and adjudicated" was never evidence that every site which needs one has one. Absence needs a different sweep than presence. Both now ask header_len_ok(). RED-verified: removing the writer guard makes a crafted file raise OverflowError: byte string is too large while attempting a 2**63-byte allocation. Also drops a now-unused MAX_HEADER_BYTES import from loading.py — the predicate moved into header_len_ok(), so only the helper is needed. --- changelog.d/pgw973.md | 11 +++++++ src/gen_worker/convert/writer.py | 3 ++ src/gen_worker/models/loading.py | 5 ++- ...est_one_safetensors_header_bound_pgw973.py | 31 ++++++++++++++++++- 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/changelog.d/pgw973.md b/changelog.d/pgw973.md index f8b87a39..21911b05 100644 --- a/changelog.d/pgw973.md +++ b/changelog.d/pgw973.md @@ -34,3 +34,14 @@ bound is `optimal_part_concurrency`'s `min(total_parts, 4)`. The semaphore is KEPT because it covers the one axis nothing else can see — an endpoint author calling `ctx.save()` from their own threads — and both docstrings now say so. +- **pgw#973: two safetensors readers had NO header bound at all, and the census + could not see them.** `convert/writer.component_stored_tensor_names` and the + deshard path in `models/loading` both did + `json.loads(f.read(header_len))` straight off an unvalidated 8-byte prefix — + present since 2026-07-24 (`6714ad8b`). Both now ask `header_len_ok()`. + RED-verified: without the guard a crafted file raises + `OverflowError: byte string is too large` attempting a 2**63-byte allocation. + **Recorded as a method defect, not just two fixes:** the §4.24 census + enumerated *bounds* and adjudicated each, so a reader carrying no bound was + invisible to it. An inventory of what is present cannot find what is absent. + diff --git a/src/gen_worker/convert/writer.py b/src/gen_worker/convert/writer.py index 0cd82bc4..fd3dca5c 100644 --- a/src/gen_worker/convert/writer.py +++ b/src/gen_worker/convert/writer.py @@ -856,6 +856,9 @@ def component_stored_tensor_names(component_dir: Path) -> frozenset[str]: for f in sorted(component_dir.glob("*.safetensors")): with open(f, "rb") as fh: header_len = struct.unpack(" Path: shard_path = snapshot_dir / shard with open(shard_path, "rb") as f: (n,) = struct.unpack("