ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Harden MLX self-heal install against supply-chain code execution (#6599)

Browse source 
* Harden MLX self-heal install against supply-chain execution

The Apple Silicon MLX self-heal runs uv pip install on a daemon thread
during Studio startup, default-on with only an env opt-out, before the
post-install stack check. Two things widened the supply-chain surface:

- it accepted source distributions, whose PEP 517 build backends run
  arbitrary code at install time; and
- it forwarded the full process environment, exposing Studio secrets to
  that code and letting a poisoned env (UV_FIND_LINKS / UV_DEFAULT_INDEX)
  repoint the install at a hostile source.

Require pre-built wheels (--only-binary=:all:) and forward only an
allowlist of variables uv needs (PATH/HOME, proxy + CA settings, cache
dir), setting UV_OVERRIDE ourselves. mlx/mlx-metal ship wheels only and
mlx-lm/mlx-vlm publish py3-none-any wheels, so a healthy self-heal is
unaffected; an unavailable wheel just leaves Studio chat-only as before.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Drop cache-dir env vars from the self-heal allowlist

Address review: a poisoned process env could set UV_CACHE_DIR / XDG_CACHE_HOME
to redirect uv at an attacker-staged cache (cache poisoning, symlink writes),
which partly undercut the index-redirect protection. Drop them from the
allowlist; uv falls back to its safe user-owned default cache, still reused
across runs, so there is no normal-path cost. Test now asserts both are
excluded from the install env.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-23 07:02:14 -07:00 • committed by GitHub
commit 61eef657e8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 135 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

62
studio/backend/tests/test_mlx_repair.py
View file

@ -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.

81
studio/backend/utils/mlx_repair.py
View file

@ -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; "