6e91d1dff8* Studio: scan HF cache snapshot loads by their repo id Inactive Hugging Face caches (legacy, default, and previously selected download locations) are loaded by their resolved snapshot path so they keep using the selected cache instead of re-downloading. That path is a local filesystem path, so evaluate_file_security exempted it with "local path; no Hub scan" and skipped Hugging Face's pickle/malware scan. Active caches load by repo id and are still scanned, so the same model could dodge the gate simply by being in an inactive cache. An HF cache snapshot keeps the canonical models--org--repo/snapshots/<rev> layout, so recover the repo id from that path and scan it instead of exempting it. Non-cache local paths (models directory, custom folders) still skip the scan, and a remote ref is still scanned by repo id. Adds a regression test that a flagged pickle in an inactive-cache snapshot path blocks the load. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: scan the exact cached commit for inactive HF caches An HF cache snapshot path encodes the commit, not just the repo id (models--org--repo/snapshots/<rev>). Recover the revision alongside the repo id and pass it to model_info and the shard-index lookup so the scan covers the exact files that will be deserialized, rather than the repo's default branch. Without this, a pickle in an older cached commit that was later removed from the branch would scan clean and still load. Extends the regression test to assert the recovered revision is forwarded to the Hub scan. --------- Co-authored-by: danielhanchen <unslothai@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
609 lines
28 KiB
Python
609 lines
28 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Malware / unsafe-file gate for model loads.
|
|
|
|
The ``trust_remote_code`` consent gate covers the ``auto_map`` Python vector; this
|
|
covers the other one -- a malicious pickle inside a weight file, which executes
|
|
during ``from_pretrained`` deserialization even with ``trust_remote_code=False``.
|
|
It reads Hugging Face's OWN scan (picklescan + ClamAV) via
|
|
``model_info(securityStatus=True).security_repo_status``. METADATA-ONLY: it never
|
|
downloads, opens, or unpickles the flagged files.
|
|
|
|
Policy:
|
|
* Hard block, non-approvable.
|
|
* Block whenever ``filesWithIssues`` lists a non-``safe`` level, regardless of
|
|
``scansDone`` (often false even for clean repos). Unknown/future levels fail
|
|
CLOSED (block) so Hub schema drift cannot silently allow a bad verdict; only a
|
|
small allowlist of clean / not-yet-scanned levels is non-blocking. The sole
|
|
fail-open path is an unavailable status (missing field / offline / error).
|
|
* Scope to the load-path RCE vector: a root-level (or load-subdir-level),
|
|
code-executing file. Inert formats (safetensors / gguf / config / text) and
|
|
subdirectory pickles that no root weight-index references are NOT loaded, so
|
|
they do not block; an index-referenced shard does, wherever it lives. This
|
|
blocks real malware (eicar's root ``*.pkl``/``*.dat``) without false-blocking
|
|
repos like ``nvidia/Nemotron-H-8B-Base-8K`` (flagged NeMo pickles under
|
|
``nemo/`` that no index lists).
|
|
* No first-party exemption (scoping is by load path/format, not org).
|
|
* Local paths are skipped (no Hub scan); a remote ``*.gguf``-named repo is still
|
|
scanned so a repo cannot dodge the gate by suffixing its name.
|
|
"""
|
|
|
|
import re
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from loggers import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
# Pickle-format weight files (plain or sharded) that execute code on load; safetensors/gguf
|
|
# are inert. Grouped by weight family so an inert safetensors only suppresses the pickle it
|
|
# actually replaces: the loader won't use an adapter's safetensors for pytorch_model.bin.
|
|
_PICKLE_WEIGHT_RE = re.compile(
|
|
r"^(model|pytorch_model|adapter_model|consolidated)(-\d+-of-\d+)?"
|
|
r"\.(bin|pt|pth|ckpt|pkl|pickle)$",
|
|
re.IGNORECASE,
|
|
)
|
|
|
|
# Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/
|
|
# malicious or a future label) blocks, so Hub schema drift fails CLOSED.
|
|
_NONBLOCKING_LEVELS = frozenset(
|
|
{"", "safe", "pending", "scanning", "queued", "unscanned", "error", "unknown", "none"}
|
|
)
|
|
|
|
# Suffixes that cannot execute code on load (tensor-only safetensors, non-pickle gguf,
|
|
# text/markup/images), so a flag on one is never an RCE vector.
|
|
_INERT_SUFFIXES = frozenset(
|
|
{
|
|
".safetensors",
|
|
".gguf",
|
|
".json",
|
|
".txt",
|
|
".md",
|
|
".rst",
|
|
".yaml",
|
|
".yml",
|
|
".png",
|
|
".jpg",
|
|
".jpeg",
|
|
".gif",
|
|
".webp",
|
|
".svg",
|
|
".bmp",
|
|
".gitattributes",
|
|
".gitignore",
|
|
}
|
|
)
|
|
|
|
# Source files are not deserialized by a weight load; executable repo code runs only
|
|
# via auto_map, which is the consent gate's domain. So a flag on a .py is not this
|
|
# gate's vector (else a flagged helper/train script would false-block).
|
|
_SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"})
|
|
|
|
|
|
# Torch-family weight indexes: from_pretrained feeds each shard they name to load_state_dict, which
|
|
# torch.load()s (pickle) any shard whose name does not end in .safetensors, whatever its stem. A
|
|
# pytorch index is superseded when a base safetensors is present (the loader prefers it); a
|
|
# safetensors index IS the chosen archive, so a non-safetensors target it names still loads. tf/flax
|
|
# indexes load via non-pickle loaders, so they are not a torch.load vector here.
|
|
_TORCH_INDEX_FILES = ("pytorch_model.bin.index.json", "model.safetensors.index.json")
|
|
|
|
# Root weight-index files. from_pretrained reads these to find sharded weights, so a
|
|
# flagged subdir pickle is a load vector iff a root index references it.
|
|
_TRANSFORMERS_INDEX_FILES = (
|
|
"pytorch_model.bin.index.json",
|
|
"model.safetensors.index.json",
|
|
"tf_model.h5.index.json",
|
|
"flax_model.msgpack.index.json",
|
|
)
|
|
|
|
|
|
def _normalize_repo_path(path: str) -> str:
|
|
"""Strip ``./`` prefixes and normalize separators for repo-relative comparison."""
|
|
p = (path or "").strip().replace("\\", "/")
|
|
while p.startswith("./"):
|
|
p = p[2:]
|
|
return p
|
|
|
|
|
|
def _file_suffix(path: str) -> str:
|
|
"""Lowercase ``.ext`` of the basename, or ``""`` if none."""
|
|
base = _normalize_repo_path(path).rsplit("/", 1)[-1]
|
|
return "." + base.rsplit(".", 1)[1].lower() if "." in base else ""
|
|
|
|
|
|
def _hf_cache_snapshot_ref(local_path: str) -> Optional[tuple]:
|
|
"""``(repo_id, revision)`` for an HF-cache snapshot path, else None. An inactive Studio
|
|
cache loads by its snapshot path but keeps the ``models--org--repo/snapshots/<rev>``
|
|
layout, so the gate recovers its provenance and scans that exact commit instead of
|
|
exempting it (an older cached commit can hold a pickle since dropped from the branch)."""
|
|
try:
|
|
path = Path(local_path).resolve(strict = False)
|
|
except (OSError, ValueError):
|
|
return None
|
|
for parent in path.parents:
|
|
if parent.name != "snapshots":
|
|
continue
|
|
encoded = parent.parent.name
|
|
if not encoded.startswith("models--"):
|
|
return None
|
|
repo_id = encoded.removeprefix("models--").replace("--", "/")
|
|
if not repo_id:
|
|
return None
|
|
return repo_id, path.relative_to(parent).parts[0] # <rev> dir under snapshots/
|
|
return None
|
|
|
|
|
|
def _load_relative_path(norm: str, load_subdirs) -> str:
|
|
"""``norm`` relative to a ``from_pretrained`` load root. Some loads read from a
|
|
snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``<snapshot>/LLM``), where a file
|
|
directly under the subdir is root-level, not nested. Strips the matching load-subdir
|
|
prefix, or returns ``norm`` unchanged when it is not under one.
|
|
"""
|
|
for subdir in load_subdirs or ():
|
|
prefix = _normalize_repo_path(subdir).strip("/")
|
|
if prefix and norm.startswith(prefix + "/"):
|
|
return norm[len(prefix) + 1 :]
|
|
return norm
|
|
|
|
|
|
def _index_prefixes(load_subdirs) -> tuple:
|
|
"""Prefixes to look for weight-index files under: repo root plus each load subdir."""
|
|
prefixes = [""]
|
|
for subdir in load_subdirs or ():
|
|
p = _normalize_repo_path(subdir).strip("/")
|
|
if p:
|
|
prefixes.append(p + "/")
|
|
return tuple(prefixes)
|
|
|
|
|
|
def _indexed_shard_paths(
|
|
model_name: str,
|
|
hf_token: Optional[str],
|
|
load_subdirs = (),
|
|
revision: Optional[str] = None,
|
|
):
|
|
"""Repo-relative weight paths a load could fetch via weight-index files. Returns a
|
|
set (empty when the repo ships no index files -- a definitive "nothing sharded"), or
|
|
None when the lookup was inconclusive (transient error) so the caller treats a
|
|
flagged subdir pickle conservatively. Reads only small JSON indexes, never weights.
|
|
Indexes are looked up at the root and each ``load_subdirs`` root, with ``weight_map``
|
|
entries re-prefixed to repo-relative paths. ``revision`` scopes to a cached commit.
|
|
"""
|
|
import json
|
|
|
|
try:
|
|
from huggingface_hub import hf_hub_download
|
|
from huggingface_hub.utils import EntryNotFoundError
|
|
from utils.hf_cache_settings import active_hf_hub_cache
|
|
except Exception:
|
|
return None
|
|
|
|
paths: set = set()
|
|
inconclusive = False
|
|
for prefix in _index_prefixes(load_subdirs):
|
|
for filename in _TRANSFORMERS_INDEX_FILES:
|
|
try:
|
|
index_path = hf_hub_download(
|
|
model_name,
|
|
prefix + filename,
|
|
revision = revision,
|
|
token = hf_token or None,
|
|
cache_dir = active_hf_hub_cache(),
|
|
)
|
|
except EntryNotFoundError:
|
|
continue # definitively absent, not an error
|
|
except Exception:
|
|
inconclusive = True # transient: an index that might exist could not be read
|
|
continue
|
|
try:
|
|
weight_map = (json.loads(open(index_path).read()) or {}).get("weight_map") or {}
|
|
for shard in weight_map.values():
|
|
shard_norm = _normalize_repo_path(str(shard))
|
|
# weight_map paths are relative to the index file's directory.
|
|
if prefix and not shard_norm.startswith(prefix):
|
|
shard_norm = prefix + shard_norm
|
|
paths.add(shard_norm)
|
|
except Exception:
|
|
inconclusive = True
|
|
# Any transient failure -> inconclusive (the shard could be listed only by the index
|
|
# we could not read), so fail closed (None) and let the caller block. Ships no index
|
|
# files -> EntryNotFoundError for each, empty set, a definitive "nothing sharded".
|
|
if inconclusive:
|
|
return None
|
|
return paths
|
|
|
|
|
|
# Two-timeout metadata fetch, mirroring hub.workers.hf_download._retry_metadata_fetch.
|
|
_REQUEST_TIMEOUT = 10.0
|
|
_RETRY_TIMEOUT = 20.0
|
|
|
|
|
|
@dataclass
|
|
class FileSecurityDecision:
|
|
"""Outcome of the Hub security scan for one model repo."""
|
|
|
|
model_name: str
|
|
blocked: bool
|
|
unsafe_files: list = field(default_factory = list) # [{"path", "level"}]
|
|
reason: str = ""
|
|
|
|
def response_payload(self) -> dict:
|
|
"""Machine-readable detail merged into the preflight payload the dialog reads."""
|
|
return {
|
|
"unsafe_files": self.unsafe_files,
|
|
"security_blocked": self.blocked,
|
|
"reason": self.reason,
|
|
}
|
|
|
|
|
|
def security_load_subdirs(model_name: str, hf_token: Optional[str] = None) -> tuple:
|
|
"""Snapshot subdirectories a load calls ``from_pretrained`` on, for scoping the scan.
|
|
Most models load from the root (``()``); Spark-TTS / BiCodec load ``<snapshot>/LLM``,
|
|
so ``LLM/`` is a load root for them. Metadata-only (tokenizer special tokens), cached.
|
|
"""
|
|
try:
|
|
from utils.models.model_config import detect_audio_type, load_model_defaults
|
|
if detect_audio_type(model_name, hf_token = hf_token) == "bicodec":
|
|
return ("LLM",)
|
|
# Tokenizer detection can fail (network/gated/unresolved alias); the YAML default
|
|
# also pins the audio type, so fall back to it (else a flagged LLM/ pickle is
|
|
# treated as an ignored subdir artifact).
|
|
if (load_model_defaults(model_name) or {}).get("audio_type") == "bicodec":
|
|
return ("LLM",)
|
|
except Exception:
|
|
pass
|
|
return ()
|
|
|
|
|
|
def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple:
|
|
"""Map a load alias to the ``(repo_id, load_subdirs)`` the load actually fetches. The
|
|
Spark-TTS / BiCodec alias ``<parent>/LLM`` is downloaded by the trainer as
|
|
``unsloth/<parent>`` and loaded from ``LLM/``, so scan that repo with ``LLM`` as a
|
|
load root (the literal alias 404s and fails open). Everything else is unchanged.
|
|
"""
|
|
try:
|
|
from utils.paths import is_local_path
|
|
if is_local_path(model_name):
|
|
return model_name, load_subdirs
|
|
except Exception:
|
|
return model_name, load_subdirs
|
|
name = (model_name or "").strip().strip("/")
|
|
# Rewrite ONLY a registry-known bicodec alias, never any repo ending in "/LLM"
|
|
# (e.g. "evil/LLM" would scan unsloth/evil and fail open on the real repo).
|
|
if name.endswith("/LLM") and name.count("/") == 1:
|
|
try:
|
|
from utils.models.model_config import load_model_defaults
|
|
if (load_model_defaults(name) or {}).get("audio_type") == "bicodec":
|
|
parent = name[: -len("/LLM")]
|
|
return f"unsloth/{parent}", tuple(dict.fromkeys((*load_subdirs, "LLM")))
|
|
except Exception:
|
|
pass
|
|
return model_name, load_subdirs
|
|
|
|
|
|
def _fetch_security_status(
|
|
model_name: str,
|
|
hf_token: Optional[str],
|
|
revision: Optional[str] = None,
|
|
):
|
|
"""``security_repo_status`` (a dict) or None if unavailable. Hub metadata only;
|
|
retries once on a transient error, then returns None so the caller fails open.
|
|
``revision`` scopes the scan to a specific cached commit (else the default branch).
|
|
"""
|
|
from huggingface_hub import model_info as hf_model_info
|
|
|
|
token_arg = hf_token if hf_token else False
|
|
last_exc = None
|
|
for attempt, timeout in enumerate((_REQUEST_TIMEOUT, _RETRY_TIMEOUT)):
|
|
try:
|
|
info = hf_model_info(
|
|
model_name,
|
|
revision = revision,
|
|
token = token_arg,
|
|
securityStatus = True,
|
|
timeout = timeout,
|
|
)
|
|
return getattr(info, "security_repo_status", None)
|
|
except Exception as exc: # network/offline/gated/404/unsupported-client
|
|
last_exc = exc
|
|
if attempt == 0:
|
|
continue
|
|
logger.debug(
|
|
"HF security scan unavailable for '%s' (%s); failing open.",
|
|
model_name,
|
|
type(last_exc).__name__ if last_exc else "unknown",
|
|
)
|
|
return None
|
|
|
|
|
|
def _st_load_roots(snapshot: Path) -> list:
|
|
"""Directories a SentenceTransformer load deserializes weights from: the snapshot root plus
|
|
each module path in modules.json. Local, no network. Mirrors the online gate (which ignores
|
|
unreferenced nested pickles ST never loads) so the offline gate doesn't over-block."""
|
|
roots = [snapshot]
|
|
try:
|
|
import json
|
|
modules = json.loads((snapshot / "modules.json").read_text())
|
|
except (OSError, ValueError):
|
|
return roots # no / invalid modules.json -> snapshot root is the only load root
|
|
for module in modules or ():
|
|
path = str((module or {}).get("path", "")).strip().strip("/")
|
|
# Relative module path only; ignore a crafted "../" escape.
|
|
if path and ".." not in path.split("/"):
|
|
candidate = snapshot / path
|
|
if candidate not in roots:
|
|
roots.append(candidate)
|
|
return roots
|
|
|
|
|
|
def _indexed_pickle_shards(index_path: Path, root: Path, snapshot: Path) -> list:
|
|
"""Shards a torch weight index points a ``from_pretrained`` load at that load_state_dict would
|
|
torch.load (pickle): every ``weight_map`` target NOT ending in ``.safetensors``, whatever its
|
|
stem (an arbitrary name like ``shards/payload`` still deserializes). Resolved relative to the
|
|
index dir (``root``) like the loader, so a shard in a nested dir is followed (iterdir misses it).
|
|
Lexical only, never ``Path.resolve()`` (HF snapshot files symlink into ``blobs/``, so resolving
|
|
escapes the snapshot and false-blocks every shard). Raises OSError -> caller fails CLOSED on an
|
|
unreadable/invalid index or a target escaping the snapshot."""
|
|
import json
|
|
import os
|
|
|
|
try:
|
|
# JSON is UTF-8 by spec; pin it so a non-ASCII index is not misdecoded (and needlessly
|
|
# blocked) under Windows' cp1252 default.
|
|
parsed = json.loads(index_path.read_text(encoding = "utf-8"))
|
|
except (OSError, ValueError) as exc:
|
|
raise OSError(f"unreadable weight index: {index_path}") from exc
|
|
weight_map = parsed.get("weight_map") if isinstance(parsed, dict) else None
|
|
if not isinstance(weight_map, dict):
|
|
return [] # no dict weight_map -> the loader resolves no shards from this index
|
|
snapshot_norm = os.path.normpath(str(snapshot))
|
|
shards = []
|
|
for shard in weight_map.values():
|
|
raw = str(shard)
|
|
if not raw:
|
|
continue
|
|
# Join the RAW weight_map value like from_pretrained's os.path.join: on POSIX a backslash is a
|
|
# literal filename char (not a separator), so normalizing it would probe a different path than
|
|
# the loader opens. normpath + containment stay platform-aware (os.sep) to block "..".
|
|
joined = os.path.normpath(os.path.join(str(root), raw))
|
|
if joined != snapshot_norm and not joined.startswith(snapshot_norm + os.sep):
|
|
raise OSError(f"weight index escapes the snapshot: {index_path}")
|
|
shard_path = Path(joined)
|
|
# Case-SENSITIVE, mirroring load_state_dict's own endswith(".safetensors"): a shard named
|
|
# payload.SAFETENSORS is not treated as safetensors by the loader and falls to torch.load.
|
|
if not shard_path.name.endswith(".safetensors") and shard_path.is_file():
|
|
shards.append(shard_path)
|
|
return shards
|
|
|
|
|
|
def _loader_resolves(root: Path, name: str) -> bool:
|
|
"""True iff from_pretrained would open ``name`` under ``root``. ``is_file()`` honors the platform
|
|
(case-sensitive on Linux, case-insensitive on Windows/macOS), so it mirrors the loader's own
|
|
lookup: an oddly-cased decoy counts as an alternative only where the loader would truly open it.
|
|
A name-fold instead would let an uppercase MODEL.SAFETENSORS suppress the scan on Linux while the
|
|
loader, asking for the canonical lowercase name, silently falls through to a pickle index."""
|
|
return (root / name).is_file()
|
|
|
|
|
|
def _cached_pickle_weight_files(snapshot: Path) -> list:
|
|
"""Pickle weight files a SentenceTransformer/Transformers load deserializes from snapshot's ST
|
|
load roots, EXCLUDING those whose weight family also ships an inert safetensors in the same dir
|
|
(the loader prefers it): a base pickle is suppressed only by a base model.safetensors, an adapter
|
|
pickle only by adapter_model.safetensors -- an unrelated safetensors is no substitute. Covers
|
|
both direct-child pickles AND pickle shards referenced by a local weight index (which the loader
|
|
follows into nested dirs, matching the online gate). Raises OSError -- caller fails CLOSED -- if
|
|
the snapshot root or a weight index is unreadable, or an index reference escapes the snapshot."""
|
|
blocked = []
|
|
seen = set()
|
|
|
|
def _add(path: Path):
|
|
key = str(path)
|
|
if key not in seen:
|
|
seen.add(key)
|
|
blocked.append(path)
|
|
|
|
for root in _st_load_roots(snapshot):
|
|
try:
|
|
entries = [p for p in root.iterdir() if p.is_file()]
|
|
except OSError:
|
|
if root == snapshot:
|
|
raise # top-level unreadable -> fail closed
|
|
continue # unreadable module subdir: nothing loadable to attest here
|
|
# Safetensors alternatives the loader would actually resolve (never a bare name-fold, which
|
|
# fails OPEN: see _loader_resolves). A base pickle is replaced only by a base safetensors, an
|
|
# adapter pickle only by an adapter one. A single model.safetensors also outranks BOTH indexes.
|
|
has_direct_base_safetensors = _loader_resolves(root, "model.safetensors")
|
|
has_base_safetensors = has_direct_base_safetensors or _loader_resolves(
|
|
root, "model.safetensors.index.json"
|
|
)
|
|
has_adapter_safetensors = _loader_resolves(root, "adapter_model.safetensors")
|
|
for path in entries:
|
|
if not _PICKLE_WEIGHT_RE.match(path.name):
|
|
continue
|
|
is_adapter = path.name.lower().startswith("adapter_model")
|
|
has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors
|
|
if not has_alternative:
|
|
_add(path)
|
|
# A torch weight index makes from_pretrained load nested shards iterdir never sees; the loader
|
|
# torch.loads any not ending in .safetensors. Probe the canonical index name with the loader's
|
|
# own lookup (_loader_resolves), so an oddly-cased artifact it would never open does not block.
|
|
# A direct model.safetensors wins over BOTH indexes; failing that a base safetensors still
|
|
# outranks the pytorch index, while a safetensors index is itself the chosen archive.
|
|
for index_name in _TORCH_INDEX_FILES:
|
|
if not _loader_resolves(root, index_name):
|
|
continue
|
|
if has_direct_base_safetensors:
|
|
continue
|
|
if index_name == "pytorch_model.bin.index.json" and has_base_safetensors:
|
|
continue
|
|
for shard_path in _indexed_pickle_shards(root / index_name, root, snapshot):
|
|
_add(shard_path)
|
|
return blocked
|
|
|
|
|
|
def _evaluate_local_only(model_name: str) -> FileSecurityDecision:
|
|
"""Offline security gate. The Hub scan is unreachable, so inspect the local cache and fail
|
|
CLOSED on an unscanned pickle weight with no inert safetensors alternative, rather than
|
|
failing open or hanging. Safetensors/gguf-only cache loads; nothing cached -> allowed."""
|
|
from utils.utils import hf_cache_snapshot_dir
|
|
|
|
try:
|
|
snapshot = hf_cache_snapshot_dir(model_name)
|
|
except Exception:
|
|
logger.warning("Offline gate: could not resolve the cache for '%s'; blocking.", model_name)
|
|
return FileSecurityDecision(
|
|
model_name, True, reason = "offline; could not inspect the local cache"
|
|
)
|
|
|
|
if snapshot is None:
|
|
return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to load")
|
|
|
|
try:
|
|
pickles = _cached_pickle_weight_files(snapshot)
|
|
except OSError:
|
|
logger.warning("Offline gate: could not read the cache for '%s'; blocking.", model_name)
|
|
return FileSecurityDecision(
|
|
model_name, True, reason = "offline; could not read the local cache"
|
|
)
|
|
|
|
if not pickles:
|
|
return FileSecurityDecision(
|
|
model_name, False, reason = "offline; cached weights are inert (safetensors/gguf)"
|
|
)
|
|
|
|
# Snapshot-relative posix paths (match the online gate; disambiguate same-named pickles).
|
|
rel_paths = sorted(p.relative_to(snapshot).as_posix() for p in pickles)
|
|
names = ", ".join(rel_paths)
|
|
logger.warning(
|
|
"Blocking offline load of '%s': cached pickle weight(s) cannot be malware-scanned "
|
|
"offline and have no safetensors alternative (%s).",
|
|
model_name,
|
|
names,
|
|
)
|
|
return FileSecurityDecision(
|
|
model_name,
|
|
True,
|
|
unsafe_files = [{"path": rel, "level": "unscanned"} for rel in rel_paths],
|
|
reason = f"offline; unscanned pickle weights with no safetensors alternative: {names}",
|
|
)
|
|
|
|
|
|
def evaluate_file_security(
|
|
model_name: str,
|
|
hf_token: Optional[str] = None,
|
|
*,
|
|
load_subdirs = (),
|
|
local_only_load: bool = False,
|
|
) -> FileSecurityDecision:
|
|
"""Block a load when HF's security scan flags unsafe serialized files.
|
|
|
|
Call UNCONDITIONALLY before any load (independent of trust_remote_code): a malicious
|
|
pickle deserializes during ``from_pretrained`` regardless. Metadata-only; fails open
|
|
when the scan is unavailable.
|
|
|
|
``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)``
|
|
for Spark-TTS / BiCodec, loading ``<snapshot>/LLM``): a flagged file directly under one
|
|
is root-level there and blocks, and an index inside it is honored when scoping shards.
|
|
|
|
``local_only_load`` marks an offline load: with the Hub scan unreachable, inspect the local
|
|
cache and fail CLOSED on an unscanned pickle weight with no safetensors alternative.
|
|
"""
|
|
# Scan the repo the load actually fetches, not the literal alias (which 404s and
|
|
# fails open): the Spark-TTS "<parent>/LLM" alias is really unsloth/<parent> from LLM/.
|
|
model_name, load_subdirs = _load_scan_target(model_name, tuple(load_subdirs))
|
|
|
|
# Local paths have no Hub scan, EXCEPT an HF-cache snapshot whose canonical path
|
|
# encodes a repo id + commit: scan that exact commit so an inactive-cache load can't
|
|
# dodge the gate. A remote ref is scanned even if named "*.gguf" (name can't dodge it).
|
|
snapshot_revision = None
|
|
try:
|
|
from utils.paths import is_local_path
|
|
if is_local_path(model_name):
|
|
cache_ref = _hf_cache_snapshot_ref(model_name)
|
|
if cache_ref is None:
|
|
return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan")
|
|
model_name, snapshot_revision = cache_ref
|
|
except Exception:
|
|
# Cannot classify the path -> do not block on that account.
|
|
return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked")
|
|
|
|
# Offline: inspect the local cache and fail closed rather than hang on model_info or fail open.
|
|
if local_only_load:
|
|
return _evaluate_local_only(model_name)
|
|
|
|
status = _fetch_security_status(model_name, hf_token, revision = snapshot_revision)
|
|
if not isinstance(status, dict):
|
|
return FileSecurityDecision(
|
|
model_name, False, reason = "scan unavailable; allowed (fail-open)"
|
|
)
|
|
|
|
# Block a non-``safe`` flagged file scoped to the load-path RCE vector (root-level,
|
|
# code-executing). Not gated on ``scansDone`` (often false even when clean; a flagged
|
|
# file is flagged regardless). Unknown levels fail closed; in-progress/clean do not.
|
|
# Subdir pickles and inert formats (safetensors/gguf) are not loaded by
|
|
# from_pretrained and do not block. Unavailable status (above) is the only fail-open.
|
|
unsafe = []
|
|
skipped = [] # flagged, but not a load-path RCE vector (subdir artifact / inert)
|
|
maybe_shard = [] # flagged subdir pickle: a load vector ONLY if a root index lists it
|
|
for entry in status.get("filesWithIssues") or []:
|
|
if not isinstance(entry, dict):
|
|
continue
|
|
level = str(entry.get("level", "")).lower()
|
|
if level in _NONBLOCKING_LEVELS:
|
|
continue
|
|
path = entry.get("path", "")
|
|
norm = _normalize_repo_path(path)
|
|
suffix = _file_suffix(norm)
|
|
# Path relative to the load root: a file under a load subdir (e.g. LLM/) is
|
|
# root-level there, not nested.
|
|
load_rel = _load_relative_path(norm, load_subdirs)
|
|
if not norm or suffix in _INERT_SUFFIXES or suffix in _SOURCE_SUFFIXES:
|
|
# Inert formats cannot execute on load; source code is the consent gate's
|
|
# domain (auto_map), not a deserialization vector.
|
|
skipped.append({"path": path, "level": level})
|
|
elif "/" not in load_rel:
|
|
unsafe.append({"path": path, "level": level}) # root pickle -> load vector
|
|
else:
|
|
# Subdir pickle: deserialized only if a weight index references it.
|
|
maybe_shard.append({"path": path, "level": level, "norm": norm})
|
|
|
|
if maybe_shard:
|
|
indexed = _indexed_shard_paths(
|
|
model_name, hf_token, load_subdirs, revision = snapshot_revision
|
|
)
|
|
for m in maybe_shard:
|
|
# Block if a root index lists this shard, or if the lookup was inconclusive
|
|
# (transient error -> stay conservative). A definitive "no index / not listed"
|
|
# stays non-blocking (e.g. NeMo nemo/*.distcp).
|
|
if indexed is None or m["norm"] in indexed:
|
|
unsafe.append({"path": m["path"], "level": m["level"]})
|
|
else:
|
|
skipped.append({"path": m["path"], "level": m["level"]})
|
|
|
|
if not unsafe:
|
|
if skipped:
|
|
# Flagged files exist, but none the load deserializes (subdir pickle or inert
|
|
# format) -> allow, but log them so they stay visible.
|
|
logger.info(
|
|
"'%s': Hugging Face flagged files, but none are a load-path RCE "
|
|
"vector (subdir/inert); allowing the load. Flagged: %s",
|
|
model_name,
|
|
", ".join(f"{s['path']}({s['level']})" for s in skipped),
|
|
)
|
|
return FileSecurityDecision(model_name, False, reason = "no unsafe files in the load path")
|
|
|
|
names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files"
|
|
logger.warning(
|
|
"Blocking load of '%s': Hugging Face security scan flagged unsafe files (%s).",
|
|
model_name,
|
|
names,
|
|
)
|
|
return FileSecurityDecision(
|
|
model_name,
|
|
True,
|
|
unsafe_files = unsafe,
|
|
reason = f"Hugging Face security scan flagged unsafe files: {names}",
|
|
)
|