ab2717afe0* 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>
119 lines
3.5 KiB
Python
119 lines
3.5 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
_backend = os.path.join(os.path.dirname(__file__), "..")
|
|
sys.path.insert(0, _backend)
|
|
|
|
from models.training import TrainingRunUpdateRequest
|
|
from routes import training_history
|
|
|
|
|
|
BASE_RUN = {
|
|
"id": "run-1",
|
|
"status": "stopped",
|
|
"model_name": "unsloth/test-model",
|
|
"dataset_name": "test-dataset",
|
|
"display_name": "Existing name",
|
|
"started_at": "2026-01-01T00:00:00Z",
|
|
"ended_at": "2026-01-01T00:01:00Z",
|
|
"total_steps": 10,
|
|
"final_step": 5,
|
|
"output_dir": "/tmp/run-1",
|
|
"resumed_later": False,
|
|
}
|
|
|
|
|
|
def _patch_run(monkeypatch: pytest.MonkeyPatch, payload: TrainingRunUpdateRequest):
|
|
stored = dict(BASE_RUN)
|
|
calls: list[str | None] = []
|
|
|
|
def fake_get_run(run_id: str):
|
|
assert run_id == "run-1"
|
|
return dict(stored)
|
|
|
|
def fake_update_run_display_name(run_id: str, display_name: str | None):
|
|
assert run_id == "run-1"
|
|
calls.append(display_name)
|
|
stored["display_name"] = display_name
|
|
|
|
monkeypatch.setattr(training_history, "get_run", fake_get_run)
|
|
monkeypatch.setattr(
|
|
training_history,
|
|
"update_run_display_name",
|
|
fake_update_run_display_name,
|
|
)
|
|
monkeypatch.setattr(training_history, "can_resume_run", lambda run: True)
|
|
|
|
result = asyncio.run(
|
|
training_history.update_training_run(
|
|
"run-1",
|
|
payload,
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
return result, calls
|
|
|
|
|
|
def test_update_run_omitted_display_name_is_noop(monkeypatch: pytest.MonkeyPatch):
|
|
result, calls = _patch_run(monkeypatch, TrainingRunUpdateRequest.model_validate({}))
|
|
|
|
assert calls == []
|
|
assert result.display_name == "Existing name"
|
|
assert result.can_resume is True
|
|
|
|
|
|
def test_update_run_explicit_null_clears_display_name(monkeypatch: pytest.MonkeyPatch):
|
|
result, calls = _patch_run(
|
|
monkeypatch,
|
|
TrainingRunUpdateRequest.model_validate({"display_name": None}),
|
|
)
|
|
|
|
assert calls == [None]
|
|
assert result.display_name is None
|
|
assert result.can_resume is True
|
|
|
|
|
|
def test_update_run_whitespace_clears_display_name(monkeypatch: pytest.MonkeyPatch):
|
|
result, calls = _patch_run(
|
|
monkeypatch,
|
|
TrainingRunUpdateRequest.model_validate({"display_name": " "}),
|
|
)
|
|
|
|
assert calls == [None]
|
|
assert result.display_name is None
|
|
|
|
|
|
def test_update_run_rejects_unknown_fields():
|
|
with pytest.raises(ValidationError):
|
|
TrainingRunUpdateRequest.model_validate({"unknown": "value"})
|
|
|
|
|
|
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
|