Check safetensors magic bytes (#198)

Co-authored-by: Sumit Dhawan <sumitdhawan@Sumits-MacBook-Air.local>

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(...)`"
     )

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