From 6ddaeb30d104a97b32492aeed0f07981074b8c44 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Tue, 19 May 2026 23:56:42 +0500 Subject: [PATCH] polish(v0.60.0 strict-safetensors): tighten header cap + input guards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to PR #198 (issue #189): - _MAX_SAFETENSORS_HEADER_BYTES: 1 GiB -> 100 MiB. Real safetensors headers are <10 MiB even for 70B-parameter models; 100 MiB is a generous defence-in-depth ceiling that still rejects an adversary's "header_len = 999 MiB" allocation attempt before fh.read() commits. - is_safetensors_magic: input-shape guards (non-string / empty / null-byte path return False, never raise). Matches project policy for detection-style helpers (mirrors v0.30.0 Candidate, v0.41.0 lr_groups, v0.53.3 is_known_vlm_base). - +2 regression tests in tests/test_v0600_part_c.py: - test_is_safetensors_magic_rejects_invalid_input (5 bad inputs) - test_max_safetensors_header_bytes_tightened (guards against re-widening to 1 GiB in a future patch) Module docstring updated to reference PR #198 / issue #189. Verified locally: - ruff check soup_cli/utils/strict_safetensors.py tests/test_v0600_part_c.py -> clean - pytest tests/test_v0600_part_c.py --no-cov -> 21 passed, 1 POSIX-skipped on Windows Closes v0.60.0 Known Limitation (6) — full magic-byte + JSON-header shape verification with hardened input surface. --- soup_cli/utils/strict_safetensors.py | 27 +++++++++++++++++++++------ tests/test_v0600_part_c.py | 24 ++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/soup_cli/utils/strict_safetensors.py b/soup_cli/utils/strict_safetensors.py index 16e0189..e2cd153 100644 --- a/soup_cli/utils/strict_safetensors.py +++ b/soup_cli/utils/strict_safetensors.py @@ -1,9 +1,11 @@ """Strict safetensors mode — refuse pickle / PyTorch-classic weights (v0.60.0 Part C). -Static-extension allowlist plus a small safetensors header check. Treats any -file with an extension in ``UNSAFE_EXTENSIONS`` as a potential -arbitrary-code-execution vector, and refuses ``.safetensors`` files whose -bytes do not begin with a plausible safetensors metadata header. Mirrors the +Static-extension allowlist plus a safetensors magic-byte header check +(PR #198 / issue #189). Treats any file with an extension in +``UNSAFE_EXTENSIONS`` as a potential arbitrary-code-execution vector, and +refuses ``.safetensors`` files whose bytes do not begin with a plausible +safetensors metadata header (rejects zip / pickle opcodes / implausible u64 +header_len, then JSON-parses the header to confirm dict shape). Mirrors the HuggingFace safetensors threat model (45% of HF repos still ship pickle weights as of late 2025). @@ -42,7 +44,10 @@ UNSAFE_EXTENSIONS = frozenset({ SAFETENSORS_EXTENSION = ".safetensors" _SAFETENSORS_HEADER_LEN_BYTES = 8 -_MAX_SAFETENSORS_HEADER_BYTES = 1 << 30 +# Real safetensors headers are <10 MiB even on 70B-parameter models; +# 100 MiB is a generous defence-in-depth ceiling that still rejects an +# adversary's "header_len = 999 MiB" allocation attempt. +_MAX_SAFETENSORS_HEADER_BYTES = 100 * (1 << 20) _ZIP_MAGIC = b"PK\x03\x04" _PICKLE_MAGIC_PREFIXES = ( b"\x80\x02", @@ -63,7 +68,17 @@ class StrictSafetensorsReport: def is_safetensors_magic(path: str) -> bool: - """Return whether ``path`` starts with a plausible safetensors header.""" + """Return whether ``path`` starts with a plausible safetensors header. + + Defensive surface — returns ``False`` (never raises) on non-string / + empty / null-byte input. Matches project policy for + detection-style helpers (mirrors v0.30.0 ``Candidate``, v0.41.0 + ``lr_groups``, v0.53.3 ``is_known_vlm_base``). + """ + if not isinstance(path, str) or not path: + return False + if "\x00" in path: + return False try: file_size = os.path.getsize(path) if file_size <= _SAFETENSORS_HEADER_LEN_BYTES: diff --git a/tests/test_v0600_part_c.py b/tests/test_v0600_part_c.py index 953a673..5494a95 100644 --- a/tests/test_v0600_part_c.py +++ b/tests/test_v0600_part_c.py @@ -90,6 +90,30 @@ class TestStrictSafetensors: assert ".pth" in UNSAFE_EXTENSIONS assert ".ckpt" in UNSAFE_EXTENSIONS + def test_is_safetensors_magic_rejects_invalid_input(self): + """Defensive surface — non-string / empty / null-byte returns False (never raises). + + Matches project policy for detection-style helpers (v0.30.0 Candidate / + v0.41.0 lr_groups / v0.53.3 is_known_vlm_base). + """ + from soup_cli.utils.strict_safetensors import is_safetensors_magic + + assert is_safetensors_magic("") is False + assert is_safetensors_magic("path\x00null") is False + assert is_safetensors_magic(None) is False # type: ignore[arg-type] + assert is_safetensors_magic(123) is False # type: ignore[arg-type] + assert is_safetensors_magic(b"bytes/path") is False # type: ignore[arg-type] + + def test_max_safetensors_header_bytes_tightened(self): + """Defence-in-depth: header cap is 100 MiB (real headers <10 MiB). + + Regression guard against a re-widening to 1 GiB which would let an + adversary trigger a ~999 MiB allocation via a crafted header_len. + """ + from soup_cli.utils.strict_safetensors import _MAX_SAFETENSORS_HEADER_BYTES + + assert _MAX_SAFETENSORS_HEADER_BYTES == 100 * (1 << 20) + def test_find_unsafe_clean(self, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) adapter = _make_safetensors_only(tmp_path)