Ver Fonte

feat: add 'init-host-keys' command to auto-write SSH host_keys

Automates the workflow previously done by 'cargo run health --fingerprint'
+ manual editing of config.toml.  For every cluster+node, fetches the
SSH host key fingerprint via fetch_host_key() and rewrites the matching
inline-table node line in config.toml with the host_key field.

Behaviour:
  cargo run init-host-keys                          # all clusters
  cargo run init-host-keys --cluster 2nodes-dac     # single cluster

Implementation details:
- src/host_keys.rs holds the logic in a dedicated module with the
  text-rewrite helpers as separate testable units.
- update_node_line handles two formats: a line whose host_key field
  is missing (insert before ' }') or already present (swap the value).
- update_host_keys_text walks the file line by line, only touching
  lines whose 'name = "X"' matches an entry in the discovered set.
  Other content (clusters without host_key, NCCL sections, etc.) is
  left untouched.
- Idempotent: re-running on a config that already has the right
  fingerprints reports 'No changes needed'.

Verified against the live cluster:
  init-host-keys -> 4 fingerprints discovered
                    config.toml rewritten for both clusters
  re-running     -> 'No changes needed'
  cargo run serve -> runs without any 'No host_key' warnings
Alexander Korolev há 4 meses atrás
pai
commit
df5e3a701e
3 ficheiros alterados com 197 adições e 0 exclusões
  1. 9 0
      src/cli.rs
  2. 184 0
      src/host_keys.rs
  3. 4 0
      src/main.rs

+ 9 - 0
src/cli.rs

@@ -77,6 +77,15 @@ pub enum Commands {
     /// Управление tmux-сессиями (альтернатива Docker)
     #[command(subcommand)]
     Tmux(TmuxCommands),
+
+    /// Auto-discover host keys and write them into config.toml
+    InitHostKeys {
+        #[arg(
+            long,
+            help = "Имя кластера (например, 4nodes-switch). Если не указано — все."
+        )]
+        cluster: Option<String>,
+    },
 }
 
 #[derive(Subcommand)]

+ 184 - 0
src/host_keys.rs

@@ -0,0 +1,184 @@
+//! Auto-discover SSH host key fingerprints for every node in the config
+//! and rewrite `config.toml` with `host_key = "SHA256:..."` entries.
+
+use anyhow::{Context, Result};
+use std::fs;
+
+use crate::config::Config;
+use crate::ssh::fetch_host_key;
+
+/// Discover host keys and update the config file.  When `cluster_filter`
+/// is `Some`, only that cluster is processed.
+pub async fn init_host_keys(
+    config: &Config,
+    cluster_filter: Option<&str>,
+    config_path: &str,
+) -> Result<()> {
+    let updates = collect_updates(config, cluster_filter).await?;
+
+    if updates.is_empty() {
+        println!("(no nodes matched)");
+        return Ok(());
+    }
+
+    println!("Discovered host keys:");
+    for (name, fp) in &updates {
+        println!("  {} -> {}", name, fp);
+    }
+
+    let content = fs::read_to_string(config_path)
+        .with_context(|| format!("Failed to read {}", config_path))?;
+    let new_content = update_host_keys_text(&content, &updates);
+
+    if content == new_content {
+        println!("\nNo changes needed in {}", config_path);
+    } else {
+        fs::write(config_path, &new_content)
+            .with_context(|| format!("Failed to write {}", config_path))?;
+        println!("\nUpdated {}", config_path);
+    }
+    Ok(())
+}
+
+/// Walk the config and return `(node_name, fingerprint)` for every node.
+async fn collect_updates(
+    config: &Config,
+    cluster_filter: Option<&str>,
+) -> Result<Vec<(String, String)>> {
+    let mut out = Vec::new();
+    for cluster in &config.clusters {
+        if let Some(name) = cluster_filter
+            && cluster.name != name
+        {
+            continue;
+        }
+        for node in &cluster.nodes {
+            match fetch_host_key(config, node).await {
+                Ok(fp) => out.push((node.name.clone(), fp)),
+                Err(e) => eprintln!("warning: failed to fetch key for {}: {}", node.name, e),
+            }
+        }
+    }
+    Ok(out)
+}
+
+/// Walk `content` line by line and rewrite any node definition whose
+/// `name = "..."` matches an entry in `updates`, setting or updating the
+/// `host_key` field.
+fn update_host_keys_text(content: &str, updates: &[(String, String)]) -> String {
+    let mut lines: Vec<String> = content.lines().map(String::from).collect();
+    for (name, fp) in updates {
+        let marker = format!("name = \"{}\"", name);
+        for line in &mut lines {
+            if is_node_line(line) && line.contains(&marker) {
+                *line = update_node_line(line, fp);
+            }
+        }
+    }
+    if content.ends_with('\n') {
+        lines.join("\n") + "\n"
+    } else {
+        lines.join("\n")
+    }
+}
+
+/// Heuristic: a node definition starts (after indent) with `{`.
+fn is_node_line(line: &str) -> bool {
+    line.trim_start().starts_with('{')
+}
+
+/// Set or replace the `host_key = "..."` field in a single node line.
+///
+/// If the field is already present, its value is swapped for `fp`.
+/// Otherwise the field is inserted before the trailing `}` with a comma
+/// if needed.
+fn update_node_line(line: &str, fp: &str) -> String {
+    let marker = "host_key = \"";
+    if let Some(start) = line.find(marker) {
+        let value_start = start + marker.len();
+        if let Some(rel_end) = line[value_start..].find('"') {
+            let value_end = value_start + rel_end;
+            let mut result = String::with_capacity(line.len());
+            result.push_str(&line[..value_start]);
+            result.push_str(fp);
+            result.push_str(&line[value_end..]);
+            return result;
+        }
+    }
+
+    // No host_key yet — insert before the closing `}`.
+    let close_idx = line.rfind('}').expect("node line must contain '}'");
+    let before_close = &line[..close_idx];
+    let after_close = &line[close_idx..];
+    let trimmed = before_close.trim_end();
+    let needs_comma = !trimmed.is_empty() && !trimmed.ends_with(',');
+    // Original line either ends with `}` or `},`. Detect which so we can
+    // keep the formatting consistent (space before `}` only when a comma
+    // follows it).
+    let has_trailing_comma = after_close.starts_with("},");
+    let mut result = String::with_capacity(line.len() + fp.len() + 20);
+    result.push_str(trimmed);
+    if needs_comma {
+        result.push(',');
+    }
+    result.push(' ');
+    result.push_str(&format!("host_key = \"{}\"", fp));
+    if has_trailing_comma {
+        result.push(' ');
+    }
+    result.push('}');
+    if has_trailing_comma {
+        result.push(',');
+    }
+    result
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn adds_host_key_when_missing() {
+        let line =
+            r#"    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },"#;
+        let result = update_node_line(line, "SHA256:abc");
+        assert_eq!(
+            result,
+            r#"    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0, host_key = "SHA256:abc" },"#
+        );
+    }
+
+    #[test]
+    fn replaces_host_key_value() {
+        let line = r#"    { name = "x", rank = 0, host_key = "OLD" },"#;
+        let result = update_node_line(line, "NEW");
+        assert_eq!(result, r#"    { name = "x", rank = 0, host_key = "NEW" },"#);
+    }
+
+    #[test]
+    fn update_host_keys_text_targets_only_named_lines() {
+        let content = "\
+[[clusters]]
+name = \"c\"
+nodes = [
+    { name = \"a\", rank = 0 },
+    { name = \"b\", rank = 1 },
+]
+";
+        let updates = vec![("a".to_string(), "SHA256:newA".to_string())];
+        let result = update_host_keys_text(content, &updates);
+        // a was updated to include host_key
+        assert!(result.contains("name = \"a\", rank = 0, host_key = \"SHA256:newA\""));
+        // b was NOT updated — original line preserved, no host_key added
+        assert!(result.contains("name = \"b\", rank = 1 }"));
+        assert!(!result.contains("name = \"b\", rank = 1, host_key"));
+    }
+
+    #[test]
+    fn update_node_line_with_no_trailing_comma() {
+        // Last field already ends with comma — no extra comma needed.
+        let line = r#"    { name = "x", rank = 0 }"#;
+        let result = update_node_line(line, "FP");
+        assert_eq!(result, r#"    { name = "x", rank = 0, host_key = "FP"}"#);
+    }
+}

+ 4 - 0
src/main.rs

@@ -6,6 +6,7 @@ mod config;
 mod docker;
 mod doctor;
 mod health;
+mod host_keys;
 mod ssh;
 mod tmux;
 
@@ -87,6 +88,9 @@ async fn main() -> Result<()> {
         Commands::Tmux(tmux_cmd) => {
             crate::tmux::handle_tmux(&config, tmux_cmd, dry_run).await?;
         }
+        Commands::InitHostKeys { cluster } => {
+            host_keys::init_host_keys(&config, cluster.as_deref(), &cli.config).await?;
+        }
     }
 
     Ok(())