polish(v0.60.0 strict-safetensors): tighten header cap + input guards

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.
This commit is contained in:
Alpamys 2026-05-19 23:56:42 +05:00
parent 914a299965
commit 6ddaeb30d1
2 changed files with 45 additions and 6 deletions

View File

@ -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:

View File

@ -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)