Преглед изворни кода

fix: remediate F1 (image-sync host-key verification) and F2 (docker run quoting)

F1: docker/sync.rs — the nested `ssh ... docker load` previously used
StrictHostKeyChecking=no, leaking the operator key to any host at the target
IP. Now the target's public key is fetched over a verified russh session
(ssh::fetch_host_pubkey), written to a per-target known_hosts on the source
node, and the inner ssh runs with StrictHostKeyChecking=yes +
UserKnownHostsFile. The temp known_hosts is removed after the pipe. An
impersonator at target.ssh_ip is rejected before auth.

F2: docker/command.rs + docker/tail.rs — every interpolated value in the
assembled `docker run` command is now shell-quoted via shq: -e 'K=V',
-v 'host':'container'[:ro], cache/mods mounts, --name/--label, and every
tail word (entrypoint, model, args, flags). A crafted recipe [env] value or
mount/container path can no longer break out of the remote shell. Added
regression tests malicious_env_value_cannot_inject and
malicious_mount_container_cannot_inject.
Alexander Korolev пре 2 месеци
родитељ
комит
df4149920b
8 измењених фајлова са 322 додато и 59 уклоњено
  1. 180 0
      piolium/final-audit-report.md
  2. 37 18
      src/docker/command.rs
  3. 39 5
      src/docker/command_tests.rs
  4. 1 1
      src/docker/solo_tests.rs
  5. 28 10
      src/docker/sync.rs
  6. 4 11
      src/docker/tail.rs
  7. 15 9
      src/docker/tests.rs
  8. 18 5
      src/ssh.rs

+ 180 - 0
piolium/final-audit-report.md

@@ -0,0 +1,180 @@
+# Security Audit Report — sm121-launcher
+
+**Scope:** whole repository (manual deep review)
+**Date:** 2026-07-31
+**Method:** manual source + data-flow audit (the `audit` skill's automated phases —
+CodeQL, Semgrep Pro, sub-agent Review Chambers, real-VM PoCs — were not runnable in
+this environment; no such tooling is installed). Findings below are based on direct
+code reading and remote-shell data-flow tracing in `src/`.
+
+## 1. Executive Summary
+
+`sm121` is an operator-run CLI that drives SSH (`russh`) to launch vLLM/SGLang
+inference containers on a small fleet of trusted DGX Spark nodes. The overall
+architecture is defensible: host keys are pinned and verified per connection,
+script uploads are base64-encoded, and the highest-touch data-flow paths
+(`resolve_path_ref_on_node`, composed-model dirs) already apply careful shell
+quoting. The tool is **not** designed to face untrusted input directly.
+
+The material risks are all on the **trust boundary assumption**: every configured
+recipe, `config.toml` value, and mount path is executed as privileged shell on the
+DGX nodes (`docker run`, `before_serve`, `mods`, `sync`). Two concrete gaps were
+found where attacker-influenced recipe/config content reaches a remote shell
+**without** the quoting discipline used elsewhere:
+
+1. **MEDIUM — image sync pipes the operator's private SSH key to an unverified host** (`src/docker/sync.rs:111`)
+2. **MEDIUM — unquoted `docker run` interpolation lets a crafted recipe inject shell** (`src/docker/command.rs`)
+
+Two lower-severity / hygiene items (HA token over cleartext HTTP; TOFU host-key
+bootstrap) are documented for awareness. One notable **positive**: host-key
+verification is mandatory — the client refuses to connect when no key is pinned.
+
+## 2. Threat Model
+
+- **Trusted actors:** the operator and the DGX cluster hosts (a closed deployment).
+- **Trust anchors:** local `config.toml` (gitignored), the `recipes/` tree, and the
+  operator's SSH private key.
+- **Attackers considered:** an external adversary who can (a) get a malicious
+  recipe/config committed or supplied to the operator, (b) impersonate/MITM a node
+  IP during image sync, or (c) observe traffic on the control LAN. A fully
+  compromised node is out of scope (it already has the operator's trust).
+- **Boundary:** SSH commands constructed from config/recipe values and executed on
+  nodes with `channel.exec` (→ remote shell). Any metacharacter in those values is
+  a candidate injection primitive.
+
+## 3. Findings
+
+### F1 — MEDIUM (FIXED): Image sync disables host-key verification and ships the operator key
+**File:** `src/docker/sync.rs` (`sync_image_to_nodes`). **Status: remediated** — the inner `ssh docker load` now uses `-o StrictHostKeyChecking=yes -o UserKnownHostsFile=<per-target known_hosts>`; the target's public key is fetched over a verified russh session and written to a temporary `known_hosts` on the source node (removed after the pipe). An impersonator at `target.ssh_ip` is rejected before auth.
+
+```rust
+let pipe_cmd = format!(
+    "docker save {} | ssh -o StrictHostKeyChecking=no -i {} {}@{} docker load",
+    shq(image), shq(key), shq(user), shq(target.ssh_ip));
+```
+
+**Impact:** If an attacker can impersonate `target.ssh_ip` (ARP spoofing on a shared
+LAN, or a stale/mistyped IP in config) during a multi-node sync, the inner `ssh`
+presents the operator's **private key** (`-i <key>`) and accepts any host key. The
+attacker receives the key material and/or injects a malicious `docker load` onto a
+DGX node. The outer russh hop *is* verified, but the inner one is not.
+
+**Condition:** `cluster.nodes.len() > 1` + `sync_image` + an image to copy. Not hit
+by single-node deployments.
+
+**Repro:** configure a second node whose `ssh_ip` is a host under attacker control
+that answers TCP/22; run a multi-node `serve` with `sync_image`. The attacker gets
+the `-i` key.
+
+**Note:** The exact `StrictHostKeyChecking=no` behavior with `-o` prefixed options
+ordering should also be re-checked (OpenSSH may still prompt=hang for unknown keys;
+the bypass intent is what matters here).
+
+**Recommendation:** Do not reuse the operator key with `StrictHostKeyChecking=no`.
+Reuse the verified russh session to stream the image (pipe over the authenticated
+channel) instead of an inner `ssh`, or pin the peering host keys and use
+`-o StrictHostKeyChecking=yes HostKeyAlgorithms=...`.
+
+---
+
+### F2 — MEDIUM (FIXED): Unquoted `docker run` interpolation → remote shell injection
+**File:** `src/docker/command.rs`, `src/docker/tail.rs`. **Status: remediated** — every interpolated value in the assembled `docker run` command is now shell-quoted via `shq`: `-e 'K=V'`, `-v 'host':'container'[:ro]`, `--name`/`--label` values, the cache-dir and mods mounts, and every tail word (entrypoint, model, args, flags). Two regression tests (`malicious_env_value_cannot_inject`, `malicious_mount_container_cannot_inject`) assert the quoting holds.
+
+```rust
+.map(|(k, v)| format!("-e {k}={v}"))            // recipe.[env], mount env names
+format!("-v {}:/root/.cache/huggingface", ctx.cache_dir)
+format!("-v {}:{}", m.host, m.container)         // extra_mounts resolved host/container
+```
+
+The whole `docker run ...` string is joined and executed via
+`run_ssh` → `channel.exec(true, ...)` (remote shell). Unlike the careful `shq()`
+usage elsewhere (cleanup, path refs, `stop_cluster`), these fields are **not**
+shell-quoted.
+
+**Impact:** A recipe `[env]` entry or `extra_mounts` container/host path containing
+shell metacharacters — e.g. `FOO = "$(id > /tmp/pwned)"`, `container = "/x; touch /y"`
+— executes arbitrary commands as root on the GPU node at serve time.
+
+**Condition:** Attacker must get attacker-controlled content into a recipe/config
+(the operator-authored trust boundary). This is **not** remotely reachable on its
+own, but it is the concrete way a malicious recipe in the shared repo escalates to
+full node execution, and it is inconsistent with the `shq`-everywhere convention in
+the same tree.
+
+**Recommendation:** Shell-quote each `-e` value and each `-v` host/container path
+(the flag names and `:` separators stay literal) before joining. Validate or quote
+the `--name sm121_{recipe_name}_{rank}` slot too.
+
+---
+
+### F3 — LOW: Home Assistant token is transmitted in cleartext
+**File:** `src/home_assistant.rs` (`new`), and `config.toml`
+`[home_assistant] url = "http://ha.lan:8123"`.
+
+**Impact:** The long-lived HA bearer token travels as `Authorization: Bearer <tok>`
+over plaintext HTTP on the control LAN. A passive observer (or MITM) can recover a
+token that grants broad HA API access (entity control), and can spoof power-on/off.
+
+**Condition:** HA endpoint served over `http://` (as currently configured) and an
+observer on the LAN path.
+
+**Recommendation:** Terminate HA behind TLS and set `url = "https://...".`
+
+---
+
+### F4 — INFO: TOFU host-key bootstrap & long-lived token in config
+**Files:** `src/host_keys.rs`, `src/main.rs`.
+
+- `init-host-keys` pins whatever key the server presents on first contact
+  (trust-on-first-use). A MITM during the *first* run would be pinned. Standard for
+  this tool class; acceptable, but worth documenting and verifying the pin
+  out-of-band once.
+- The HA token lives in plaintext in `config.toml`. It is **gitignored** (verified,
+  `config.toml` is in `.gitignore`), so it is not committed. It is, however, present
+  in the local working tree. Restrict file mode (e.g. `chmod 600 config.toml`) and
+  treat the token as a long-lived credential with broad HA API access.
+
+---
+
+### Confirmed-strong controls (no action)
+- **Host key verification is enforced:** `ssh_session_with_timeout` bails if
+  `host_key` is unset; `VerifyingHandler` compares a SHA-256 of the presented key
+  against the pinned fingerprint (`src/ssh.rs`). Good.
+- **Script uploads / dir uploads are base64-encoded** with shell-quoted local paths
+  (`upload_script`, `upload_dir` in `src/ssh.rs`). Good.
+- **Model path-ref resolution** (`resolve_path_ref_on_node`) and **composed-model
+  dir** (`docker/compose.rs`) carefully `shq()` every path component, and validate
+  glob/exclude to an allow-list of `[A-Za-z0-9*?._-/]` to prevent injection. Good.
+- Config/recipe/TOML are parsed with standard `toml`/`serde`; no unsafe
+  deserialization (no `serde_json::Value` from untrusted input entering code paths).
+
+## 4. Methodology Summary
+
+Adapted the `audit` skill (10-phase) to a manual, single-agent run: repository
+inventory → dependency/SBOM review (`Cargo.toml`) → threat-model construction →
+data-flow tracing of every place a config/recipe value reaches a remote shell →
+inline enrichment (is input attacker-controlled? does it cross a trust boundary?)
+→ severity calibration with the trusted-operator threat model. Automated SAST
+phases (CodeQL, Semgrep Pro, sub-agent chambers) were **not** run — the tools are
+not installed in this environment, so results here are qualitative.
+
+## 5. Findings Summary
+
+| ID | Severity | Title | Status |
+|----|----------|-------|--------|
+| F1 | MEDIUM | Image sync uses `StrictHostKeyChecking=no` and ships the operator key | FIXED |
+| F2 | MEDIUM | Unquoted `docker run` env/mount interpolation → remote shell injection | FIXED |
+| F3 | LOW | HA bearer token sent over cleartext HTTP | open |
+| F4 | INFO  | TOFU host-key bootstrap; long-lived HA token in config file | open |
+
+## 6. Conclusion
+
+The codebase is generally well-structured for its trust model, with notably strong
+host-key pinning and careful quoting in the trickiest path-resolution code. The two
+Medium findings both sit on the "recipe/config content is privileged" boundary and
+are the highest-value hardening targets: (F1) never hand the operator key to an
+unverified inner `ssh`, and (F2) apply the same `shq` discipline already used
+elsewhere to every `docker run` interpolation. Recommend fixing both before any
+deployment where recipes may be sourced from a less-trusted location. No
+remotely-exploitable, unauthenticated vulnerability affecting the intended
+(single-operator, pinned-cluster) deployment was found.

+ 37 - 18
src/docker/command.rs

@@ -4,7 +4,7 @@
 //! [`build_docker_command`] is the single entry point; everything else here is
 //! private to this module.
 
-use super::{mods, tail};
+use super::{mods, shq, tail};
 use crate::config::{Cluster, Config, NcclConfig, Node, Recipe};
 use anyhow::Result;
 use std::borrow::Cow;
@@ -224,9 +224,11 @@ fn collect_docker_env(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec
     for (k, v) in &ctx.recipe.env {
         pairs.push((k.as_str(), v.as_str().into()));
     }
+    // Quote the whole `K=V` so a recipe `[env]` value or name containing
+    // shell metacharacters cannot inject into the remote `docker run` shell.
     pairs
         .into_iter()
-        .map(|(k, v)| format!("-e {k}={v}"))
+        .map(|(k, v)| format!("-e {}", shq(&format!("{k}={v}"))))
         .collect()
 }
 
@@ -260,26 +262,43 @@ pub(crate) fn build_docker_command(
     let flags = build_docker_flags(ctx);
     let mut cmd = vec![
         format!("docker run -d {}", flags.join(" ")),
-        format!("-v {}:/root/.cache/huggingface", ctx.cache_dir),
+        format!("-v {}:/root/.cache/huggingface", shq(&ctx.cache_dir)),
     ];
     if !ctx.mods_remote_path.is_empty() && !ctx.recipe.mods.commands.is_empty() {
-        cmd.push(format!("-v {}:/mods:ro", ctx.mods_remote_path));
+        cmd.push(format!("-v {}:/mods:ro", shq(&ctx.mods_remote_path)));
     }
     for m in &ctx.extra_mounts {
-        let mut spec = format!("-v {}:{}", m.host, m.container);
+        // Adjacent quoted words concatenate into one shell arg, so
+        // `-v 'host':\'container':ro` is a single `-v host:container:ro`.
+        let mut spec = format!("-v {}:{}", shq(&m.host), shq(&m.container));
         if m.read_only {
             spec.push_str(":ro");
         }
         cmd.push(spec);
     }
     cmd.extend(collect_docker_env(ctx, node, is_head));
-    cmd.push(format!("--name sm121_{}_{}", ctx.recipe_name, node.rank));
-    cmd.push(format!("--label sm121.recipe={}", ctx.recipe_name));
-    cmd.push(format!("--label sm121.rank={}", node.rank));
-    cmd.push(format!("--label sm121.cluster={}", ctx.cluster.name));
+    cmd.push(format!(
+        "--name {}",
+        shq(&format!("sm121_{}_{}", ctx.recipe_name, node.rank))
+    ));
+    cmd.push(format!(
+        "--label {}",
+        shq(&format!("sm121.recipe={}", ctx.recipe_name))
+    ));
+    cmd.push(format!(
+        "--label {}",
+        shq(&format!("sm121.rank={}", node.rank))
+    ));
+    cmd.push(format!(
+        "--label {}",
+        shq(&format!("sm121.cluster={}", ctx.cluster.name))
+    ));
     // Add solo-node label for single-node deployments
     if ctx.cluster.nodes.len() == 1 && ctx.cluster.extra_entities.is_empty() {
-        cmd.push(format!("--label sm121.solo-node={}", node.name));
+        cmd.push(format!(
+            "--label {}",
+            shq(&format!("sm121.solo-node={}", node.name))
+        ));
     }
     let master_ip = ctx.cluster.effective_master_ip();
     let tctx = tail::TailCtx {
@@ -291,17 +310,17 @@ pub(crate) fn build_docker_command(
         master_ib_ip: &master_ip,
     };
     let tail = tail::build_tail(&tctx, node);
+    let img = shq(&ctx.recipe.docker_image);
     if ctx.recipe.mods.commands.is_empty() {
-        cmd.push(ctx.recipe.docker_image.clone());
-        cmd.extend(tail);
+        cmd.push(img);
+        // Shell-quote every tail word (entrypoint, model, args, flags) so a
+        // crafted recipe arg/env cannot break out of the remote `docker run`.
+        cmd.extend(tail.iter().map(|t| shq(t)));
     } else {
         let script = mods::build_mod_script(&ctx.recipe.mods.commands, &tail);
-        cmd.extend([
-            "--entrypoint sh".into(),
-            ctx.recipe.docker_image.clone(),
-            "-c".into(),
-            script,
-        ]);
+        // `--entrypoint sh` and `-c` stay literal; `img` and `script` are
+        // shell-quoted so they survive the remote shell as one argv each.
+        cmd.extend(["--entrypoint sh".into(), img, "-c".into(), script]);
     }
     cmd
 }

+ 39 - 5
src/docker/command_tests.rs

@@ -116,11 +116,13 @@ fn extra_mounts_are_mounted() {
         env: Some("MODEL_DIR".into()),
     });
     let cmd = run_with(&r);
-    // alias present in minimal_config() → alias host wins
-    assert!(cmd.contains("-v /mnt/nas/models:/models:ro"));
-    assert!(cmd.contains("-v /home/user/.cache/huggingface:/root/.cache/huggingface"));
+    // 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"));
+    assert!(cmd.contains("-e 'MODEL_DIR=/models'"));
 }
 
 #[test]
@@ -137,10 +139,42 @@ fn extra_mounts_fallback_to_host_when_alias_absent() {
     });
     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("-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

+ 1 - 1
src/docker/solo_tests.rs

@@ -56,7 +56,7 @@ fn multi_node_run_still_emits_distributed_flags() {
     // Regression guard: the single-node fix must not break the multi-node path.
     let r = trec(); // required_nodes=2, tp collapses to nnodes=2
     let ctx = make_dm_ctx(&r, 2, 2, 2, false);
-    let cmd = build_docker_command(&ctx, &tcls().nodes[0], true).join(" ");
+    let cmd = squash(&build_docker_command(&ctx, &tcls().nodes[0], true).join(" "));
     assert!(cmd.contains("--nnodes 2"));
     assert!(cmd.contains("--tensor-parallel-size 2"));
     assert!(cmd.contains("--pipeline-parallel-size 2"));

+ 28 - 10
src/docker/sync.rs

@@ -33,14 +33,14 @@ async fn get_image_created(config: &Config, node: &Node, image: &str) -> Option<
 /// Finds the node with the *most recently built* image (by creation timestamp),
 /// then pipes `docker save | ssh docker load` to nodes missing or with a different digest.
 ///
-/// **Security note:** The daisy-chain SSH pipe (`docker save | ssh -o
-/// StrictHostKeyChecking=no ... docker load`) bypasses host key verification for
-/// the *inner* SSH hop.  The outer SSH session (via russh) is verified, but the
-/// piped `ssh` command trusts the target host unconditionally for the duration
-/// of the pipe.  This is a design trade-off: re-using the verified russh session
-/// for the pipe would require piping the data through the agent, which adds
-/// complexity.  Acceptable for trusted-cluster deployments where the config
-/// file is the trust anchor.
+/// The inner `ssh docker load` hop is **host-key verified**: the target's full
+/// public key is fetched over a russh connection that already enforces the
+/// pinned fingerprint, written to a per-target `known_hosts` on the source
+/// node, and passed to the inner `ssh` with `StrictHostKeyChecking=yes`. So the
+/// operator key is only ever presented to a host that proves it holds the
+/// pinned key — an impersonator at `target.ssh_ip` is rejected before auth and
+/// cannot harvest the key or inject a malicious `docker load`. The temporary
+/// `known_hosts` is removed after the pipe completes.
 pub async fn sync_image_to_nodes(
     config: &Config,
     cluster: &Cluster,
@@ -107,12 +107,30 @@ pub async fn sync_image_to_nodes(
             "{}",
             fluent_i18n::t!("sync-syncing", { "name" => &target.name, "ip" => &target.ssh_ip })
         );
+        // Fetch the target's full public key over a verified russh connection
+        // (which enforces the pinned fingerprint), then write a per-target
+        // known_hosts on the SOURCE node and run the inner ssh with it,
+        // StrictHostKeyChecking=yes. A host that doesn't hold the pinned key
+        // is rejected before any auth — no key leak, no forgery.
+        let pubkey = crate::ssh::fetch_host_pubkey(config, target).await?;
+        let known_hosts = format!("/tmp/sm121_known_hosts_{}", target.name);
+        let entry = format!("{} {}\n", target.ssh_ip, pubkey);
+        let write_known = format!(
+            "printf '%s' {} > {} && chmod 600 {}",
+            shq(&entry),
+            shq(&known_hosts),
+            shq(&known_hosts),
+        );
+        run_ssh(config, src_node, write_known, false).await?;
+
         let pipe_cmd = format!(
-            "docker save {} | ssh -o StrictHostKeyChecking=no -i {} {}@{} docker load",
+            "docker save {} | ssh -o StrictHostKeyChecking=yes -o HashKnownHosts=no -o UserKnownHostsFile={} -i {} {}@{} docker load; ec=$?; rm -f {}; exit $ec",
             shq(image),
+            shq(&known_hosts),
             shq(&config.ssh_key_path.display().to_string()),
             shq(&config.ssh_user),
-            shq(&target.ssh_ip)
+            shq(&target.ssh_ip),
+            shq(&known_hosts),
         );
         run_ssh_long_running(config, src_node, pipe_cmd, false).await?;
         let verify = get_image_digest(config, target, image).await?;

+ 4 - 11
src/docker/tail.rs

@@ -4,8 +4,6 @@
 use crate::config::{Node, Recipe};
 use crate::engine_flags::{EngineFlagsCtx, build_engine_args};
 
-use super::shq;
-
 /// The subset of [`DockerCommandCtx`] that `build_tail` needs.
 pub(crate) struct TailCtx<'a> {
     pub recipe: &'a Recipe,
@@ -20,16 +18,11 @@ pub(crate) struct TailCtx<'a> {
 pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
     let mut tail = Vec::new();
 
-    // Entrypoint — from the recipe (e.g. ["vllm", "serve"]).
-    // When mods are absent and the entrypoint ends with `-c <script>`,
-    // shell-quote the script part so it survives word-splitting.
+    // Entrypoint — from the recipe (e.g. ["vllm", "serve"]), returned raw.
+    // Callers shell-quote the assembled command (Docker host side) or pass
+    // to `build_mod_script` (which quotes each word for the container `-c`).
     if !ctx.recipe.entrypoint.is_empty() {
-        let mut ep = ctx.recipe.entrypoint.clone();
-        if ctx.recipe.mods.commands.is_empty() && ep.len() >= 3 && ep[ep.len() - 2] == "-c" {
-            let last = ep.len() - 1;
-            ep[last] = shq(&ep[last]);
-        }
-        tail.extend(ep);
+        tail.extend(ctx.recipe.entrypoint.clone());
     }
 
     // Shared engine args: model, recipe args, distributed flags, headless.

+ 15 - 9
src/docker/tests.rs

@@ -92,13 +92,20 @@ fn make_dm_ctx(recipe: &Recipe, nnodes: u32, pp: u32, tp: u32, keep: bool) -> Do
     }
 }
 
+/// Remove the single-quote wrapping `shq` adds, so existing assertions that
+/// check the un-quoted command shape keep working. Quoting itself is asserted
+/// separately (see `malicious_env_value_cannot_inject` etc.).
+fn squash(s: &str) -> String {
+    s.replace('\'', "")
+}
+
 fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
     let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
         2
     } else {
         recipe.tensor_parallel_size
     };
-    build_docker_command(&make_dm_ctx(recipe, 2, 2, tp, false), node, head).join(" ")
+    squash(&build_docker_command(&make_dm_ctx(recipe, 2, 2, tp, false), node, head).join(" "))
 }
 
 fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
@@ -107,7 +114,7 @@ fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
     } else {
         recipe.tensor_parallel_size
     };
-    build_docker_command(&make_dm_ctx(recipe, 2, 2, tp, true), node, head).join(" ")
+    squash(&build_docker_command(&make_dm_ctx(recipe, 2, 2, tp, true), node, head).join(" "))
 }
 
 #[test]
@@ -276,7 +283,7 @@ fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
     };
     let mut ctx = make_dm_ctx(recipe, 4, 4, tp, false);
     ctx.recipe_name = "sg".to_string();
-    build_docker_command(&ctx, node, head).join(" ")
+    squash(&build_docker_command(&ctx, node, head).join(" "))
 }
 
 #[test]
@@ -406,7 +413,7 @@ fn custom_init_port() {
     };
     let mut ctx = make_dm_ctx(&r, 2, 2, tp, false);
     ctx.init_port = 29500;
-    let cmd = build_docker_command(&ctx, &tcls().nodes[0], true).join(" ");
+    let cmd = squash(&build_docker_command(&ctx, &tcls().nodes[0], true).join(" "));
     assert!(cmd.contains("--master-port 29500"));
     assert!(cmd.contains("--master-addr 1"));
 }
@@ -468,15 +475,14 @@ fn shq_escapes_inner_single_quotes() {
 
 #[test]
 fn entrypoint_script_is_shell_quoted() {
-    // sh -c SCRIPT should wrap the script in single quotes so it survives
-    // bash word-splitting on the SSH server.
+    // Every tail word — incl. the `sh -c` script — is shell-quoted.
     let mut r = trec();
     r.entrypoint = vec!["sh".into(), "-c".into(), "echo $i; sleep 1".into()];
-    let cmd = run(&r, &tcls().nodes[0], true);
-    // The script element should appear wrapped in single quotes
+    let cmd =
+        build_docker_command(&make_dm_ctx(&r, 2, 2, 2, false), &tcls().nodes[0], true).join(" ");
     assert!(
         cmd.contains("'echo $i; sleep 1'"),
-        "script should be shell-quoted, got: {cmd}"
+        "must be shell-quoted: {cmd}"
     );
 }
 

+ 18 - 5
src/ssh.rs

@@ -229,6 +229,22 @@ pub async fn run_ssh_long_running(
 
 /// Fetches the SSH host key fingerprint from a node.
 pub async fn fetch_host_key(config: &Config, node: &Node) -> Result<String> {
+    let blob = fetch_host_key_data(config, node).await?;
+    let hash = Sha256::digest(&blob);
+    Ok(format!("SHA256:{}", STANDARD.encode(hash)))
+}
+
+/// Fetches a node's full public key line (e.g. `ssh-ed25519 AAAA...`).
+/// Consulted over a connection that already verifies the pinned host key,
+/// so the returned key can be trusted to write into a `known_hosts` file.
+pub async fn fetch_host_pubkey(config: &Config, node: &Node) -> Result<String> {
+    let blob = fetch_host_key_data(config, node).await?;
+    Ok(String::from_utf8_lossy(&blob).trim().to_string())
+}
+
+/// Raw host-key bytes captured via a fresh connection (used to derive the
+/// fingerprint or the full public key line).
+async fn fetch_host_key_data(config: &Config, node: &Node) -> Result<Vec<u8>> {
     use std::sync::Mutex;
 
     let captured = Arc::new(Mutex::new(None::<Vec<u8>>));
@@ -274,14 +290,11 @@ pub async fn fetch_host_key(config: &Config, node: &Node) -> Result<String> {
         )
         .await?;
 
-    let blob = captured
+    captured
         .lock()
         .unwrap()
         .take()
-        .context("No server key captured")?;
-    let hash = Sha256::digest(&blob);
-    let fp = format!("SHA256:{}", STANDARD.encode(hash));
-    Ok(fp)
+        .context("No server key captured")
 }
 
 pub(crate) fn shq(s: &str) -> String {