From 94ac9ec8670ea1a2b76378c1cf31855b05cf531a Mon Sep 17 00:00:00 2001 From: Aly Raffauf Date: Wed, 5 Aug 2026 00:38:39 -0400 Subject: [PATCH] use openssh bindings for ssh ops; add --accept-new-host-keys for onboardings --- Cargo.lock | 78 +++++++++++++++++++++++++++++++++++++- Cargo.toml | 1 + README.md | 1 + src/cli.rs | 4 ++ src/main.rs | 12 +++++- src/nix.rs | 100 ++++++++++++++++++++++++++----------------------- src/process.rs | 36 +++++++++++++++--- src/ssh.rs | 85 +++++++++++++++++++++++++++++++++++++++++ 8 files changed, 262 insertions(+), 55 deletions(-) create mode 100644 src/ssh.rs diff --git a/Cargo.lock b/Cargo.lock index d60e890..a70e07d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -73,6 +73,12 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + [[package]] name = "blzrd" version = "0.1.0" @@ -82,6 +88,7 @@ dependencies = [ "env_logger", "futures", "log", + "openssh", "serde", "serde_json", "tokio", @@ -158,7 +165,7 @@ version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" dependencies = [ - "bitflags", + "bitflags 1.3.2", "defmt-macros", ] @@ -216,6 +223,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "futures" version = "0.3.33" @@ -386,6 +399,12 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "log" version = "0.4.33" @@ -421,6 +440,20 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "openssh" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d534c4bfecb0ed71dea4db444a5922a294d15cf40e700548f27295e1feb0ef18" +dependencies = [ + "libc", + "once_cell", + "shell-escape", + "tempfile", + "thiserror", + "tokio", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -495,6 +528,19 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags 2.13.1", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -544,6 +590,12 @@ dependencies = [ "zmij", ] +[[package]] +name = "shell-escape" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45bb67a18fa91266cc7807181f62f9178a6873bfad7dc788c42e6430db40184f" + [[package]] name = "signal-hook-registry" version = "1.4.8" @@ -560,6 +612,16 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + [[package]] name = "strsim" version = "0.11.1" @@ -588,6 +650,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + [[package]] name = "thiserror" version = "2.0.19" @@ -619,6 +694,7 @@ dependencies = [ "mio", "pin-project-lite", "signal-hook-registry", + "socket2", "tokio-macros", "windows-sys", ] diff --git a/Cargo.toml b/Cargo.toml index d716626..59edd2f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,4 +23,5 @@ anyhow = "1" futures = "0.3" log = "0.4" env_logger = "0.11" +openssh = "0.11" uuid = { version = "1", features = ["v4"] } diff --git a/README.md b/README.md index 0f80d0b..1c70172 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ blzrd switch --flake - `--flake `: Flake URL or path (env: `FLAKE`, default: `.`). Global; can be placed before or after the subcommand. - `--build-host `: Build closures on this remote host instead of locally (default: `localhost`). +- `--accept-new-host-keys`: Accept and record previously unknown SSH host keys. Changed keys are still rejected. - `--skip `: Comma-separated nodes to skip. - positional `nodes`: Nodes to deploy (default: all). Accepts a comma- or space-separated list. - `RUST_LOG=debug`: Enable debug output showing every subprocess command and its output. diff --git a/src/cli.rs b/src/cli.rs index d1a94a9..497c597 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -28,6 +28,10 @@ pub struct CommonArgs { /// Build closures on a remote host instead of locally. #[arg(long, default_value = "localhost")] pub build_host: String, + + /// Accept and record previously unknown SSH host keys. + #[arg(long)] + pub accept_new_host_keys: bool, } #[derive(Subcommand, Debug)] diff --git a/src/main.rs b/src/main.rs index 3250d36..ddf69fd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,12 +3,14 @@ mod models; mod nix; mod op; mod process; +mod ssh; use std::collections::HashSet; use futures::future::join_all; use crate::cli::{Command, CommonArgs}; +use crate::ssh::HostKeyPolicy; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -96,12 +98,18 @@ async fn run_deploy( log::info!("Operation {} is valid for all nodes", op); + let host_key_policy = if common.accept_new_host_keys { + HostKeyPolicy::AcceptNew + } else { + HostKeyPolicy::Strict + }; + // Build each node's closure locally or on a remote builder. log::info!("Building {} output(s)...", jobs.len()); let mut outs: std::collections::HashMap = std::collections::HashMap::with_capacity(jobs.len()); for (name, spec) in &jobs { - let (out, _debug) = nix::build_closure(spec, &common.build_host).await?; + let (out, _debug) = nix::build_closure(spec, &common.build_host, host_key_policy).await?; log::info!(" ✔ {name} ({})", spec.system); outs.insert(name.clone(), out); } @@ -116,7 +124,7 @@ async fn run_deploy( let name = name.clone(); let spec = spec.clone(); async move { - let result = nix::deploy_closure(&spec, &out, op).await; + let result = nix::deploy_closure(&spec, &out, op, host_key_policy).await; (name, spec, result) } }) diff --git a/src/nix.rs b/src/nix.rs index 4d85b0c..4df4fb0 100644 --- a/src/nix.rs +++ b/src/nix.rs @@ -8,7 +8,8 @@ use uuid::Uuid; use crate::models::{BuildResult, DebugInfo, JobSpec, NixEvalJobsResult, SystemType}; use crate::op::Operation; -use crate::process::{get_config_attr, run, run_json}; +use crate::process::{get_config_attr, run, run_json, run_json_with_env, run_with_env}; +use crate::ssh::{run as run_ssh, HostKeyPolicy}; /// Evaluate the flake's `blzrd.nodes` and return enriched `JobSpec`s. pub async fn eval_deployments(cfg: &str) -> Result<(HashMap, Vec)> { @@ -169,7 +170,11 @@ async fn build_job_specs( } /// Build a job's derivation and return the `out` store path. -pub async fn build_closure(spec: &JobSpec, build_host: &str) -> Result<(String, DebugInfo)> { +pub async fn build_closure( + spec: &JobSpec, + build_host: &str, + host_key_policy: HostKeyPolicy, +) -> Result<(String, DebugInfo)> { let drv = format!("{}^*", spec.drv_path); if build_host == "localhost" { @@ -185,16 +190,19 @@ pub async fn build_closure(spec: &JobSpec, build_host: &str) -> Result<(String, // Remote builder branch. let store = format!("ssh-ng://{build_host}"); + let nix_ssh_options = nix_ssh_options(host_key_policy); + let nix_env = [("NIX_SSHOPTS", nix_ssh_options.as_str())]; // 1. Copy the derivation to the builder. - let (_out, _debug) = run("nix", &["copy", "--to", &store, &spec.drv_path]) + let (_out, _debug) = run_with_env("nix", &["copy", "--to", &store, &spec.drv_path], &nix_env) .await .with_context(|| format!("copy to {build_host}"))?; // 2. Build on the builder. - let (results, debug): (Vec, _) = run_json( + let (results, debug): (Vec, _) = run_json_with_env( "nix", &["build", "--no-link", "--json", "--store", &store, &drv], + &nix_env, ) .await .with_context(|| format!("build on {build_host}"))?; @@ -206,16 +214,20 @@ pub async fn build_closure(spec: &JobSpec, build_host: &str) -> Result<(String, .context("build result missing 'out' output")?; // 3. Copy the out path back. - let (_out, _debug2) = run("nix", &["copy", "--from", &store, &out, "--no-check-sigs"]) - .await - .with_context(|| format!("copy from {build_host}"))?; + let (_out, _debug2) = run_with_env( + "nix", + &["copy", "--from", &store, &out, "--no-check-sigs"], + &nix_env, + ) + .await + .with_context(|| format!("copy from {build_host}"))?; Ok((out, debug)) } /// Poll a remote `blzrd-activate-*` transient systemd unit until it reaches a /// terminal state, then return. Exponential backoff 5 -> 60 seconds, 5-minute deadline. -async fn poll_activation(target: &str, unit: &str) -> Result<()> { +async fn poll_activation(target: &str, unit: &str, host_key_policy: HostKeyPolicy) -> Result<()> { let mut sleep_val = Duration::from_secs(5); let deadline = Instant::now() + Duration::from_secs(300); @@ -224,17 +236,9 @@ async fn poll_activation(target: &str, unit: &str) -> Result<()> { anyhow::bail!("activation on {target}: timed out waiting for {unit} unit"); } - let args: Vec<&str> = vec![ - "-o", - "ConnectTimeout=5", - "-o", - "BatchMode=yes", - target, - "systemctl", - "show", - unit, - ]; - let result = run("ssh", &args).await.ok(); + let result = run_ssh(target, "systemctl", &["show", unit], host_key_policy) + .await + .ok(); let mut sub_state: Option = None; let mut exec_status: Option = None; @@ -283,7 +287,12 @@ async fn poll_activation(target: &str, unit: &str) -> Result<()> { Ok(()) } -pub async fn deploy_closure(spec: &JobSpec, out_path: &str, op: Operation) -> Result { +pub async fn deploy_closure( + spec: &JobSpec, + out_path: &str, + op: Operation, + host_key_policy: HostKeyPolicy, +) -> Result { let target = spec.target(); let path = out_path.to_string(); @@ -296,21 +305,15 @@ pub async fn deploy_closure(spec: &JobSpec, out_path: &str, op: Operation) -> Re match (sys, op) { (SystemType::Darwin, Operation::Switch) => { cmds.push(vec![ - "ssh".into(), - target.clone(), - "PATH=/run/current-system/sw/bin:$PATH".into(), - "sudo".into(), - "nix-env".into(), + "/run/current-system/sw/bin/sudo".into(), + "/run/current-system/sw/bin/nix-env".into(), "-p".into(), "/nix/var/nix/profiles/system".into(), "--set".into(), path.clone(), ]); cmds.push(vec![ - "ssh".into(), - target.clone(), - "PATH=/run/current-system/sw/bin:$PATH".into(), - "sudo".into(), + "/run/current-system/sw/bin/sudo".into(), format!("{path}/activate"), ]); } @@ -324,20 +327,16 @@ pub async fn deploy_closure(spec: &JobSpec, out_path: &str, op: Operation) -> Re (SystemType::Nixos, Operation::Switch) => { cmds.push(vec![ - "ssh".into(), - target.clone(), - "sudo".into(), - "nix-env".into(), + "/run/current-system/sw/bin/sudo".into(), + "/run/current-system/sw/bin/nix-env".into(), "-p".into(), "/nix/var/nix/profiles/system".into(), "--set".into(), path.clone(), ]); cmds.push(vec![ - "ssh".into(), - target.clone(), - "sudo".into(), - "systemd-run".into(), + "/run/current-system/sw/bin/sudo".into(), + "/run/current-system/sw/bin/systemd-run".into(), "--unit".into(), unit.clone(), "--remain-after-exit".into(), @@ -350,19 +349,15 @@ pub async fn deploy_closure(spec: &JobSpec, out_path: &str, op: Operation) -> Re (SystemType::Nixos, Operation::Boot) => { cmds.push(vec![ - "ssh".into(), - target.clone(), - "sudo".into(), - "nix-env".into(), + "/run/current-system/sw/bin/sudo".into(), + "/run/current-system/sw/bin/nix-env".into(), "-p".into(), "/nix/var/nix/profiles/system".into(), "--set".into(), path.clone(), ]); cmds.push(vec![ - "ssh".into(), - target.clone(), - "sudo".into(), + "/run/current-system/sw/bin/sudo".into(), format!("{path}/bin/switch-to-configuration"), op.to_string(), ]); @@ -370,7 +365,9 @@ pub async fn deploy_closure(spec: &JobSpec, out_path: &str, op: Operation) -> Re } // 1. Copy the closure to the target. - let (_out, debug) = run( + let nix_ssh_options = nix_ssh_options(host_key_policy); + let nix_env = [("NIX_SSHOPTS", nix_ssh_options.as_str())]; + let (_out, debug) = run_with_env( "nix", &[ "copy", @@ -379,6 +376,7 @@ pub async fn deploy_closure(spec: &JobSpec, out_path: &str, op: Operation) -> Re &path, "--no-check-sigs", ], + &nix_env, ) .await .with_context(|| format!("copy to {target}"))?; @@ -386,16 +384,24 @@ pub async fn deploy_closure(spec: &JobSpec, out_path: &str, op: Operation) -> Re // 2. Run each activation command in order. for cmd in &cmds { let args: Vec<&str> = cmd[1..].iter().map(String::as_str).collect(); - let (_out, _d) = run(&cmd[0], &args) + let (_out, _d) = run_ssh(&target, &cmd[0], &args, host_key_policy) .await .with_context(|| format!("activation on {target}"))?; } if matches!((sys, op), (SystemType::Nixos, Operation::Switch)) { - poll_activation(&target, &unit) + poll_activation(&target, &unit, host_key_policy) .await .with_context(|| format!("activation on {target}"))?; } Ok(debug) } + +fn nix_ssh_options(host_key_policy: HostKeyPolicy) -> String { + let host_key_option = host_key_policy.nix_ssh_option(); + match std::env::var("NIX_SSHOPTS") { + Ok(existing) if !existing.trim().is_empty() => format!("{existing} {host_key_option}"), + _ => host_key_option.to_string(), + } +} diff --git a/src/process.rs b/src/process.rs index 52772e7..6d4c1ff 100644 --- a/src/process.rs +++ b/src/process.rs @@ -5,10 +5,23 @@ use tokio::process::Command; use crate::models::DebugInfo; pub async fn run(cmd: &str, args: &[&str]) -> Result<(Vec, DebugInfo)> { + run_with_env(cmd, args, &[]).await +} + +pub async fn run_with_env( + cmd: &str, + args: &[&str], + envs: &[(&str, &str)], +) -> Result<(Vec, DebugInfo)> { let display = format!("{} {}", cmd, args.join(" ")); - let output = Command::new(cmd) - .args(args) + let mut command = Command::new(cmd); + command.args(args); + for (key, value) in envs { + command.env(key, value); + } + + let output = command .output() .await .with_context(|| format!("failed to spawn {display}"))?; @@ -40,10 +53,23 @@ pub async fn run(cmd: &str, args: &[&str]) -> Result<(Vec, DebugInfo)> { /// Run `cmd args...` and deserialize the JSON output into a struct of type `T`. pub async fn run_json(cmd: &str, args: &[&str]) -> Result<(T, DebugInfo)> { + run_json_with_env(cmd, args, &[]).await +} + +pub async fn run_json_with_env( + cmd: &str, + args: &[&str], + envs: &[(&str, &str)], +) -> Result<(T, DebugInfo)> { let display = format!("{} {}", cmd, args.join(" ")); - let output = Command::new(cmd) - .args(args) + let mut command = Command::new(cmd); + command.args(args); + for (key, value) in envs { + command.env(key, value); + } + + let output = command .output() .await .with_context(|| format!("failed to spawn {display}"))?; @@ -76,7 +102,7 @@ pub async fn run_json(cmd: &str, args: &[&str]) -> Result<( } /// Emit a `DebugInfo` at the debug log level (only visible with `--debug`). -fn log_debug(debug: &DebugInfo) { +pub(crate) fn log_debug(debug: &DebugInfo) { log::debug!("$ {}", debug.command); if !debug.std_out.trim().is_empty() { log::debug!("stdout:\n{}", debug.std_out.trim()); diff --git a/src/ssh.rs b/src/ssh.rs new file mode 100644 index 0000000..c9728e0 --- /dev/null +++ b/src/ssh.rs @@ -0,0 +1,85 @@ +use std::time::Duration; + +use anyhow::{Context, Result}; +use openssh::{KnownHosts, SessionBuilder}; + +use crate::models::DebugInfo; +use crate::process::log_debug; + +const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); + +#[derive(Debug, Clone, Copy)] +pub enum HostKeyPolicy { + Strict, + AcceptNew, +} + +impl HostKeyPolicy { + fn known_hosts(self) -> KnownHosts { + match self { + Self::Strict => KnownHosts::Strict, + Self::AcceptNew => KnownHosts::Add, + } + } + + pub fn nix_ssh_option(self) -> &'static str { + match self { + Self::Strict => "-o StrictHostKeyChecking=yes", + Self::AcceptNew => "-o StrictHostKeyChecking=accept-new", + } + } +} + +/// Run one remote command over a fresh SSH session. +pub async fn run( + target: &str, + program: &str, + args: &[&str], + host_key_policy: HostKeyPolicy, +) -> Result<(Vec, DebugInfo)> { + let display = format!("ssh {target} {program} {}", args.join(" ")); + + let mut builder = SessionBuilder::default(); + builder + .known_hosts_check(host_key_policy.known_hosts()) + .connect_timeout(CONNECT_TIMEOUT); + + let session = builder + .connect(target) + .await + .with_context(|| format!("failed to connect to {target}"))?; + + let output_result = session + .command(program) + .args(args) + .output() + .await + .with_context(|| format!("failed to run {display}")); + + if let Err(error) = session.close().await { + log::debug!("failed to close SSH session to {target}: {error}"); + } + + let output = output_result?; + + let std_out = String::from_utf8_lossy(&output.stdout).into_owned(); + let std_err = String::from_utf8_lossy(&output.stderr).into_owned(); + let debug = DebugInfo { + command: display.clone(), + std_out, + std_err: std_err.clone(), + }; + + log_debug(&debug); + + if !output.status.success() { + let detail = if std_err.trim().is_empty() { + &debug.std_out + } else { + &debug.std_err + }; + anyhow::bail!("{display} failed: {}", detail.trim()); + } + + Ok((output.stdout, debug)) +} -- 2.51.2