launcher/src/docker/command_tests.rs

use super::*;
use crate::config::{
    Cluster, DeployTarget, DistributedBackend, ExtraMount, NcclConfig, Node, Recipe, RecipeArgs,
    VenvSetup,
};
use std::collections::HashMap;

fn tcls() -> Cluster {
    Cluster {
        name: "t".into(),
        master_ib_ip: Some("1".into()),
        nodes: vec![Node {
            name: "n0".into(),
            ssh_ip: "10.0.0.1".into(),
            ib_ip: Some("1".into()),
            rank: 0,
            host_key: None,
            entity: None,
        }],
        extra_entities: vec![],
    }
}

fn tnccl() -> NcclConfig {
    NcclConfig {
        ib_hca: "r".into(),
        socket_ifname: "e".into(),
        ib_gid_index: Some(3),
        ib_disable: 0,
        gloo_socket_ifname: None,
        ucx_net_devices: None,
        nccl_cross_nic: None,
        nccl_proto: None,
        nccl_p2p_level: None,
    }
}

fn minimal_config() -> Config {
    Config {
        ssh_user: "u".into(),
        ssh_key_path: "/tmp/key".into(),
        clusters: vec![],
        nccl: HashMap::new(),
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        global_mods_dir: None,
        drop_caches: true,
        mount_aliases: HashMap::from([("models".to_string(), "/mnt/nas/models".to_string())]),
        recipe_sources: Vec::new(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
    }
}

fn run_with(recipe: &Recipe) -> String {
    let resolved = resolve_extra_mounts(&minimal_config(), recipe).unwrap();
    let ctx = DockerCommandCtx {
        cluster: tcls(),
        nccl: Some(tnccl()),
        recipe: recipe.clone(),
        recipe_name: "r".to_string(),
        pp_size: 2,
        tp_size: 1,
        nnodes: 2,
        cache_dir: "/home/user/.cache/huggingface".to_string(),
        keep: false,
        init_port: 25000,
        mods_remote_path: String::new(),
        extra_mounts: resolved.static_mounts,
    };
    build_docker_command(&ctx, &tcls().nodes[0], true).join(" ")
}

fn base_recipe() -> Recipe {
    Recipe {
        engine: "vllm".into(),
        name: "T".into(),
        docker_image: "img".into(),
        model: "mdl".into(),
        port: 8000,
        required_nodes: Some(2),
        pipeline_parallel_size: Some(2),
        tensor_parallel_size: 1,
        args: RecipeArgs::List(vec!["--x".into(), "4096".into()]),
        env: HashMap::new(),
        cluster: None,
        entrypoint: vec![],
        distributed_backend: Some(DistributedBackend::Mp),
        before_serve: vec![],
        after_stop: vec![],
        solo: None,
        deploy: DeployTarget::Docker,
        venv: None,
        venv_setup: VenvSetup::default(),
        venv_setup_script: None,
        mods: Default::default(),
        shm_size: String::new(),
        sync_image: true,
        docker_image_dir: None,
        extra_mounts: vec![],
        build: Default::default(),
        composed_model: None,
    }
}

#[test]
fn extra_mounts_are_mounted() {
    let mut r = base_recipe();
    r.extra_mounts.push(ExtraMount {
        alias: Some("models".into()),
        host: Some("~/.cache/models".into()),
        repo: None,
        subfolder: None,
        container: "/models".into(),
        read_only: true,
        env: Some("MODEL_DIR".into()),
    });
    let cmd = run_with(&r);
    // alias present in minimal_config() → alias host wins. Host/container
    // and env are now shell-quoted so a crafted value can't inject.
    assert!(cmd.contains("-v '/mnt/nas/models':'/models':ro"));
    assert!(cmd.contains("-v '/home/user/.cache/huggingface':/root/.cache/huggingface"));
    assert!(!cmd.contains("-v /mnt/nas/models:/models:ro")); // unquoted form gone
    // env injected pointing at the container path, before recipe env
    assert!(cmd.contains("-e 'MODEL_DIR=/models'"));
}

#[test]
fn extra_mounts_fallback_to_host_when_alias_absent() {
    let mut r = base_recipe();
    r.extra_mounts.push(ExtraMount {
        alias: Some("no_such_alias".into()),
        host: Some("~/.cache/jit".into()),
        repo: None,
        subfolder: None,
        container: "/root/.cache/jit".into(),
        read_only: false,
        env: None,
    });
    let cmd = run_with(&r);
    // alias not in config → host default used with ~ expansion (ssh_user=u)
    assert!(cmd.contains("-v '/home/u/.cache/jit':'/root/.cache/jit'"));
    assert!(!cmd.contains(":ro"));
}

#[test]
fn malicious_env_value_cannot_inject() {
    // A crafted `[env]` value with shell metachars must stay a single
    // quoted argument — never break out of the remote `docker run`.
    let mut r = base_recipe();
    r.env
        .insert("EVIL".to_string(), "x; touch /tmp/pwned".to_string());
    let cmd = run_with(&r);
    assert!(cmd.contains("-e 'EVIL=x; touch /tmp/pwned'"));
    // No raw `; touch` escape: the whole K=V is inside quotes.
    assert!(!cmd.contains("-e EVIL=x;"));
}

#[test]
fn malicious_mount_container_cannot_inject() {
    let mut r = base_recipe();
    r.extra_mounts.push(ExtraMount {
        alias: None,
        host: Some("/safe".into()),
        repo: None,
        subfolder: None,
        container: "/x; touch /tmp/pwned".into(),
        read_only: false,
        env: None,
    });
    let cmd = run_with(&r);
    // Container path is quoted; the metacharacter is inert inside quotes.
    assert!(cmd.contains("'/x; touch /tmp/pwned'"));
    // Never appear as the unquoted concatenation `-v /safe:/x; touch`.
    assert!(!cmd.contains("-v /safe:/x; touch"));
}

#[test]
fn model_path_ref_is_partitioned_not_mounted() {
    // repo-based mounts are path references, NOT bind-mounts. They are
    // partitioned into path_refs and resolved per-node via substitution.
    let mut r = base_recipe();
    r.extra_mounts.push(ExtraMount {
        alias: None,
        host: None,
        repo: Some("aidendle94/GLM-5.2-Hybrid-FP8-MXPF4".into()),
        subfolder: Some("mtp-draft".into()),
        container: "mtp_draft".into(),
        read_only: true,
        env: None,
    });
    let resolved = resolve_extra_mounts(&minimal_config(), &r).unwrap();
    assert_eq!(resolved.path_refs.len(), 1);
    assert_eq!(resolved.static_mounts.len(), 0);
    let pr = &resolved.path_refs[0];
    assert_eq!(pr.repo, "aidendle94/GLM-5.2-Hybrid-FP8-MXPF4");
    assert_eq!(pr.subfolder.as_deref(), Some("mtp-draft"));
    assert_eq!(pr.name, "mtp_draft");
    // No -v flag for path refs in build_docker_command (substitution only)
    let cmd = run_with(&r);
    assert!(!cmd.contains("-v /models/mtp-draft"));
}

#[test]
fn substitute_path_refs_replaces_placeholders() {
    let mut r = base_recipe();
    r.args = RecipeArgs::List(vec![
        "--speculative-config".into(),
        "{\"model\":\"{{mtp_draft}}\"}".into(),
    ]);
    let path_refs = vec![(
        "mtp_draft".to_string(),
        "/root/.cache/huggingface/hub/models--aidendle94--GLM-5.2-Hybrid-FP8-MXPF4/snapshots/bb0d0858/mtp-draft".to_string(),
    )];
    let substituted = substitute_path_refs(&r, &path_refs);
    if let RecipeArgs::List(args) = &substituted.args {
        assert_eq!(
            args[1],
            "{\"model\":\"/root/.cache/huggingface/hub/models--aidendle94--GLM-5.2-Hybrid-FP8-MXPF4/snapshots/bb0d0858/mtp-draft\"}"
        );
    } else {
        panic!("expected list args");
    }
}

#[test]
fn substitute_path_refs_in_env() {
    let mut r = base_recipe();
    r.env
        .insert("MTP_PATH".to_string(), "{{mtp_draft}}".to_string());
    let path_refs = vec![(
        "mtp_draft".to_string(),
        "/root/.cache/huggingface/hub/models--aidendle94--GLM-5.2-Hybrid-FP8-MXPF4/snapshots/bb0d0858/mtp-draft".to_string(),
    )];
    let substituted = substitute_path_refs(&r, &path_refs);
    assert_eq!(
        substituted.env["MTP_PATH"],
        "/root/.cache/huggingface/hub/models--aidendle94--GLM-5.2-Hybrid-FP8-MXPF4/snapshots/bb0d0858/mtp-draft"
    );
}

#[test]
fn clear_paths_dedups_and_sorts() {
    let m = |host: &str| ResolvedMount {
        host: host.into(),
        container: "/c".into(),
        read_only: false,
        env: None,
    };
    let mounts = vec![m("/b/.cache/jit"), m("/a/.cache/jit"), m("/b/.cache/jit")];
    assert_eq!(clear_paths(&mounts), vec!["/a/.cache/jit", "/b/.cache/jit"]);
    assert!(clear_paths(&[]).is_empty());
}