diff --git a/studio/backend/tests/test_mlx_repair.py b/studio/backend/tests/test_mlx_repair.py index 7e60c9c3b9..bb8dfce049 100644 --- a/studio/backend/tests/test_mlx_repair.py +++ b/studio/backend/tests/test_mlx_repair.py @@ -119,6 +119,68 @@ def test_repair_install_pins_transformers_and_cleans_up(monkeypatch): assert env.get("UV_OVERRIDE", "").endswith("overrides-darwin-arm64.txt") +def test_install_requires_prebuilt_wheels(monkeypatch): + # A source distribution's PEP 517 build backend runs arbitrary code at install + # time, before the post-install stack check. The unattended self-heal must + # require pre-built wheels so a malicious resolver-selected sdist cannot execute + # during ordinary Studio startup. mlx/mlx-metal ship wheels only and + # mlx-lm/mlx-vlm publish py3-none-any wheels, so a healthy self-heal still works. + pytest.importorskip("transformers") + captured = {} + + class _Result: + returncode = 0 + stdout = "" + + monkeypatch.setattr(mr, "_uv_executable", lambda: "/usr/bin/uv") + monkeypatch.setattr( + mr.subprocess, "run", lambda cmd, **k: captured.update(cmd = cmd) or _Result() + ) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: True) + + assert mr.attempt_mlx_repair() is True + assert mr._ONLY_BINARY_ARG in captured["cmd"] + + +def test_install_env_drops_secrets_and_source_redirects(monkeypatch): + # The unattended self-heal must not hand resolver/build code the full Studio + # environment: secrets and package-source redirects are dropped, while the + # variables uv genuinely needs are forwarded. + monkeypatch.setenv("HF_TOKEN", "secret-hf") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret-aws") + monkeypatch.setenv("WANDB_API_KEY", "secret-wandb") + monkeypatch.setenv("UV_FIND_LINKS", "/tmp/evil") + monkeypatch.setenv("UV_DEFAULT_INDEX", "file:///tmp/evil-index") + monkeypatch.setenv("UV_INDEX_URL", "https://evil.example/simple") + monkeypatch.setenv("PIP_INDEX_URL", "https://evil.example/simple") + monkeypatch.setenv("UV_CACHE_DIR", "/tmp/evil-cache") + monkeypatch.setenv("XDG_CACHE_HOME", "/tmp/evil-xdg-cache") + monkeypatch.setenv("PATH", "/usr/bin:/bin") + monkeypatch.setenv("HOME", "/home/studio") + + env = mr._mlx_install_env() + + # Secrets never reach a (potentially malicious) build/install hook. + for secret in ("HF_TOKEN", "AWS_SECRET_ACCESS_KEY", "WANDB_API_KEY"): + assert secret not in env + # A poisoned process env cannot repoint the install at a hostile source or + # an attacker-staged cache (cache poisoning / symlink writes). + for redirect in ( + "UV_FIND_LINKS", + "UV_DEFAULT_INDEX", + "UV_INDEX_URL", + "PIP_INDEX_URL", + "UV_CACHE_DIR", + "XDG_CACHE_HOME", + ): + assert redirect not in env + # What uv genuinely needs is still forwarded. + assert env["PATH"] == "/usr/bin:/bin" + assert env["HOME"] == "/home/studio" + # UV_OVERRIDE is set by us (not inherited), so a poisoned one is ignored. + assert env.get("UV_OVERRIDE", "").endswith("overrides-darwin-arm64.txt") + + def test_repair_rejects_inadequate_stack(monkeypatch): # A successful uv run that still leaves an old/missing mlx-vlm must NOT clear # chat-only: attempt_mlx_repair returns False so Train/Export stay disabled. diff --git a/studio/backend/utils/mlx_repair.py b/studio/backend/utils/mlx_repair.py index 06d0289166..7e980deb82 100644 --- a/studio/backend/utils/mlx_repair.py +++ b/studio/backend/utils/mlx_repair.py @@ -49,6 +49,56 @@ MLX_PACKAGES = tuple(f"{name}>={version}" for name, version in _MLX_MIN_VERSIONS _MLX_REINSTALL_ARGS = tuple( arg for name in _MLX_PACKAGE_NAMES for arg in ("--reinstall-package", name) ) +# Require pre-built wheels for the unattended self-heal. A source distribution's +# PEP 517 build backend runs arbitrary code at install time, and this install is +# default-on, resolver-driven, and runs before the post-install stack check can +# reject anything. mlx/mlx-metal ship wheels only (no sdist on PyPI) and +# mlx-lm/mlx-vlm publish py3-none-any wheels, so requiring wheels does not break a +# healthy self-heal; if a wheel is genuinely unavailable the install fails and +# Studio stays chat-only (the existing safe fallback) until `unsloth studio update`. +_ONLY_BINARY_ARG = "--only-binary=:all:" +# Allowlist of environment variables forwarded to the install subprocess. The +# self-heal runs without confirmation on the default startup path, so it must not +# hand resolver/build code the full Studio environment. Everything outside this +# set is dropped, which excludes three dangerous classes by construction: +# * secrets (HF_TOKEN, AWS_*, WANDB_API_KEY, ...) that a malicious wheel/sdist +# build hook would otherwise read straight out of os.environ; +# * package-source redirects (UV_INDEX*, UV_DEFAULT_INDEX, UV_FIND_LINKS, +# PIP_INDEX_URL, ...) so a poisoned process env cannot silently repoint the +# install at an attacker-controlled index/find-links; +# * cache-dir redirects (UV_CACHE_DIR, XDG_CACHE_HOME) so a poisoned env cannot +# point uv at an attacker-staged cache (cache poisoning / symlink writes). uv +# falls back to its safe user-owned default cache, reused across runs anyway. +# uv still honours on-disk config (uv.toml / pip.conf), so a corporate mirror +# configured there keeps working; only process-env redirects are dropped. We set +# UV_OVERRIDE ourselves in _mlx_install_env, so a poisoned one here is ignored. +_MLX_ENV_ALLOWLIST = frozenset( + { + "PATH", + "HOME", + "USER", + "LOGNAME", + "TMPDIR", + "TMP", + "TEMP", + "LANG", + "LC_ALL", + "LC_CTYPE", + # proxies + custom CA bundles so installs behind a corporate gateway work + "HTTP_PROXY", + "HTTPS_PROXY", + "NO_PROXY", + "ALL_PROXY", + "http_proxy", + "https_proxy", + "no_proxy", + "all_proxy", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + "REQUESTS_CA_BUNDLE", + "CURL_CA_BUNDLE", + } +) _REPAIR_TIMEOUT_S = 900 # Attempt at most once per process; success is sticky (mlx then imports and the @@ -134,13 +184,22 @@ def _uv_install_cmd(*args: str) -> list[str] | None: def _mlx_install_env() -> dict[str, str]: - """Environment for the mlx install. Mirror the main installer - (install_python_stack.py) by pointing UV_OVERRIDE at overrides-darwin-arm64.txt, - which relaxes mlx-vlm/mlx-lm's transformers>=5 requirement to >=4.57.6. Without - it, uv keeps the Studio transformers pin only by silently backtracking mlx-vlm - to an old, unsupported version (uv honours UV_OVERRIDE; plain pip ignores it, - so the transformers constraint below is the pip-path safety net).""" - env = dict(os.environ) + """Minimal, allowlisted environment for the unattended mlx install. + + The self-heal runs without confirmation on the default startup path, so it + forwards only the variables uv genuinely needs (see _MLX_ENV_ALLOWLIST) instead + of the full Studio environment: secrets and package-source redirects in + os.environ are dropped so a malicious resolver-selected artifact cannot read + Studio secrets or be steered to a hostile index. + + Mirror the main installer (install_python_stack.py) by pointing UV_OVERRIDE at + overrides-darwin-arm64.txt, which relaxes mlx-vlm/mlx-lm's transformers>=5 + requirement to >=4.57.6. Without it, uv keeps the Studio transformers pin only + by silently backtracking mlx-vlm to an old, unsupported version (uv honours + UV_OVERRIDE; plain pip ignores it, so the transformers constraint below is the + pip-path safety net). We set UV_OVERRIDE ourselves, so a poisoned one in the + process env is ignored.""" + env = {key: os.environ[key] for key in _MLX_ENV_ALLOWLIST if key in os.environ} override = ( Path(__file__).resolve().parents[1] / "requirements" @@ -191,7 +250,13 @@ def attempt_mlx_repair(*, timeout: int = _REPAIR_TIMEOUT_S) -> bool: constraint_path = None try: constraint_args, constraint_path = _transformers_constraint_args() - cmd = _uv_install_cmd("--upgrade", *_MLX_REINSTALL_ARGS, *constraint_args, *MLX_PACKAGES) + cmd = _uv_install_cmd( + "--upgrade", + _ONLY_BINARY_ARG, + *_MLX_REINSTALL_ARGS, + *constraint_args, + *MLX_PACKAGES, + ) if cmd is None: logger.warning( "MLX self-heal requires uv so Studio can apply dependency overrides; "