mirror of https://github.com/razor-ai/soup.git
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:
parent
914a299965
commit
6ddaeb30d1
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue