mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
Run the malware gate on the RAG embedding model before it loads (#6887)
* Run the malware gate on the RAG embedding model before it loads Setting the RAG embedding model through PUT /api/settings/embedding-model persisted an arbitrary repo and later handed it straight to SentenceTransformer, which deserializes pickle weights. Unlike the normal model-load paths, this route never ran evaluate_file_security, and force skipped verification entirely, so a repo Hugging Face flags as unsafe (or any repo under force) could be downloaded and loaded in the backend process without a scan. Run the malware/pickle scan at both ends: the settings endpoint now scans before persisting and returns 409 on a flagged repo even under force (force still only skips the is-embedding-model type check for offline or local repos), and the embedder scans again at the load sink so a name that arrives via env or default is covered too. Local paths and unreachable scans fail open inside evaluate_file_security, and the sink never bricks the embedder on a gate error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Thread the load token into the embedding scan and hard-fail on a block The load-sink scan ran without a token, so evaluate_file_security (which passes token=False when none is given) could not reach a gated or private repo and failed open for exactly the model SentenceTransformer would still load. Resolve the loader's own token (HF_TOKEN env or the cached login) and pass it to the sink scan, and fall back to it in the settings endpoint when the request omits one. The sink previously raised a plain RuntimeError, which the llama-server fallback in encode() and _build_st_backend_or_fallback() swallowed as a routine ST failure, silently switching backends instead of blocking. Raise a distinct UnsafeEmbeddingModelError that both fallback paths re-raise, so a flagged model hard-fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scan sentence-transformers module dirs and scope the embedding pickle gate to the ST backend Extend the RAG embedding malware gate so a poisoned pickle under a SentenceTransformer module dir (for example 0_Transformer/pytorch_model.bin) blocks. Those dirs are read from the repo's modules.json and passed as load roots to evaluate_file_security at both the settings endpoint and the load sink, so such a pickle is treated as root-level there instead of an unreferenced nested shard that was previously allowed. Scope the ST pickle scan to the sentence-transformers backend. On the llama-server backend the embedder loads GGUF files (inert) from the -GGUF companion repo, never the ST repo's pickle, so a custom ST repo with a flagged pickle and a clean GGUF companion is no longer rejected. The existing GGUF availability checks already cover that path. Return 403 for the hard security block instead of 409. The settings UI routes every 409 into the forceable save-anyway flow, but this block cannot be bypassed by force, so it now uses a distinct status the client treats as non-forceable. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Base the embedding pickle scan on the actual backend, not just the resolver _llama_backend_active only consulted the auto resolver, so on a GPU box where auto resolves to sentence-transformers but the process already fell back to the llama-server backend at runtime (a torch or CUDA load/encode failure), it returned False and the settings endpoint hard-blocked a save whose ST pickle is flagged even though the process loads only inert GGUF. Add active_backend_is_llama, which reflects the actual built backend (True when the cached backend is a LlamaServerBackend, including a runtime fallback) and otherwise defers to the resolver as a fresh process would, and delegate _llama_backend_active to it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Report the cached embedding backend verbatim, not the resolver active_backend_is_llama() fell through to the config resolver whenever a backend was already built but was not llama-server, so a live sentence-transformers backend could report llama=True once the resolver picked llama (GPU heuristic or a runtime config change) and wrongly skip its pickle scan. Once a backend exists, return isinstance(backend, LlamaServerBackend) directly; only defer to the resolver before any backend is built. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
bdb958e052
commit
414503745e
6 changed files with 573 additions and 15 deletions
|
|
@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None:
|
|||
install_torchao_windows_rocm_stub()
|
||||
|
||||
|
||||
class UnsafeEmbeddingModelError(RuntimeError):
|
||||
"""Raised when the embedding model repo is flagged unsafe. A distinct type so the
|
||||
llama-server fallback paths re-raise it instead of masking a security block as a
|
||||
routine ST failure."""
|
||||
|
||||
|
||||
def _ambient_hf_token() -> str | None:
|
||||
"""The HF token the loader itself would use (HF_TOKEN env or the cached login), so
|
||||
the scan can reach a gated/private repo instead of failing open. None if unavailable."""
|
||||
try:
|
||||
from huggingface_hub import get_token
|
||||
return get_token()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
|
||||
"""The module directories a SentenceTransformer load reads weights from, taken from
|
||||
the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``).
|
||||
ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the
|
||||
security scan: a flagged pickle directly under one must block. Returns () on any
|
||||
failure (no modules.json, offline, malformed) so the guard never bricks the embedder.
|
||||
"""
|
||||
try:
|
||||
import json
|
||||
|
||||
from utils.paths import is_local_path
|
||||
|
||||
if is_local_path(name):
|
||||
from pathlib import Path
|
||||
from utils.paths import normalize_path
|
||||
|
||||
path = Path(normalize_path(name)).expanduser() / "modules.json"
|
||||
if not path.is_file():
|
||||
return ()
|
||||
data = json.loads(path.read_text())
|
||||
else:
|
||||
from huggingface_hub import hf_hub_download
|
||||
from huggingface_hub.utils import EntryNotFoundError
|
||||
|
||||
try:
|
||||
local = hf_hub_download(name, "modules.json", token = token or None)
|
||||
except EntryNotFoundError:
|
||||
return ()
|
||||
data = json.loads(open(local).read())
|
||||
subdirs = []
|
||||
for module in data or ():
|
||||
sub = str((module or {}).get("path", "")).strip().strip("/")
|
||||
if sub:
|
||||
subdirs.append(sub)
|
||||
return tuple(dict.fromkeys(subdirs))
|
||||
except Exception:
|
||||
return ()
|
||||
|
||||
|
||||
def _guard_model_security(name: str) -> None:
|
||||
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
|
||||
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
|
||||
/settings gate (a name can also arrive via env/default); local paths and unreachable
|
||||
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
|
||||
"""
|
||||
try:
|
||||
from utils.security import evaluate_file_security, security_load_subdirs
|
||||
|
||||
token = _ambient_hf_token()
|
||||
# Union the audio-model load roots with the ST module dirs so a flagged pickle
|
||||
# directly under a Transformer module dir (0_Transformer/) blocks instead of
|
||||
# passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
|
||||
)
|
||||
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
|
||||
except Exception:
|
||||
return
|
||||
if blocked:
|
||||
raise UnsafeEmbeddingModelError(
|
||||
f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
|
||||
"scan; refusing to load. Set a different RAG embedding model."
|
||||
)
|
||||
|
||||
|
||||
def _get(model_name: str | None = None):
|
||||
"""Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16
|
||||
for a ~1.5x speedup at negligible accuracy loss."""
|
||||
|
|
@ -75,6 +156,7 @@ def _get(model_name: str | None = None):
|
|||
|
||||
device = _device()
|
||||
logger.info("loading embedding model %s on %s", name, device)
|
||||
_guard_model_security(name)
|
||||
_model = SentenceTransformer(
|
||||
name, device = device, model_kwargs = {"torch_dtype": "float16"}
|
||||
)
|
||||
|
|
@ -159,6 +241,8 @@ class _SentenceTransformersBackend:
|
|||
):
|
||||
try:
|
||||
return _st_encode(texts, model_name = model_name, normalize = normalize)
|
||||
except UnsafeEmbeddingModelError:
|
||||
raise # a security block must hard-fail, not fall back to llama-server
|
||||
except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure
|
||||
# ST loaded but this encode blew up; swap the process to the llama-server
|
||||
# embedder (so later encodes stay in one space) and retry.
|
||||
|
|
@ -222,6 +306,8 @@ def _build_st_backend_or_fallback():
|
|||
try:
|
||||
backend.warm(model_name = None)
|
||||
return backend
|
||||
except UnsafeEmbeddingModelError:
|
||||
raise # a security block must hard-fail, not fall back to llama-server
|
||||
except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure
|
||||
fallback = _try_make_llama_backend()
|
||||
if fallback is None:
|
||||
|
|
@ -290,6 +376,37 @@ def _reset_backend() -> None:
|
|||
_backend_key = None
|
||||
|
||||
|
||||
def active_backend_is_llama() -> bool:
|
||||
"""True when this process actually embeds via the llama-server (GGUF) backend.
|
||||
|
||||
Reflects the ACTUAL built backend once one exists: an ``auto`` install that
|
||||
resolves to sentence-transformers but then falls back to llama-server at
|
||||
runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or
|
||||
``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so
|
||||
callers gating on the ST pickle must see llama here. Before any backend is
|
||||
built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw
|
||||
key) exactly as a fresh process would. Never raises: a backend probe must not
|
||||
block saving a model."""
|
||||
try:
|
||||
with _backend_lock:
|
||||
backend = _backend
|
||||
if backend is not None:
|
||||
# A backend exists: report what it ACTUALLY is. A concrete
|
||||
# sentence-transformers backend must return False even if the
|
||||
# resolver would now pick llama, so its pickle stays gated. If the
|
||||
# llama import fails we cannot be llama, so fall to the safe False.
|
||||
try:
|
||||
from .embed_llama_server import LlamaServerBackend
|
||||
except Exception: # noqa: BLE001 - llama plumbing import must never block
|
||||
return False
|
||||
return isinstance(backend, LlamaServerBackend)
|
||||
raw = (config.EMBED_BACKEND or "auto").strip().lower()
|
||||
key = _resolve_auto() if raw in _AUTO_ALIASES else raw
|
||||
return key in _LLAMA_ALIASES
|
||||
except Exception: # noqa: BLE001 - a backend probe must never block saving
|
||||
return False
|
||||
|
||||
|
||||
def warm(model_name: str | None = None) -> None:
|
||||
"""Eagerly load the embedder so the first real request isn't slow."""
|
||||
_get_backend().warm(model_name = model_name)
|
||||
|
|
|
|||
|
|
@ -260,17 +260,29 @@ def _embedding_model_response() -> EmbeddingModelResponse:
|
|||
)
|
||||
|
||||
|
||||
def _llama_backend_active() -> bool:
|
||||
"""True when this install embeds via the llama-server (GGUF) backend."""
|
||||
from core.rag import config as rag_config
|
||||
from core.rag import embeddings
|
||||
|
||||
def _ambient_hf_token() -> Optional[str]:
|
||||
"""The HF token the loader would use (HF_TOKEN env or the cached login), so a gated
|
||||
repo is scanned rather than failing open. None if unavailable."""
|
||||
try:
|
||||
raw = (rag_config.EMBED_BACKEND or "auto").strip().lower()
|
||||
key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw
|
||||
from huggingface_hub import get_token
|
||||
return get_token()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _llama_backend_active() -> bool:
|
||||
"""True when this install actually embeds via the llama-server (GGUF) backend.
|
||||
|
||||
Delegates to the embeddings module so a runtime fallback from
|
||||
sentence-transformers to llama-server (after a torch/CUDA load or encode
|
||||
failure) is honored: in that state the process loads only inert GGUF, so the
|
||||
ST pickle gate below must not hard-block a repo whose GGUF companion is clean.
|
||||
Before any backend is built this still reflects the resolver."""
|
||||
from core.rag import embeddings
|
||||
try:
|
||||
return embeddings.active_backend_is_llama()
|
||||
except Exception: # noqa: BLE001 - backend probe must never block saving
|
||||
return False
|
||||
return key in embeddings._LLAMA_ALIASES
|
||||
|
||||
|
||||
def _resolves_as_local_gguf(model: str) -> bool:
|
||||
|
|
@ -357,6 +369,8 @@ def update_embedding_model(
|
|||
"""Set the RAG embedding model. Unless ``force`` is set, the repo is verified
|
||||
to be an embedding model via HF metadata; an unverifiable model (wrong type,
|
||||
typo, gated repo, or no network) returns 409 so the UI can offer "save anyway".
|
||||
A repo flagged unsafe by HF's security scan returns 403 instead: a hard block
|
||||
that ``force`` cannot bypass, so the UI must not offer "save anyway".
|
||||
Documents indexed under the previous model must be re-uploaded."""
|
||||
from utils.models import is_embedding_model
|
||||
|
||||
|
|
@ -370,15 +384,51 @@ def update_embedding_model(
|
|||
event = "settings.update_embedding_model_failed",
|
||||
log = logger,
|
||||
) from exc
|
||||
hf_token = (payload.hf_token or "").strip() or None
|
||||
# The env/default model needs no verification; saving it is a no-op override.
|
||||
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
|
||||
# what the backend loads, and HF metadata cannot verify a local path.
|
||||
if (
|
||||
model != default_embedding_model()
|
||||
and not payload.force
|
||||
and not (_llama_backend_active() and _resolves_as_local_gguf(model))
|
||||
):
|
||||
hf_token = (payload.hf_token or "").strip() or None
|
||||
is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model)
|
||||
# The pickle gate only matters for the sentence-transformers backend, which is what
|
||||
# deserializes pickles. On the llama-server backend the embedder loads GGUF files
|
||||
# (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would
|
||||
# wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability
|
||||
# checks below cover that path instead.
|
||||
scan_st_pickle = (
|
||||
model != default_embedding_model() and not is_local_gguf and not _llama_backend_active()
|
||||
)
|
||||
if scan_st_pickle:
|
||||
# Malware/pickle gate before we persist a repo the embedder later loads with
|
||||
# SentenceTransformer. Runs even under force (force only skips the is-embedding
|
||||
# type check for offline/local repos HF cannot verify); local paths and
|
||||
# unreachable scans fail open inside evaluate_file_security.
|
||||
from utils.security import evaluate_file_security, security_load_subdirs
|
||||
from core.rag.embeddings import _st_module_subdirs
|
||||
|
||||
# Fall back to the loader's own token so a gated/private repo is actually scanned
|
||||
# (a token-less scan fails open for exactly the repo that would still load).
|
||||
scan_token = hf_token or _ambient_hf_token()
|
||||
# Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
|
||||
# one blocks instead of passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys(
|
||||
(
|
||||
*security_load_subdirs(model, scan_token),
|
||||
*_st_module_subdirs(model, scan_token),
|
||||
)
|
||||
)
|
||||
)
|
||||
if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
|
||||
# 403, not 409: the client routes every 409 into the forceable "save anyway"
|
||||
# flow, but this block is a hard, non-forceable security refusal.
|
||||
raise HTTPException(
|
||||
status_code = 403,
|
||||
detail = (
|
||||
f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
|
||||
"cannot be used as the embedding model."
|
||||
),
|
||||
)
|
||||
if model != default_embedding_model() and not payload.force and not is_local_gguf:
|
||||
from core.rag import config as rag_config
|
||||
|
||||
# A GGUF-named repo on the llama-server backend is loaded from its .gguf
|
||||
|
|
|
|||
365
studio/backend/tests/test_embedding_model_security_gate.py
Normal file
365
studio/backend/tests/test_embedding_model_security_gate.py
Normal file
|
|
@ -0,0 +1,365 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""The RAG embedding model must pass the malware/pickle gate before it is persisted or
|
||||
loaded. A flagged repo (or any repo saved with force) previously reached
|
||||
SentenceTransformer unscanned, bypassing the normal model-load protections."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types as _types
|
||||
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.settings as settings
|
||||
|
||||
|
||||
class _Decision:
|
||||
def __init__(self, blocked):
|
||||
self.blocked = blocked
|
||||
|
||||
|
||||
def _security_stub(blocked):
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.evaluate_file_security = lambda *a, **k: _Decision(blocked)
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
return mod
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch):
|
||||
# The settings scan unions in the ST module dirs read from modules.json; keep it
|
||||
# offline and deterministic for the endpoint tests that use this fixture.
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
return TestClient(app, raise_server_exceptions = False), saved
|
||||
|
||||
|
||||
def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch):
|
||||
c, saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
r = c.put(
|
||||
"/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True}
|
||||
)
|
||||
# 403, not the forceable 409, so the client does not offer "save anyway".
|
||||
assert r.status_code == 403
|
||||
assert "model" not in saved # force must not persist a flagged repo
|
||||
|
||||
|
||||
def test_flagged_repo_is_blocked_without_force(client, monkeypatch):
|
||||
c, saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
|
||||
assert r.status_code == 403
|
||||
assert "model" not in saved
|
||||
|
||||
|
||||
def test_hard_block_uses_non_forceable_status(client, monkeypatch):
|
||||
# The forceable verification path uses 409; the hard security block must be distinct
|
||||
# (403) so the frontend never routes it into the "save anyway" force flow.
|
||||
c, _saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
|
||||
assert blocked.status_code == 403
|
||||
|
||||
# A verification failure (not-an-embedding-model) stays forceable at 409.
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False)
|
||||
import utils.models as _models
|
||||
|
||||
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
|
||||
unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"})
|
||||
assert unverified.status_code == 409
|
||||
|
||||
|
||||
def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
|
||||
# On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
|
||||
# pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
# force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped.
|
||||
called = {"scanned": False}
|
||||
mod = _types.ModuleType("utils.security")
|
||||
|
||||
def _fail(*a, **k):
|
||||
called["scanned"] = True
|
||||
return _Decision(True)
|
||||
|
||||
mod.evaluate_file_security = _fail
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.put(
|
||||
"/embedding-model",
|
||||
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert called["scanned"] is False # the ST pickle scan never ran on the llama path
|
||||
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
|
||||
|
||||
|
||||
def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch):
|
||||
# auto resolves to sentence-transformers (GPU present) but the embedder fell back to
|
||||
# llama-server at runtime (torch/CUDA load or encode failure), so the process now loads
|
||||
# only inert GGUF. The real _llama_backend_active() must reflect that cached fallback,
|
||||
# so a flagged ST repo with a clean GGUF companion must not be hard-blocked here.
|
||||
import core.rag.embeddings as embeddings
|
||||
from core.rag.embed_llama_server import LlamaServerBackend
|
||||
|
||||
# Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even
|
||||
# though the auto resolver would still say sentence-transformers.
|
||||
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
|
||||
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
|
||||
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
# Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the
|
||||
# real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored.
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
|
||||
called = {"scanned": False}
|
||||
mod = _types.ModuleType("utils.security")
|
||||
|
||||
def _fail(*a, **k):
|
||||
called["scanned"] = True
|
||||
return _Decision(True)
|
||||
|
||||
mod.evaluate_file_security = _fail
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.put(
|
||||
"/embedding-model",
|
||||
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback
|
||||
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
|
||||
|
||||
|
||||
def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch):
|
||||
# active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers
|
||||
# to the resolver (fresh-process behavior) when none has been built yet.
|
||||
import core.rag.embeddings as embeddings
|
||||
import core.rag.config as rag_config
|
||||
from core.rag.embed_llama_server import LlamaServerBackend
|
||||
|
||||
# A cached llama backend wins even when auto would resolve to sentence-transformers.
|
||||
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto")
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
|
||||
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
|
||||
assert embeddings.active_backend_is_llama() is True
|
||||
|
||||
# A cached ST backend reports False even when the resolver now picks llama, so its
|
||||
# pickle stays gated (the cached backend, not the resolver, is what actually embeds).
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
|
||||
monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend())
|
||||
assert embeddings.active_backend_is_llama() is False
|
||||
|
||||
# No cached backend -> the resolver decides, unchanged from before.
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
|
||||
monkeypatch.setattr(embeddings, "_backend", None)
|
||||
assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers
|
||||
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
|
||||
assert embeddings.active_backend_is_llama() is True # auto -> llama-server
|
||||
|
||||
# An explicit (non-auto) key is honored verbatim without a cached backend.
|
||||
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server")
|
||||
assert embeddings.active_backend_is_llama() is True
|
||||
|
||||
|
||||
def test_settings_scan_scopes_module_subdirs(monkeypatch):
|
||||
# The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a
|
||||
# pickle directly under one blocks; assert those subdirs reach evaluate_file_security.
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(
|
||||
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
|
||||
)
|
||||
seen = {}
|
||||
|
||||
def _capture(*a, **k):
|
||||
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
|
||||
return _Decision(False)
|
||||
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
mod.evaluate_file_security = _capture
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.put(
|
||||
"/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True}
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert "0_Transformer" in seen["subdirs"]
|
||||
|
||||
|
||||
def test_clean_repo_saves_under_force(client, monkeypatch):
|
||||
c, saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True})
|
||||
assert r.status_code == 200
|
||||
assert saved.get("model") == "acme/clean-embed"
|
||||
|
||||
|
||||
def test_load_sink_refuses_flagged_model(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
import core.rag.embeddings as embeddings
|
||||
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
|
||||
embeddings._guard_model_security("attacker/malicious-embed")
|
||||
|
||||
|
||||
def test_load_sink_allows_clean_model(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
import core.rag.embeddings as embeddings
|
||||
embeddings._guard_model_security("acme/clean-embed") # no raise
|
||||
|
||||
|
||||
def test_sink_threads_ambient_token_into_scan(monkeypatch):
|
||||
# A gated repo set via env/default has no request token; the guard must feed the
|
||||
# loader's own token to the scan, or it fails open for the repo that still loads.
|
||||
seen = {}
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.security_load_subdirs = (
|
||||
lambda name, token = None: seen.setdefault("subdirs_token", token) or ()
|
||||
)
|
||||
mod.evaluate_file_security = lambda *a, **k: seen.setdefault(
|
||||
"scan_token", k.get("hf_token")
|
||||
) or _Decision(False)
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient")
|
||||
embeddings._guard_model_security("acme/gated-embed")
|
||||
assert seen["scan_token"] == "hf_ambient"
|
||||
assert seen["subdirs_token"] == "hf_ambient"
|
||||
|
||||
|
||||
def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch):
|
||||
# A flagged pickle directly under a Transformer module dir (0_Transformer/) must
|
||||
# reach the scan as a load root; assert the guard unions the module dirs into
|
||||
# load_subdirs so evaluate_file_security treats such a pickle as root-level.
|
||||
seen = {}
|
||||
|
||||
def _capture(*a, **k):
|
||||
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
|
||||
return _Decision(False)
|
||||
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.security_load_subdirs = lambda name, token = None: ()
|
||||
mod.evaluate_file_security = _capture
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
|
||||
)
|
||||
embeddings._guard_model_security("acme/embed-with-module-dir")
|
||||
assert "0_Transformer" in seen["subdirs"]
|
||||
|
||||
|
||||
def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch):
|
||||
# The helper must parse each module's non-empty "path" from a local repo's
|
||||
# modules.json and drop the root-level ("") Transformer entry.
|
||||
import json
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
(tmp_path / "modules.json").write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."},
|
||||
{"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."},
|
||||
{"idx": 2, "name": "2", "path": "", "type": "..."},
|
||||
]
|
||||
)
|
||||
)
|
||||
subdirs = embeddings._st_module_subdirs(str(tmp_path), None)
|
||||
assert subdirs == ("0_Transformer", "1_Pooling")
|
||||
|
||||
|
||||
def test_st_module_subdirs_swallows_errors(monkeypatch):
|
||||
# Any failure (no modules.json, offline, malformed) returns () so the guard never
|
||||
# bricks the embedder.
|
||||
import huggingface_hub
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise RuntimeError("offline")
|
||||
|
||||
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom)
|
||||
assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == ()
|
||||
|
||||
|
||||
def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch):
|
||||
# The ST encode fallback must re-raise a security block, not swap to llama-server.
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise embeddings.UnsafeEmbeddingModelError("flagged")
|
||||
|
||||
monkeypatch.setattr(embeddings, "_st_encode", _boom)
|
||||
monkeypatch.setattr(
|
||||
embeddings,
|
||||
"_switch_to_llama_fallback",
|
||||
lambda err: pytest.fail("security block must not fall back to llama-server"),
|
||||
)
|
||||
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
|
||||
embeddings._SentenceTransformersBackend().encode(["hi"])
|
||||
|
|
@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base():
|
|||
if runs_gate and not resolves_base:
|
||||
offenders.append(f"{rel} runs a load gate but never resolves the LoRA base")
|
||||
assert not offenders, "\n".join(offenders)
|
||||
|
||||
|
||||
def test_rag_embedding_path_runs_the_malware_gate():
|
||||
"""The RAG embedding model is set through /settings and later loaded by
|
||||
SentenceTransformer, which deserializes pickles; both sites must run the malware gate
|
||||
or a flagged repo loads unscanned (bypassing the normal model-load protections)."""
|
||||
offenders = []
|
||||
for rel in ("routes/settings.py", "core/rag/embeddings.py"):
|
||||
if "evaluate_file_security(" not in (_BACKEND / rel).read_text():
|
||||
offenders.append(
|
||||
f"{rel} loads/persists an embedding model without evaluate_file_security"
|
||||
)
|
||||
assert not offenders, "\n".join(offenders)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,10 @@ type ApiEmbeddingModelSettings = {
|
|||
* (wrong type, gated repo, or offline). Retry with force to save anyway. */
|
||||
export class EmbeddingModelVerificationError extends Error {}
|
||||
|
||||
/** 403 from the backend: the repo is flagged unsafe by Hugging Face's security scan.
|
||||
* A hard block; force cannot bypass it, so it must not enter the "save anyway" flow. */
|
||||
export class EmbeddingModelBlockedError extends Error {}
|
||||
|
||||
function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings {
|
||||
return {
|
||||
embeddingModel: settings.embedding_model,
|
||||
|
|
@ -56,6 +60,11 @@ export async function updateEmbeddingModelSettings(
|
|||
force: options?.force ?? false,
|
||||
}),
|
||||
});
|
||||
if (res.status === 403) {
|
||||
throw new EmbeddingModelBlockedError(
|
||||
await readFastApiError(res, "This model is blocked by a security scan"),
|
||||
);
|
||||
}
|
||||
if (res.status === 409) {
|
||||
throw new EmbeddingModelVerificationError(
|
||||
await readFastApiError(res, "Could not verify the embedding model"),
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ import {
|
|||
updatePreviewSharing,
|
||||
} from "../api/preview-sharing";
|
||||
import {
|
||||
EmbeddingModelBlockedError,
|
||||
type EmbeddingModelSettings,
|
||||
EmbeddingModelVerificationError,
|
||||
loadEmbeddingModelSettings,
|
||||
|
|
@ -410,7 +411,10 @@ export function GeneralTab() {
|
|||
description: t("settings.general.rag.reindexWarning"),
|
||||
});
|
||||
} catch (error) {
|
||||
if (error instanceof EmbeddingModelVerificationError) {
|
||||
// A hard security block cannot be forced; keep the "save anyway" action hidden.
|
||||
if (error instanceof EmbeddingModelBlockedError) {
|
||||
setEmbeddingModelNeedsForce(false);
|
||||
} else if (error instanceof EmbeddingModelVerificationError) {
|
||||
setEmbeddingModelNeedsForce(true);
|
||||
}
|
||||
setEmbeddingModelError(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue