diff --git a/scripts/mlx_image_server.py b/scripts/mlx_image_server.py index 8bc27fdaf..f955299a1 100644 --- a/scripts/mlx_image_server.py +++ b/scripts/mlx_image_server.py @@ -327,7 +327,12 @@ def list_models(): @app.post("/v1/images/generations") def generate(req: ImageRequest): - model = req.model or _args.model + # The served model is the one this process was launched with. `req.model` + # is accepted for OpenAI wire compatibility and ignored, matching + # scripts/diffusion_server.py: honouring it would let a caller point the + # generator at any local directory or Hugging Face repo, and the HiDream + # branch runs a python script from inside that directory. + model = _args.model width, height = _size(req.size) out_images = [] count = max(1, min(int(req.n or 1), 4)) @@ -393,7 +398,7 @@ async def edit_image( size: str = Form("1024x1024"), response_format: str = Form("b64_json"), ): - active_model = model or _args.model + active_model = _args.model # pinned; see generate() if _is_lama_inpaint(active_model) or _is_ddcolor(active_model): image_raw = await image.read() mask_raw = await mask.read() if mask is not None else None diff --git a/tests/test_mlx_image_server_security.py b/tests/test_mlx_image_server_security.py new file mode 100755 index 000000000..f87ed5e55 --- /dev/null +++ b/tests/test_mlx_image_server_security.py @@ -0,0 +1,140 @@ +"""Pin the mlx_image_server caller-chosen-model + DNS-rebinding regressions. + +Background: scripts/mlx_image_server.py used to resolve the model per request +(``req.model or _args.model``) instead of serving the model the process was +launched with. ``_is_hidream`` is a substring test and ``_snapshot_path`` +accepts either a local directory or a Hugging Face repo id, so a caller could +name any directory / repo and the HiDream branch would then run +``/scripts/hidream_o1/generate_hidream_o1_mlx.py`` under +``sys.executable``. The server has no auth, and the cookbook binds it to +``0.0.0.0`` whenever it is serving to a remote host, so that was reachable +code execution. + +The fix pins both request paths to ``_args.model``, matching +scripts/diffusion_server.py. +""" + +import argparse +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_SCRIPT = Path(__file__).resolve().parent.parent / "scripts" / "mlx_image_server.py" + +_BASE_URL = "http://127.0.0.1" + + +def _load_module(): + """Fresh import of the server module. Unlike diffusion_server it pulls in no + heavy runtime (mlx / torch imports all live inside the request handlers), so + the real module is imported rather than AST-extracted.""" + spec = importlib.util.spec_from_file_location("mlx_image_server_under_test", _SCRIPT) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def server(monkeypatch): + """Server module launched with a pinned, non-HiDream model.""" + module = _load_module() + module._args = argparse.Namespace( + model="mlx-community/pinned-model", + host="127.0.0.1", + port=8100, + steps=0, + width=512, + height=512, + base_model="", + lora_style="", + lora_paths=[], + lora_scales=[], + vlm_model="", + ) + return module + + +def _client(module): + from fastapi.testclient import TestClient + + return TestClient(module.app, base_url=_BASE_URL) + + +def _plant_hidream_model_dir(tmp_path: Path) -> tuple[Path, Path]: + """A directory that satisfies _is_hidream() and carries the script the + HiDream branch executes. The script writes a marker so the test can tell + whether it ran.""" + model_dir = tmp_path / "hidream-planted" + generator = model_dir / "scripts" / "hidream_o1" + generator.mkdir(parents=True) + marker = model_dir / "executed.txt" + (generator / "generate_hidream_o1_mlx.py").write_text( + f"open({str(marker)!r}, 'w').write('ran')\n", encoding="utf-8" + ) + return model_dir, marker + + +def test_generate_does_not_run_code_from_a_caller_named_model_dir(server, tmp_path): + """The regression: naming a local directory as the model must not execute + the generator script inside it.""" + model_dir, marker = _plant_hidream_model_dir(tmp_path) + + _client(server).post( + "/v1/images/generations", + json={"model": str(model_dir), "prompt": "x", "size": "64x64"}, + ) + + assert not marker.exists(), ( + "code inside the caller-named model directory ran; the request model " + "must not select the generator" + ) + + +def test_generate_does_not_fetch_a_caller_named_repo(server, monkeypatch): + """The remote half of the same defect: the caller's string must never reach + the Hugging Face downloader.""" + downloaded = [] + stub = types.ModuleType("huggingface_hub") + stub.snapshot_download = lambda repo: downloaded.append(repo) + monkeypatch.setitem(sys.modules, "huggingface_hub", stub) + + _client(server).post( + "/v1/images/generations", + json={"model": "attacker-account/hidream-anything", "prompt": "x"}, + ) + + assert "attacker-account/hidream-anything" not in downloaded + + +def test_edits_does_not_run_code_from_a_caller_named_model_dir(server, tmp_path): + """/v1/images/edits resolved the model the same way and must be pinned too. + A "lama" name reaches the inpaint bridge, so the caller's model string is + what picks the branch here.""" + model_dir = tmp_path / "lama-planted" + model_dir.mkdir() + called = [] + server._run_inpaint_bridge = lambda *a, **kw: called.append(a) + server._run_ddcolor_bridge = lambda *a, **kw: called.append(a) + + resp = _client(server).post( + "/v1/images/edits", + data={"model": str(model_dir), "prompt": "x"}, + files={"image": ("i.png", b"not-a-real-png", "image/png")}, + ) + + assert not called, "caller-supplied model selected the edit branch" + assert resp.status_code == 422, "pinned non-edit model should be refused" + + +def test_pinned_hidream_model_is_still_served(server, tmp_path, monkeypatch): + """Behaviour preservation: pinning must not break a server that was actually + launched with a HiDream model.""" + model_dir, marker = _plant_hidream_model_dir(tmp_path) + server._args.model = str(model_dir) + + _client(server).post("/v1/images/generations", json={"model": "", "prompt": "x"}) + + assert marker.exists(), "the model this server was launched with must still run"