From 4c8137a7af016ea451d4cec81cb2958de3f619fd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 28 Jun 2026 05:51:34 +0000 Subject: [PATCH] Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. --- .../core/inference/diffusion_prequant.py | 40 ++++++++++- studio/backend/models/inference.py | 4 +- .../backend/tests/test_diffusion_prequant.py | 69 +++++++++++++++++++ 3 files changed, 109 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 064647fae3..3aecc65a6d 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -33,6 +33,26 @@ from typing import Any, Optional # on-disk structure changes so an old/foreign artifact is rejected rather than mis-loaded. PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1" +# Loading a checkpoint ends in ``torch.load(weights_only=False)``, which executes arbitrary +# code embedded in the pickle. A hosted family *repo* checkpoint is first-party and trusted, +# but a ``source.kind == "path"`` can originate from the ``transformer_prequant_path`` field +# of a load request -- i.e. an authenticated API caller naming an arbitrary local file. +# Unpickling that is remote code execution, so the local-path branch is refused unless an +# operator explicitly opts in via this env var. The trusted hosted-repo path is unaffected. +ALLOW_LOCAL_PREQUANT_PATH_ENV = "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH" + + +def _local_prequant_path_allowed() -> bool: + """Whether a request-supplied local pre-quant *path* may be unpickled (operator opt-in).""" + import os + + return (os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV) or "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + @dataclass(frozen = True) class PrequantSource: @@ -93,6 +113,21 @@ def load_prequantized_transformer( ordinary unavailable artifact. """ try: + # weights_only=False (required below) executes pickle code, so a caller-supplied + # local path is refused unless an operator opted in. The hosted family repo is + # first-party and always allowed. + if source.kind == "path" and not _local_prequant_path_allowed(): + _warn( + logger, + f"{scheme}:path", + RuntimeError( + "request-supplied local pre-quant path refused (unpickling an " + f"arbitrary file is unsafe); set {ALLOW_LOCAL_PREQUANT_PATH_ENV}=1 " + "to allow trusted local checkpoints", + ), + ) + return None + path = _resolve_checkpoint_path(source, hf_token) if path is None: return None @@ -100,9 +135,8 @@ def load_prequantized_transformer( import torch # torchao weight subclasses are not safetensors-serializable, so the checkpoint is - # a torch.save pickle. weights_only=False is required to rebuild those subclasses; - # only a configured family repo (first-party) or an explicit local path reaches - # here, which is the trust signal -- this never loads an arbitrary remote pickle. + # a torch.save pickle. weights_only=False is required to rebuild those subclasses. + # The local-path branch is gated above; the repo branch is a first-party artifact. ckpt = torch.load(path, weights_only = False, map_location = "cpu") if not _validate_checkpoint(ckpt, scheme, base, logger): return None diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 928df30358..237b9e90f7 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1744,7 +1744,9 @@ class DiffusionLoadRequest(BaseModel): "scripts/build_prequant_checkpoint.py) for the requested transformer_quant " "scheme. Loads the already-quantized weights with the dense bf16 never on the " "GPU (~half the load VRAM and a smaller download). null uses the family's hosted " - "checkpoint if configured, else quantises the dense transformer at load time.", + "checkpoint if configured, else quantises the dense transformer at load time. " + "Loading a local path unpickles the file (arbitrary code execution), so it is " + "ignored unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1.", ) diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index ac5aa4ca91..0a6a6dc67f 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -143,9 +143,16 @@ def _load( scheme = "fp8", load_raises = False, exists = True, + allow_local = True, ): _FakeTransformer.calls = {} _stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises) + # The local-path branch is opt-in (it unpickles an arbitrary file); these tests + # exercise the load mechanics, so enable it unless a test is checking the gate. + if allow_local: + monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, "1") + else: + monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) path = tmp_path / "ckpt.pt" if exists: path.write_bytes(b"x") @@ -195,3 +202,65 @@ def test_load_scheme_mismatch_is_none(monkeypatch, tmp_path): def test_load_base_mismatch_is_none(monkeypatch, tmp_path): assert _load(monkeypatch, tmp_path, _good_ckpt(base = "other/model")) is None + + +# ── local-path opt-in gate (RCE guard) ─────────────────────────────────────────── +def test_load_local_path_refused_by_default(monkeypatch, tmp_path): + # A valid checkpoint at a real file is still refused: torch.load must never run on a + # request-supplied path without the operator opt-in. + called = {"load": False} + + def _explode(*a, **k): + called["load"] = True + raise AssertionError("torch.load must not run on a refused local path") + + torch = types.ModuleType("torch") + torch.load = _explode + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) + + path = tmp_path / "ckpt.pt" + path.write_bytes(b"x") + source = PrequantSource(kind = "path", location = str(path), filename = None) + result = load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + logger = None, + ) + assert result is None + assert called["load"] is False + + +def test_load_local_path_allowed_with_optin(monkeypatch, tmp_path): + assert _load(monkeypatch, tmp_path, _good_ckpt(), allow_local = True) is not None + + +def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path): + # The hosted-repo branch is first-party and trusted: it loads with no opt-in env set. + _FakeTransformer.calls = {} + _stub_torch_accelerate(monkeypatch, _good_ckpt()) + monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) + + downloaded = tmp_path / "transformer_fp8.pt" + downloaded.write_bytes(b"x") + hub = types.ModuleType("huggingface_hub") + hub.hf_hub_download = lambda repo_id, filename, token = None: str(downloaded) + monkeypatch.setitem(sys.modules, "huggingface_hub", hub) + + source = PrequantSource(kind = "repo", location = "org/hosted-fp8", filename = "transformer_fp8.pt") + result = load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + logger = None, + ) + assert result is not None