diff --git a/soup_cli/utils/strict_safetensors.py b/soup_cli/utils/strict_safetensors.py index 35d3fe5..16e0189 100644 --- a/soup_cli/utils/strict_safetensors.py +++ b/soup_cli/utils/strict_safetensors.py @@ -1,12 +1,15 @@ """Strict safetensors mode — refuse pickle / PyTorch-classic weights (v0.60.0 Part C). -Static-extension allowlist. Treats any file with an extension in -``UNSAFE_EXTENSIONS`` as a potential arbitrary-code-execution vector and -refuses to proceed when ``strict=True``. Mirrors the HuggingFace safetensors -threat model (45% of HF repos still ship pickle weights as of late 2025). +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 +HuggingFace safetensors threat model (45% of HF repos still ship pickle +weights as of late 2025). Public surface: - ``UNSAFE_EXTENSIONS`` frozenset. +- ``is_safetensors_magic(path)`` -> bool. - ``find_unsafe_weight_files(model_dir)`` -> tuple of offending paths. - ``check_strict_safetensors(model_dir, *, strict=False)`` -> ``StrictSafetensorsReport``. @@ -16,6 +19,7 @@ so CI pipelines can grep specifically for strict-safetensors failures. from __future__ import annotations +import json import os from dataclasses import dataclass, field from typing import Tuple @@ -36,6 +40,17 @@ UNSAFE_EXTENSIONS = frozenset({ ".msgpack", # ambiguous binary blob — many loaders unpickle from this }) +SAFETENSORS_EXTENSION = ".safetensors" +_SAFETENSORS_HEADER_LEN_BYTES = 8 +_MAX_SAFETENSORS_HEADER_BYTES = 1 << 30 +_ZIP_MAGIC = b"PK\x03\x04" +_PICKLE_MAGIC_PREFIXES = ( + b"\x80\x02", + b"\x80\x03", + b"\x80\x04", + b"\x80\x05", +) + @dataclass(frozen=True) class StrictSafetensorsReport: @@ -47,8 +62,42 @@ class StrictSafetensorsReport: reason: str = field(default="") +def is_safetensors_magic(path: str) -> bool: + """Return whether ``path`` starts with a plausible safetensors header.""" + try: + file_size = os.path.getsize(path) + if file_size <= _SAFETENSORS_HEADER_LEN_BYTES: + return False + with open(path, "rb") as fh: + head = fh.read(16) + if head.startswith(_ZIP_MAGIC): + return False + if any(head.startswith(prefix) for prefix in _PICKLE_MAGIC_PREFIXES): + return False + if len(head) < _SAFETENSORS_HEADER_LEN_BYTES: + return False + header_len = int.from_bytes( + head[:_SAFETENSORS_HEADER_LEN_BYTES], "little" + ) + if header_len <= 0: + return False + if header_len > _MAX_SAFETENSORS_HEADER_BYTES: + return False + if header_len > file_size - _SAFETENSORS_HEADER_LEN_BYTES: + return False + fh.seek(_SAFETENSORS_HEADER_LEN_BYTES) + header = fh.read(header_len) + if len(header) != header_len: + return False + decoded = header.decode("utf-8") + parsed = json.loads(decoded) + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + return False + return isinstance(parsed, dict) + + def find_unsafe_weight_files(model_dir: str) -> Tuple[str, ...]: - """Walk ``model_dir`` and return every file matching ``UNSAFE_EXTENSIONS``. + """Walk ``model_dir`` and return unsafe or invalid weight files. Returns relative paths inside ``model_dir`` sorted alphabetically. Does NOT raise on a missing directory — that's the caller's job to gate. @@ -72,8 +121,10 @@ def find_unsafe_weight_files(model_dir: str) -> Tuple[str, ...]: continue for filename in files: ext = os.path.splitext(filename)[1].lower() + full = os.path.join(root, filename) if ext in UNSAFE_EXTENSIONS: - full = os.path.join(root, filename) + offenders.append(full) + elif ext == SAFETENSORS_EXTENSION and not is_safetensors_magic(full): offenders.append(full) return tuple(sorted(offenders)) @@ -117,7 +168,7 @@ def check_strict_safetensors( first = offenders[0] rel = os.path.relpath(first, model_dir) reason = ( - f"unsafe weight file (pickle / PyTorch-classic): {rel!r}; " + f"unsafe weight file (pickle / PyTorch-classic / invalid safetensors): {rel!r}; " "re-save as safetensors via " "`from safetensors.torch import save_file; save_file(...)`" ) diff --git a/tests/test_v0600_part_c.py b/tests/test_v0600_part_c.py index 355c789..953a673 100644 --- a/tests/test_v0600_part_c.py +++ b/tests/test_v0600_part_c.py @@ -20,9 +20,15 @@ from soup_cli.cli import app def _make_safetensors_only(tmp_path: Path) -> Path: + import numpy as np + from safetensors.numpy import save_file + target = tmp_path / "safe_adapter" target.mkdir() - (target / "adapter_model.safetensors").write_bytes(b"weights") + save_file( + {"weight": np.array([1.0], dtype=np.float32)}, + str(target / "adapter_model.safetensors"), + ) (target / "adapter_config.json").write_text('{"r": 8}', encoding="utf-8") return target @@ -35,6 +41,32 @@ def _make_with_pickle(tmp_path: Path) -> Path: return target +def _make_with_renamed_pickle(tmp_path: Path) -> Path: + target = tmp_path / "renamed_pickle_adapter" + target.mkdir() + (target / "adapter_model.safetensors").write_bytes(b"\x80\x04pickled-bytes") + (target / "adapter_config.json").write_text('{"r": 8}', encoding="utf-8") + return target + + +def _make_with_renamed_zip(tmp_path: Path) -> Path: + target = tmp_path / "renamed_zip_adapter" + target.mkdir() + (target / "adapter_model.safetensors").write_bytes(b"PK\x03\x04zipped-bytes") + (target / "adapter_config.json").write_text('{"r": 8}', encoding="utf-8") + return target + + +def _make_with_corrupt_safetensors_header(tmp_path: Path) -> Path: + target = tmp_path / "corrupt_safetensors_adapter" + target.mkdir() + (target / "adapter_model.safetensors").write_bytes( + (1024).to_bytes(8, "little") + b"{}" + ) + (target / "adapter_config.json").write_text('{"r": 8}', encoding="utf-8") + return target + + class TestStrictSafetensors: def test_imports(self): from soup_cli.utils.strict_safetensors import ( @@ -42,9 +74,11 @@ class TestStrictSafetensors: StrictSafetensorsReport, check_strict_safetensors, find_unsafe_weight_files, + is_safetensors_magic, ) assert callable(check_strict_safetensors) assert callable(find_unsafe_weight_files) + assert callable(is_safetensors_magic) assert isinstance(UNSAFE_EXTENSIONS, frozenset) assert dataclasses.is_dataclass(StrictSafetensorsReport) @@ -99,6 +133,32 @@ class TestStrictSafetensors: assert report.ok is False assert len(report.unsafe_files) == 1 + def test_find_unsafe_flags_pickle_renamed_safetensors(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + adapter = _make_with_renamed_pickle(tmp_path) + from soup_cli.utils.strict_safetensors import find_unsafe_weight_files + + found = find_unsafe_weight_files(str(adapter)) + assert len(found) == 1 + assert found[0].endswith("adapter_model.safetensors") + + def test_check_strict_zip_renamed_safetensors_raises(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + adapter = _make_with_renamed_zip(tmp_path) + from soup_cli.utils.strict_safetensors import check_strict_safetensors + + with pytest.raises(ValueError, match="(?i)safetensors|unsafe"): + check_strict_safetensors(str(adapter), strict=True) + + def test_check_lenient_corrupt_safetensors_returns_report(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + adapter = _make_with_corrupt_safetensors_header(tmp_path) + from soup_cli.utils.strict_safetensors import check_strict_safetensors + + report = check_strict_safetensors(str(adapter), strict=False) + assert report.ok is False + assert len(report.unsafe_files) == 1 + def test_check_outside_cwd_rejected(self, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) from soup_cli.utils.strict_safetensors import check_strict_safetensors