odysseus/tests/test_mlx_image_server_security.py

"""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
``<model>/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"


@pytest.mark.parametrize("prompt", [
    "ordinary prompt with spaces and 'quotes'",
    "image; touch SHELL_MARKER",
    "image && touch SHELL_MARKER || echo failed",
    "$(touch SHELL_MARKER) `touch SHELL_MARKER`",
    "image\ntouch SHELL_MARKER\n",
    "image > SHELL_MARKER | touch SHELL_MARKER",
    "--output=SHELL_MARKER",
])
def test_hidream_prompt_reaches_real_child_as_one_literal_argument(
    server, tmp_path, monkeypatch, prompt,
):
    import base64
    import json

    model_dir, _ = _plant_hidream_model_dir(tmp_path)
    server._args.model = str(model_dir)
    script = model_dir / "scripts" / "hidream_o1" / "generate_hidream_o1_mlx.py"
    script.write_text(
        "import json, pathlib, sys\n"
        "argv = sys.argv[1:]\n"
        "out = pathlib.Path(argv[argv.index('--output') + 1])\n"
        "out.write_text(json.dumps(argv), encoding='utf-8')\n",
        encoding="utf-8",
    )
    marker = tmp_path / "shell-executed"
    shell_startup = tmp_path / "shell-startup"
    shell_startup.write_text(f"touch {str(marker)!r}\n", encoding="utf-8")
    monkeypatch.setenv("BASH_ENV", str(shell_startup))
    monkeypatch.setenv("ENV", str(shell_startup))
    monkeypatch.setenv("SHELL", "/bin/sh")
    prompt = prompt.replace("SHELL_MARKER", str(marker))

    response = _client(server).post(
        "/v1/images/generations",
        json={"model": "attacker/hidream-other", "prompt": prompt, "size": "64x64"},
    )

    assert response.status_code == 200
    argv = json.loads(base64.b64decode(response.json()["data"][0]["b64_json"]))
    assert argv[argv.index("--prompt") + 1] == prompt
    assert len(argv) == 13
    assert argv[argv.index("--model-path") + 1] == str(model_dir)
    assert not marker.exists()


def test_mflux_prompt_is_one_argv_item_without_shell(server, monkeypatch):
    import subprocess

    prompt = "a picture; $(touch /tmp/should-not-exist)\n--output=/etc/passwd"
    seen = []
    monkeypatch.setattr(server, "_resolve_cli", lambda name: "/operator/bin/mflux-generate")

    def run(cmd, **kwargs):
        seen.append((cmd, kwargs))
        Path(cmd[cmd.index("--output") + 1]).write_bytes(b"image")
        return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")

    monkeypatch.setattr(server.subprocess, "run", run)
    response = _client(server).post(
        "/v1/images/generations", json={"model": "/attacker/flux", "prompt": prompt},
    )
    assert response.status_code == 200
    cmd, kwargs = seen[0]
    assert cmd[0] == "/operator/bin/mflux-generate"
    assert cmd[cmd.index("--model") + 1] == server._args.model
    assert cmd[cmd.index("--prompt") + 1] == prompt
    assert len(cmd) == 13
    assert not kwargs.get("shell", False)
    assert "executable" not in kwargs