spark-control/image/tests/test_launch_command.py

"""build_launch_command: argument assembly + the shell-injection invariant.

The security-critical property is that every user-controllable value (repo,
vllm_args, knobs) is shlex-quoted at the sink, so `shlex.split` cleanly reverses
the command back into the exact token list. The vLLM pre-flight validator
(validate.py) depends on this round-trip — these tests lock it in.
"""
import shlex

import pytest
from pydantic import ValidationError

from app.models import Defaults, ModelDef, build_launch_command

DEFAULTS = Defaults(port=8888, host="0.0.0.0")


def _model(**kw) -> ModelDef:
    base = dict(display_name="X", repo="org/name", size_gb=1.0, mode="solo")
    base.update(kw)
    return ModelDef(**base)


def test_solo_model_emits_solo_flag_and_ordered_args():
    cmd = build_launch_command("k", _model(vllm_args=["--max-model-len=1000"]), DEFAULTS)
    assert cmd == (
        "./launch-cluster.sh --solo -d exec vllm serve org/name "
        "--port=8888 --host=0.0.0.0 --max-model-len=1000"
    )


def test_cluster_model_omits_solo_flag():
    cmd = build_launch_command("k", _model(mode="cluster", vllm_args=["-tp=2"]), DEFAULTS)
    assert " --solo " not in cmd
    assert cmd.startswith("./launch-cluster.sh -d exec vllm serve org/name")


def test_knob_overrides_matching_bundled_flag():
    # bundled arg sets max-model-len; the knob must win (single occurrence).
    m = _model(vllm_args=["--max-model-len=1000"], knobs={"max_model_len": 65536})
    cmd = build_launch_command("k", m, DEFAULTS)
    assert "--max-model-len=65536" in cmd
    assert "--max-model-len=1000" not in cmd


def test_repo_with_shell_metacharacters_is_quoted_not_executed():
    # build_launch_command quotes even a hostile repo (validate_repo guards the
    # API boundary; this proves the sink itself is safe in depth).
    evil = "org/name; rm -rf ~ #"
    cmd = build_launch_command("k", _model(repo=evil), DEFAULTS)
    # The raw metacharacters must not appear unquoted...
    assert "; rm -rf" not in cmd.replace(shlex.quote(evil), "")
    # ...and shlex.split must recover the repo as one literal token.
    tokens = shlex.split(cmd)
    assert evil in tokens


def test_command_string_round_trips_through_shlex_split():
    # The invariant validate.py relies on: every arg survives quote -> split intact.
    args = ["--max-model-len=32768", "--load-format=fastsafetensors", "--note=a b c"]
    cmd = build_launch_command("k", _model(vllm_args=args), DEFAULTS)
    tokens = shlex.split(cmd)
    for a in args:
        assert a in tokens


def test_injection_via_vllm_arg_stays_literal():
    payload = "--foo=$(touch /tmp/pwned)"
    cmd = build_launch_command("k", _model(vllm_args=[payload]), DEFAULTS)
    assert payload in shlex.split(cmd)  # preserved as one inert token


# ---- local / fine-tuned models (served by directory, not HF repo) ----

def test_local_model_bind_mounts_dir_and_serves_the_path():
    m = _model(repo="", local_path="/home/u/models/ft-v2", vllm_args=["--max-model-len=2048"])
    cmd = build_launch_command("k", m, DEFAULTS)
    tokens = shlex.split(cmd)
    # The launch script's hook bind-mounts the host dir at the SAME container path.
    assert tokens[0] == (
        "VLLM_SPARK_EXTRA_DOCKER_ARGS=-v /home/u/models/ft-v2:/home/u/models/ft-v2"
    )
    # vLLM is pointed at the directory, not an HF repo id.
    i = tokens.index("serve")
    assert tokens[i + 1] == "/home/u/models/ft-v2"
    assert "--max-model-len=2048" in tokens


def test_local_model_chat_template_arg_survives_round_trip():
    m = _model(
        repo="",
        local_path="/m/ft",
        vllm_args=["--chat-template=/m/ft/chat_template.jinja"],
    )
    cmd = build_launch_command("k", m, DEFAULTS)
    assert "--chat-template=/m/ft/chat_template.jinja" in shlex.split(cmd)


def test_local_path_with_metacharacters_is_quoted_not_executed():
    # The validator rejects a hostile path at the boundary; bypass it with
    # model_construct to prove the quote_arg sink is safe in depth even if a bad
    # value somehow reaches build_launch_command.
    evil = "/m/ft; rm -rf ~"
    m = ModelDef.model_construct(
        display_name="X", repo="", local_path=evil, size_gb=1.0, mode="solo",
        vllm_args=[], knobs=None, custom=False, capabilities=[],
        expected_ready_seconds=300, description=None,
    )
    cmd = build_launch_command("k", m, DEFAULTS)
    tokens = shlex.split(cmd)
    i = tokens.index("serve")
    assert tokens[i + 1] == evil  # recovered as one literal token, not executed
    assert tokens[0] == f"VLLM_SPARK_EXTRA_DOCKER_ARGS=-v {evil}:{evil}"


def test_model_requires_exactly_one_source():
    with pytest.raises(ValidationError):
        ModelDef(display_name="x", size_gb=1, mode="solo")  # neither repo nor local_path
    with pytest.raises(ValidationError):
        ModelDef(display_name="x", repo="o/n", local_path="/p", size_gb=1, mode="solo")  # both


def test_local_model_rejects_chat_template_outside_dir():
    # Only local_path is mounted into the container, so a chat-template elsewhere
    # would silently 404 inside vLLM — reject it up front.
    with pytest.raises(ValidationError):
        ModelDef(
            display_name="x", repo="", local_path="/m/ft", size_gb=1, mode="solo",
            vllm_args=["--chat-template=/other/dir/t.jinja"],
        )


def test_invalid_local_path_rejected_by_model():
    with pytest.raises(ValidationError):
        ModelDef(display_name="x", repo="", local_path="/m/../etc", size_gb=1, mode="solo")


def test_merge_overrides_loads_local_and_skips_invalid(monkeypatch):
    # YAML/override-added local models get the same validation as the API; a single
    # bad entry is skipped (logged) rather than breaking the whole catalog load.
    from app import models as M
    monkeypatch.setattr(M, "load_overrides", lambda: {"knobs": {}, "custom": [
        {"key": "good", "display_name": "G", "local_path": "/home/u/m", "size_gb": 1, "mode": "solo"},
        {"key": "bad", "display_name": "B", "local_path": "/home/u/../etc", "size_gb": 1, "mode": "solo"},
    ]})
    cat = M._merge_overrides(M.Catalog(models={}))
    assert cat.models["good"].is_local and cat.models["good"].source == "/home/u/m"
    assert "bad" not in cat.models  # traversal path skipped, not catalog-fatal