ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/tests/test_embedding_model_security_gate.py
Daniel Han 414503745e
 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>
2026-07-07 04:30:21 -07:00

365 lines
16 KiB
Python
Raw Blame History

# 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"])
Reference in a new issue View git blame Copy permalink