refactor: comprehensive rewrite — concurrent SSH, labels, config-driven clusters, host key verification

Major improvements:

- Container lifecycle: add --name and --label flags to docker run
  for unambiguous identification; use label-based docker ps/stop/logs
  filters for targeted operations
- Concurrent SSH: replace sequential node iteration with
  tokio::spawn + futures::join_all in start_cluster and health_check,
  with automatic failure collection and summary reporting
- Retry logic: add run_ssh_with_retry with exponential backoff
  (2 retries for start, 1 for logs/health) to handle transient SSH drops
- Recipe-cluster binding: add recipe.cluster field to recipe.toml
  with fallback chain: CLI --cluster > recipe.cluster > auto
- required_nodes validation: fail early if cluster size is insufficient
  for the recipe
- Engine entrypoint: add recipe.entrypoint field to inject custom
  commands between image and model (e.g. SGLang --host)
- Host key verification: add VerifyingHandler that checks SSH
  host key fingerprints (SHA256) against config.toml; add
  health --fingerprint command for key discovery
- Stop error reporting: collect per-node failures instead of
  failing silently; warn and exit with error if any node fails
- Dependencies: add sha2 0.10, base64 0.22, futures 0.3
- Rename recipies/ to recipes/ (fix directory typo)
- Update AGENTS.md with new fields and commands

AGENTS.md

@@ -7,7 +7,7 @@ Rust CLI tool for launching vLLM/SGLang containers across a DGX Spark cluster. S
 - `Cargo.toml` — binary crate `vllm-spark-launcher`. Edition 2024.
 - `src/main.rs` — single-file binary: CLI parsing, SSH, docker orchestration.
 - `config.toml` — runtime config: SSH key, cluster definitions, per-cluster NCCL settings.
-- `recipies/<name>/recipe.toml` — recipe files (note: directory is misspelled "recipies").
+- `recipes/<name>/recipe.toml` — recipe files.
 
 No `build.rs`, no workspace, no additional modules. Everything is in `main.rs`.
 

Cargo.lock

@@ -216,6 +216,12 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fd307490d624467aa6f74b0eabb77633d1f758a7b25f12bceb0b22e08d9726f6"
 
+[[package]]
+name = "base64"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
+
 [[package]]
 name = "base64ct"
 version = "1.8.3"
@@ -2526,10 +2532,13 @@ name = "sm121-launcher"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "base64",
  "clap",
+ "futures",
  "russh",
  "russh-keys",
  "serde",
+ "sha2 0.10.9",
  "tokio",
  "toml",
  "tracing",

Cargo.toml

@@ -12,4 +12,7 @@ anyhow = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 toml = "1"
 tracing = "0.1"
-tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+sha2 = { version = "0.10", features = ["std"] }
+base64 = "0.22"
+futures = "0.3"

src/main.rs

@@ -1,8 +1,11 @@
 use anyhow::{Context, Result};
+use base64::engine::general_purpose::STANDARD;
+use base64::Engine;
 use clap::{Parser, Subcommand};
 use russh::keys::{PrivateKeyWithHashAlg, load_secret_key as russh_load_secret_key};
 use russh::*;
 use serde::Deserialize;
+use sha2::{Digest, Sha256};
 use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::fs;
@@ -55,6 +58,8 @@ enum Commands {
             help = "Имя кластера (например, 4nodes-switch). Если не указано - проверяет все."
         )]
         cluster: Option<String>,
+        #[arg(long, help = "Показать хост-ключи узлов (для добавления в config.toml)")]
+        fingerprint: bool,
     },
 }
 
@@ -64,6 +69,8 @@ struct Node {
     ssh_ip: String,
     ib_ip: String,
     rank: u32,
+    #[serde(default)]
+    host_key: Option<String>,
 }
 
 #[derive(Clone, Deserialize)]
@@ -89,7 +96,7 @@ struct Config {
     nccl: HashMap<String, NcclConfig>,
 }
 
-#[derive(Deserialize)]
+#[derive(Clone, Deserialize)]
 struct Recipe {
     #[serde(default = "default_engine")]
     engine: String,
@@ -108,6 +115,10 @@ struct Recipe {
     args: Vec<String>,
     #[serde(default)]
     env: HashMap<String, String>,
+    #[serde(default)]
+    cluster: Option<String>,
+    #[serde(default)]
+    entrypoint: Vec<String>,
 }
 
 fn default_engine() -> String {
@@ -128,23 +139,42 @@ async fn main() -> Result<()> {
             serve(&config, &recipe, &cluster, dry_run).await?;
         }
         Commands::Logs { recipe, follow } => {
-            let cluster = get_cluster_for_recipe(&config, &recipe).await?;
-            show_logs(&config, &cluster, follow, dry_run).await?;
+            let (cluster, recipe_name) = get_cluster_for_recipe(&config, &recipe).await?;
+            show_logs(&config, &cluster, &recipe_name, follow, dry_run).await?;
         }
         Commands::Stop { recipe } => {
-            let cluster = get_cluster_for_recipe(&config, &recipe).await?;
-            stop_cluster(&config, &cluster, dry_run).await?;
+            let (cluster, recipe_name) = get_cluster_for_recipe(&config, &recipe).await?;
+            stop_cluster(&config, &cluster, &recipe_name, dry_run).await?;
         }
         Commands::Health {
             cluster: cluster_name,
-        } => match cluster_name {
-            Some(ref name) => {
-                let cluster = select_cluster_by_name(&config, name)?;
-                health_check(&config, &cluster, dry_run).await?;
-            }
-            None => {
-                for cluster in &config.clusters {
-                    health_check(&config, cluster, dry_run).await?;
+            fingerprint,
+        } => {
+            if fingerprint {
+                let clusters_to_check: Vec<Cluster> = match &cluster_name {
+                    Some(name) => vec![select_cluster_by_name(&config, name)?],
+                    None => config.clusters.clone(),
+                };
+                for cluster in &clusters_to_check {
+                    println!("\n=== Кластер: {} ===", cluster.name);
+                    for node in &cluster.nodes {
+                        match fetch_host_key(&config, node).await {
+                            Ok(fp) => println!("{} ({}) -> host_key = \"{}\"", node.name, node.ssh_ip, fp),
+                            Err(e) => println!("{} ({}) -> error: {}", node.name, node.ssh_ip, e),
+                        }
+                    }
+                }
+            } else {
+                match cluster_name {
+                    Some(ref name) => {
+                        let cluster = select_cluster_by_name(&config, name)?;
+                        health_check(&config, &cluster, dry_run).await?;
+                    }
+                    None => {
+                        for cluster in &config.clusters {
+                            health_check(&config, cluster, dry_run).await?;
+                        }
+                    }
                 }
             }
         },
@@ -160,30 +190,50 @@ async fn serve(
     cluster_name: &str,
     dry_run: bool,
 ) -> Result<()> {
-    let recipe_path = format!("recipies/{}/recipe.toml", recipe_name);
+    let recipe_path = format!("recipes/{}/recipe.toml", recipe_name);
     let recipe: Recipe = toml::from_str(&fs::read_to_string(&recipe_path).await?)?;
 
     info!("🚀 Запуск рецепта: {} ({})", recipe.name, recipe.engine);
 
-    let (cluster, nccl) = select_cluster(config, cluster_name)?;
+    let cluster_choice = if cluster_name != "auto" {
+        cluster_name.to_string()
+    } else if let Some(ref rc) = recipe.cluster {
+        rc.clone()
+    } else {
+        "auto".to_string()
+    };
+    let (cluster, nccl) = select_cluster(config, &cluster_choice)?;
+
+    if let Some(required) = recipe.required_nodes {
+        let actual = cluster.nodes.len() as u32;
+        if actual != required {
+            anyhow::bail!(
+                "Recipe '{}' requires {} nodes, but cluster '{}' has {} nodes",
+                recipe.name, required, cluster.name, actual
+            );
+        }
+    }
 
-    start_cluster(config, &cluster, nccl, &recipe, dry_run).await?;
+    start_cluster(config, &cluster, nccl, recipe_name, &recipe, dry_run).await?;
 
     if !dry_run {
         info!("\n📜 Последние логи контейнеров:");
-        show_logs(config, &cluster, false, dry_run).await?;
+        show_logs(config, &cluster, recipe_name, false, dry_run).await?;
     }
 
     Ok(())
 }
 
 // ====================== Логи ======================
-async fn show_logs(config: &Config, cluster: &Cluster, follow: bool, dry_run: bool) -> Result<()> {
+async fn show_logs(config: &Config, cluster: &Cluster, recipe_name: &str, follow: bool, dry_run: bool) -> Result<()> {
     if follow {
         let head_node = cluster.nodes.first().context("No nodes in cluster")?;
         info!("📜 Following logs for '{}' (head node)...", head_node.name);
-        let cmd = "docker ps -q | head -n 1 | xargs -r docker logs -f";
-        run_ssh_follow(config, head_node, cmd.to_string(), dry_run).await?;
+        let cmd = format!(
+            "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank=0\" | xargs -r docker logs -f",
+            recipe_name
+        );
+        run_ssh_follow(config, head_node, cmd, dry_run).await?;
     } else {
         let flag = "--tail=100";
         info!(
@@ -192,15 +242,29 @@ async fn show_logs(config: &Config, cluster: &Cluster, follow: bool, dry_run: bo
             cluster.nodes.len()
         );
 
+        let mut handles = Vec::new();
         for node in &cluster.nodes {
             let cmd = format!(
-                "docker ps -q | head -n 1 | xargs -r docker logs {} 2>&1",
-                flag
+                "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank={}\" | xargs -r docker logs {} 2>&1",
+                recipe_name, node.rank, flag
             );
-            let output = run_ssh(config, cluster, node, cmd, dry_run).await?;
-            if !output.trim().is_empty() {
-                println!("\n--- {} ---", node.name);
-                println!("{}", output);
+            let config = config.clone();
+            let cluster = cluster.clone();
+            let node = node.clone();
+            let dry = dry_run;
+            handles.push(tokio::spawn(async move {
+                let output = run_ssh_with_retry(&config, &cluster, &node, cmd, dry, 1).await;
+                (node.name, output)
+            }));
+        }
+
+        let results = futures::future::join_all(handles).await;
+        for result in results {
+            if let Ok((name, Ok(output))) = result {
+                if !output.trim().is_empty() {
+                    println!("\n--- {} ---", name);
+                    println!("{}", output);
+                }
             }
         }
     }
@@ -208,26 +272,37 @@ async fn show_logs(config: &Config, cluster: &Cluster, follow: bool, dry_run: bo
 }
 
 // ====================== Stop ======================
-async fn stop_cluster(config: &Config, cluster: &Cluster, dry_run: bool) -> Result<()> {
+async fn stop_cluster(config: &Config, cluster: &Cluster, recipe_name: &str, dry_run: bool) -> Result<()> {
     info!("🛑 Останавливаем кластер '{}'...", cluster.name);
+    let mut failures = Vec::new();
+
     for node in &cluster.nodes {
-        let _ = run_ssh(
-            config,
-            cluster,
-            node,
-            "docker ps -q | xargs -r docker stop".to_string(),
-            dry_run,
-        )
-        .await;
-        let _ = run_ssh(
-            config,
-            cluster,
-            node,
-            "docker ps -q | xargs -r docker rm -f".to_string(),
-            dry_run,
-        )
-        .await;
+        let stop_cmd = format!(
+            "docker ps -q --filter \"label=sm121.recipe={}\" | xargs -r docker stop",
+            recipe_name
+        );
+        if let Err(e) = run_ssh(config, cluster, node, stop_cmd, dry_run).await {
+            tracing::warn!("Stop failed on {}: {}", node.name, e);
+            failures.push((node.name.clone(), "stop".to_string(), e));
+        }
+
+        let rm_cmd = format!(
+            "docker ps -qa --filter \"label=sm121.recipe={}\" | xargs -r docker rm",
+            recipe_name
+        );
+        if let Err(e) = run_ssh(config, cluster, node, rm_cmd, dry_run).await {
+            tracing::warn!("Remove failed on {}: {}", node.name, e);
+            failures.push((node.name.clone(), "rm".to_string(), e));
+        }
+    }
+
+    if !failures.is_empty() {
+        for (name, action, e) in &failures {
+            tracing::warn!("Failed to {} on {}: {}", action, name, e);
+        }
+        anyhow::bail!("Failed to stop containers on {} node(s)", failures.len());
     }
+
     info!("✅ Все контейнеры остановлены.");
     Ok(())
 }
@@ -236,40 +311,60 @@ async fn stop_cluster(config: &Config, cluster: &Cluster, dry_run: bool) -> Resu
 async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) -> Result<()> {
     println!("\n=== Кластер: {} ===", cluster.name);
 
-    let mut all_healthy = true;
-
+    let mut handles = Vec::new();
     for node in &cluster.nodes {
-        print!("{} ({}): ", node.name, node.ib_ip);
-
-        let docker_check = run_ssh(
-            config,
-            cluster,
-            node,
-            "docker ps --format '{{.Names}}' 2>&1".to_string(),
-            dry_run,
-        )
-        .await?;
-
-        if docker_check.is_empty() && dry_run {
-            continue;
-        }
+        let config = config.clone();
+        let cluster = cluster.clone();
+        let node = node.clone();
+        let dry = dry_run;
+        handles.push(tokio::spawn(async move {
+            let result = run_ssh_with_retry(
+                &config,
+                &cluster,
+                &node,
+                "docker ps --filter \"label=sm121\" --format '{{.Names}} {{.Status}}' 2>&1".to_string(),
+                dry,
+                1,
+            ).await;
+            (node.name, node.ib_ip, result)
+        }));
+    }
 
-        if docker_check.starts_with("docker:") || docker_check.contains("permission denied") {
-            println!("✗ Docker недоступен");
-            if let Some(line) = docker_check.lines().next() {
-                println!("  └─ {}", line);
+    let mut all_healthy = true;
+    let results = futures::future::join_all(handles).await;
+    for result in results {
+        match result {
+            Ok((name, ib_ip, Ok(docker_check))) => {
+                print!("{} ({}): ", name, ib_ip);
+                if docker_check.is_empty() && dry_run {
+                    continue;
+                }
+                if docker_check.starts_with("docker:") || docker_check.contains("permission denied") {
+                    println!("✗ Docker недоступен");
+                    if let Some(line) = docker_check.lines().next() {
+                        println!("  └─ {}", line);
+                    }
+                    all_healthy = false;
+                } else {
+                    let containers: Vec<&str> = docker_check.lines().filter(|l| !l.is_empty()).collect();
+                    if containers.is_empty() {
+                        println!("✗ Нет контейнеров");
+                        all_healthy = false;
+                    } else {
+                        println!("✓ {} контейнер(ов)", containers.len());
+                        for c in containers {
+                            println!("  └─ {}", c);
+                        }
+                    }
+                }
             }
-            all_healthy = false;
-        } else {
-            let containers: Vec<&str> = docker_check.lines().collect();
-            if containers.is_empty() || (containers.len() == 1 && containers[0].is_empty()) {
-                println!("✗ Нет контейнеров");
+            Ok((name, ib_ip, Err(e))) => {
+                println!("{} ({}): ✗ SSH ошибка: {}", name, ib_ip, e);
+                all_healthy = false;
+            }
+            Err(join_err) => {
+                println!("✗ Task join error: {}", join_err);
                 all_healthy = false;
-            } else {
-                println!("✓ {} контейнер(ов)", containers.len());
-                for c in containers {
-                    println!("  └─ {}", c);
-                }
             }
         }
     }
@@ -294,11 +389,12 @@ fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster
 }
 
 // ====================== Вспомогательные функции ======================
-async fn get_cluster_for_recipe(config: &Config, recipe_name: &str) -> Result<Cluster> {
-    let recipe_path = format!("recipies/{}/recipe.toml", recipe_name);
-    let _recipe: Recipe = toml::from_str(&fs::read_to_string(&recipe_path).await?)?;
-    let (cluster, _) = select_cluster(config, "auto")?;
-    Ok(cluster)
+async fn get_cluster_for_recipe(config: &Config, recipe_name: &str) -> Result<(Cluster, String)> {
+    let recipe_path = format!("recipes/{}/recipe.toml", recipe_name);
+    let recipe: Recipe = toml::from_str(&fs::read_to_string(&recipe_path).await?)?;
+    let cluster_choice = recipe.cluster.as_deref().unwrap_or("auto");
+    let (cluster, _) = select_cluster(config, cluster_choice)?;
+    Ok((cluster, recipe_name.to_string()))
 }
 
 fn select_cluster<'c>(config: &'c Config, cluster_name: &str) -> Result<(Cluster, &'c NcclConfig)> {
@@ -330,6 +426,7 @@ async fn start_cluster(
     config: &Config,
     cluster: &Cluster,
     nccl: &NcclConfig,
+    recipe_name: &str,
     recipe: &Recipe,
     dry_run: bool,
 ) -> Result<()> {
@@ -337,20 +434,51 @@ async fn start_cluster(
     let pp_size = recipe.pipeline_parallel_size.unwrap_or(0);
     let pp_size = if pp_size == 0 { nnodes } else { pp_size };
 
+    let recipe_clone = (*recipe).clone();
+
+    let mut handles = Vec::new();
     for node in &cluster.nodes {
         let is_head = node.rank == 0;
+        let cleanup_cmd = format!("docker rm -f sm121_{}_{} 2>/dev/null; true", recipe_name, node.rank);
+
+        let config = config.clone();
+        let cluster = cluster.clone();
+        let nccl = nccl.clone();
+        let node = node.clone();
+        let recipe = recipe_clone.clone();
+        let recipe_name = recipe_name.to_string();
+        let dry = dry_run;
+        let engine = recipe.engine.clone();
+
+        handles.push(tokio::spawn(async move {
+            // Cleanup leftover container
+            if !dry {
+                let _ = run_ssh(&config, &cluster, &node, cleanup_cmd, dry).await;
+            }
 
-        let docker_cmd =
-            build_docker_command(cluster, nccl, recipe, node, is_head, pp_size, nnodes);
-        let full_cmd = docker_cmd.join(" ");
+            let docker_cmd = build_docker_command(
+                &cluster, &nccl, &recipe, &recipe_name, &node, is_head, pp_size, nnodes,
+            );
+            let full_cmd = docker_cmd.join(" ");
 
-        run_ssh(config, cluster, node, full_cmd, dry_run).await?;
-        info!(
-            "✅ {} запущен на {} (rank {})",
-            recipe.engine.to_uppercase(),
-            node.name,
-            node.rank
-        );
+            let result = run_ssh_with_retry(&config, &cluster, &node, full_cmd, dry, 2).await;
+            (node.name, node.rank, engine, result)
+        }));
+    }
+
+    let results = futures::future::join_all(handles).await;
+    for result in results {
+        match result {
+            Ok((name, rank, engine, Ok(_output))) => {
+                info!("✅ {} запущен на {} (rank {})", engine.to_uppercase(), name, rank);
+            }
+            Ok((name, rank, engine, Err(e))) => {
+                anyhow::bail!("Failed to start {} on {} (rank {}): {}", engine.to_uppercase(), name, rank, e);
+            }
+            Err(join_err) => {
+                anyhow::bail!("Task join error: {}", join_err);
+            }
+        }
     }
 
     if !dry_run {
@@ -367,6 +495,7 @@ fn build_docker_command(
     cluster: &Cluster,
     nccl: &NcclConfig,
     recipe: &Recipe,
+    recipe_name: &str,
     node: &Node,
     is_head: bool,
     pp_size: u32,
@@ -386,11 +515,21 @@ fn build_docker_command(
     }
 
     if is_head {
-        cmd.push(format!("-p {}:{}", recipe.port, recipe.port));
+        cmd.push(format!("-p {}:{}" , recipe.port, recipe.port));
     }
 
+    // Container name and labels
+    cmd.push(format!("--name sm121_{}_{}", recipe_name, node.rank));
+    cmd.push(format!("--label sm121.recipe={}", recipe_name));
+    cmd.push(format!("--label sm121.rank={}", node.rank));
+
     cmd.push(recipe.docker_image.clone());
 
+    // Engine entrypoint (optional — empty for vLLM)
+    if !recipe.entrypoint.is_empty() {
+        cmd.extend(recipe.entrypoint.clone());
+    }
+
     cmd.push(recipe.model.clone());
     cmd.extend(recipe.args.clone());
     cmd.extend(vec![
@@ -411,13 +550,20 @@ fn build_docker_command(
     cmd
 }
 
-async fn ssh_session(config: &Config, node: &Node) -> Result<client::Handle<SimpleHandler>> {
+async fn ssh_session(config: &Config, node: &Node) -> Result<client::Handle<VerifyingHandler>> {
     let key = load_secret_key(&config.ssh_key_path)?;
     let ssh_config = russh::client::Config::default();
+    let expected_key = Arc::new(std::sync::Mutex::new(node.host_key.clone()));
+
+    if expected_key.lock().unwrap().is_none() {
+        tracing::warn!("No host_key set for {} — skipping verification (insecure)", node.name);
+    }
+
+    let handler = VerifyingHandler { expected_key };
     let mut session = russh::client::connect(
         Arc::new(ssh_config),
         (node.ssh_ip.as_str(), 22),
-        SimpleHandler,
+        handler,
     )
     .await?;
 
@@ -431,6 +577,47 @@ async fn ssh_session(config: &Config, node: &Node) -> Result<client::Handle<Simp
     Ok(session)
 }
 
+async fn fetch_host_key(config: &Config, node: &Node) -> Result<String> {
+    use std::sync::Mutex;
+
+    let captured = Arc::new(Mutex::new(None::<Vec<u8>>));
+    let captured_clone = captured.clone();
+
+    struct CaptureHandler {
+        captured: Arc<Mutex<Option<Vec<u8>>>>,
+    }
+    impl client::Handler for CaptureHandler {
+        type Error = russh::Error;
+        async fn check_server_key(&mut self, key: &russh::keys::PublicKey) -> Result<bool, Self::Error> {
+            *self.captured.lock().unwrap() = Some(key.to_string().into_bytes());
+            Ok(true)
+        }
+    }
+
+    let key = load_secret_key(&config.ssh_key_path)?;
+    let ssh_config = russh::client::Config::default();
+    let handler = CaptureHandler { captured: captured_clone };
+    let mut session = russh::client::connect(
+        Arc::new(ssh_config),
+        (node.ssh_ip.as_str(), 22),
+        handler,
+    )
+    .await?;
+
+    session
+        .authenticate_publickey(
+            &config.ssh_user,
+            PrivateKeyWithHashAlg::new(Arc::new(key), None),
+        )
+        .await?;
+
+    let blob = captured.lock().unwrap().take()
+        .context("No server key captured")?;
+    let hash = sha2::Sha256::digest(&blob);
+    let fp = format!("SHA256:{}", base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &hash));
+    Ok(fp)
+}
+
 async fn run_ssh(
     config: &Config,
     _cluster: &Cluster,
@@ -456,6 +643,36 @@ async fn run_ssh(
     Ok(output)
 }
 
+async fn run_ssh_with_retry(
+    config: &Config,
+    cluster: &Cluster,
+    node: &Node,
+    command: String,
+    dry_run: bool,
+    max_retries: u32,
+) -> Result<String> {
+    let mut last_err = None;
+    for attempt in 0..=max_retries {
+        match run_ssh(config, cluster, node, command.clone(), dry_run).await {
+            Ok(output) => return Ok(output),
+            Err(e) => {
+                tracing::warn!(
+                    "SSH attempt {}/{} to {} failed: {}",
+                    attempt + 1,
+                    max_retries + 1,
+                    node.name,
+                    e
+                );
+                last_err = Some(e);
+                if attempt < max_retries {
+                    sleep(Duration::from_secs(2u64.pow(attempt))).await;
+                }
+            }
+        }
+    }
+    Err(last_err.unwrap())
+}
+
 async fn run_ssh_follow(
     config: &Config,
     node: &Node,
@@ -479,13 +696,25 @@ async fn run_ssh_follow(
     Ok(())
 }
 
-struct SimpleHandler;
-impl client::Handler for SimpleHandler {
+struct VerifyingHandler {
+    expected_key: Arc<std::sync::Mutex<Option<String>>>,
+}
+
+impl client::Handler for VerifyingHandler {
     type Error = russh::Error;
     async fn check_server_key(
         &mut self,
-        _server_public_key: &russh::keys::PublicKey,
+        server_public_key: &russh::keys::PublicKey,
     ) -> Result<bool, Self::Error> {
+        let expected = self.expected_key.lock().unwrap().clone();
+        if let Some(ref expected_fp) = expected {
+            let key_bytes = server_public_key.to_string().into_bytes();
+            let hash = Sha256::digest(&key_bytes);
+            let fp = format!("SHA256:{}", STANDARD.encode(&hash));
+            if &fp != expected_fp {
+                return Err(russh::Error::Kex);
+            }
+        }
         Ok(true)
     }
 }