feat: ssh workflow
This commit is contained in:
+352
@@ -0,0 +1,352 @@
|
||||
use std::fs::File;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use anyhow::{Context, Result, bail};
|
||||
use russh::ChannelMsg;
|
||||
use russh::client::{self, AuthResult, Handle};
|
||||
use russh::keys::known_hosts;
|
||||
use russh::keys::{PrivateKeyWithHashAlg, PublicKey, load_secret_key};
|
||||
use ssh2_config::{HostParams, ParseRule, SshConfig};
|
||||
|
||||
pub struct Target {
|
||||
pub host_alias: String,
|
||||
pub user: Option<String>,
|
||||
pub port: Option<u16>,
|
||||
pub key_path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct HostConfig {
|
||||
host_name: Option<String>,
|
||||
user: Option<String>,
|
||||
port: Option<u16>,
|
||||
identity_file: Option<Vec<PathBuf>>,
|
||||
}
|
||||
|
||||
impl From<HostParams> for HostConfig {
|
||||
fn from(params: HostParams) -> Self {
|
||||
Self {
|
||||
host_name: params.host_name,
|
||||
user: params.user,
|
||||
port: params.port,
|
||||
identity_file: params.identity_file,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Resolved {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub user: String,
|
||||
pub key_path: PathBuf,
|
||||
}
|
||||
|
||||
fn resolve_with(
|
||||
target: &Target,
|
||||
ssh_cfg: &HostConfig,
|
||||
fallback_user: Option<&str>,
|
||||
) -> Result<Resolved> {
|
||||
let host = ssh_cfg
|
||||
.host_name
|
||||
.clone()
|
||||
.unwrap_or_else(|| target.host_alias.clone());
|
||||
let port = target.port.or(ssh_cfg.port).unwrap_or(22);
|
||||
let user = target
|
||||
.user
|
||||
.clone()
|
||||
.or_else(|| ssh_cfg.user.clone())
|
||||
.or_else(|| fallback_user.map(str::to_string))
|
||||
.context("cannot determine SSH user: set 'user' in deploy.toml or ~/.ssh/config")?;
|
||||
let key_path = target
|
||||
.key_path
|
||||
.clone()
|
||||
.or_else(|| {
|
||||
ssh_cfg
|
||||
.identity_file
|
||||
.as_ref()
|
||||
.and_then(|files| files.first().cloned())
|
||||
})
|
||||
.unwrap_or_else(default_key);
|
||||
Ok(Resolved {
|
||||
host,
|
||||
port,
|
||||
user,
|
||||
key_path,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve(target: &Target) -> Result<Resolved> {
|
||||
resolve_with(
|
||||
target,
|
||||
&query_ssh_config(&target.host_alias),
|
||||
env_user().as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn default_key() -> PathBuf {
|
||||
home::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ssh")
|
||||
.join("id_rsa")
|
||||
}
|
||||
|
||||
fn env_user() -> Option<String> {
|
||||
std::env::var("USERNAME")
|
||||
.or_else(|_| std::env::var("USER"))
|
||||
.ok()
|
||||
}
|
||||
|
||||
fn ssh_config_path() -> Option<PathBuf> {
|
||||
let path = home::home_dir()?.join(".ssh").join("config");
|
||||
path.is_file().then_some(path)
|
||||
}
|
||||
|
||||
fn query_ssh_config(alias: &str) -> HostConfig {
|
||||
let Some(path) = ssh_config_path() else {
|
||||
return HostConfig::default();
|
||||
};
|
||||
let file = match File::open(&path) {
|
||||
Ok(file) => file,
|
||||
Err(_) => return HostConfig::default(),
|
||||
};
|
||||
let mut reader = std::io::BufReader::new(file);
|
||||
match SshConfig::default().parse(&mut reader, ParseRule::ALLOW_UNKNOWN_FIELDS) {
|
||||
Ok(config) => config.query(alias).into(),
|
||||
Err(err) => {
|
||||
eprintln!("warning: cannot parse {}: {err}", path.display());
|
||||
HostConfig::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ClientHandler {
|
||||
host: String,
|
||||
port: u16,
|
||||
}
|
||||
|
||||
impl client::Handler for ClientHandler {
|
||||
type Error = russh::Error;
|
||||
|
||||
async fn check_server_key(&mut self, server_key: &PublicKey) -> Result<bool, Self::Error> {
|
||||
let key = server_key;
|
||||
match known_hosts::check_known_hosts(&self.host, self.port, key) {
|
||||
Ok(true) => Ok(true),
|
||||
Ok(false) => {
|
||||
eprintln!(
|
||||
"note: '{}' is not in known_hosts yet - trusting and remembering it (TOFU)",
|
||||
format_host_port(&self.host, self.port)
|
||||
);
|
||||
if let Err(err) = known_hosts::learn_known_hosts(&self.host, self.port, key) {
|
||||
eprintln!("warning: could not update known_hosts: {err}");
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
Err(russh::keys::Error::KeyChanged { line }) => {
|
||||
eprintln!(
|
||||
"SECURITY WARNING: host key for '{}' changed (known_hosts line {})! \
|
||||
If this is expected (VPS reinstall), remove that line from \
|
||||
~/.ssh/known_hosts and retry.",
|
||||
format_host_port(&self.host, self.port),
|
||||
line + 1
|
||||
);
|
||||
Ok(false)
|
||||
}
|
||||
Err(other) => {
|
||||
eprintln!("warning: cannot verify known_hosts entry: {other}");
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn format_host_port(host: &str, port: u16) -> String {
|
||||
if port == 22 {
|
||||
host.to_string()
|
||||
} else {
|
||||
format!("[{host}]:{port}")
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Session {
|
||||
handle: Handle<ClientHandler>,
|
||||
}
|
||||
|
||||
fn runtime() -> &'static tokio::runtime::Runtime {
|
||||
static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
|
||||
RT.get_or_init(|| {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("failed to start async runtime for SSH")
|
||||
})
|
||||
}
|
||||
|
||||
pub fn connect(resolved: &Resolved) -> Result<Session> {
|
||||
println!(
|
||||
"connecting to {}@{} (key {}) ...",
|
||||
resolved.user,
|
||||
format_host_port(&resolved.host, resolved.port),
|
||||
resolved.key_path.display()
|
||||
);
|
||||
runtime().block_on(async {
|
||||
let handler = ClientHandler {
|
||||
host: resolved.host.clone(),
|
||||
port: resolved.port,
|
||||
};
|
||||
let mut handle = client::connect(
|
||||
Arc::new(client::Config::default()),
|
||||
(resolved.host.as_str(), resolved.port),
|
||||
handler,
|
||||
)
|
||||
.await
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"SSH connection/handshake with {} failed",
|
||||
format_host_port(&resolved.host, resolved.port)
|
||||
)
|
||||
})?;
|
||||
authenticate(&mut handle, resolved).await?;
|
||||
Ok(Session { handle })
|
||||
})
|
||||
}
|
||||
|
||||
async fn authenticate(handle: &mut Handle<ClientHandler>, r: &Resolved) -> Result<()> {
|
||||
let secret = load_secret_key(&r.key_path, None).map_err(|err| {
|
||||
anyhow::anyhow!(
|
||||
"cannot load private key {}: {err} (passphrase-protected keys are not supported yet)",
|
||||
r.key_path.display()
|
||||
)
|
||||
})?;
|
||||
let best_hash = handle
|
||||
.best_supported_rsa_hash()
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.flatten();
|
||||
let key_alg = PrivateKeyWithHashAlg::new(Arc::new(secret), best_hash);
|
||||
match handle
|
||||
.authenticate_publickey(&r.user, key_alg)
|
||||
.await
|
||||
.context("publickey authentication exchange failed")?
|
||||
{
|
||||
AuthResult::Success => Ok(()),
|
||||
AuthResult::Failure { .. } => bail!(
|
||||
"SSH auth rejected for '{}' using key {} (is the public part in authorized_keys on the VPS?)",
|
||||
r.user,
|
||||
r.key_path.display()
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub fn exec(&mut self, cmd: &str) -> Result<()> {
|
||||
println!("[Remote] running: {cmd} ...");
|
||||
runtime().block_on(async {
|
||||
use std::io::Write;
|
||||
|
||||
let mut channel = self
|
||||
.handle
|
||||
.channel_open_session()
|
||||
.await
|
||||
.context("failed to open SSH channel")?;
|
||||
channel
|
||||
.exec(true, cmd)
|
||||
.await
|
||||
.with_context(|| format!("failed to start '{cmd}' on the remote host"))?;
|
||||
|
||||
let mut status: Option<u32> = None;
|
||||
let mut stderr = String::new();
|
||||
while let Some(msg) = channel.wait().await {
|
||||
match msg {
|
||||
ChannelMsg::Data { data } => {
|
||||
print!("{}", String::from_utf8_lossy(&data));
|
||||
let _ = std::io::stdout().flush();
|
||||
}
|
||||
ChannelMsg::ExtendedData { data, ext: 1 } => {
|
||||
stderr.push_str(&String::from_utf8_lossy(&data));
|
||||
}
|
||||
ChannelMsg::ExitStatus { exit_status } => status = Some(exit_status),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let status =
|
||||
status.context("remote command finished without reporting an exit status")?;
|
||||
if status != 0 {
|
||||
let tail = if stderr.trim().is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("\nstderr:\n{stderr}")
|
||||
};
|
||||
bail!("[Remote] running: {cmd} ... FAILED (exit code {status}){tail}");
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn target(user: Option<&str>, port: Option<u16>) -> Target {
|
||||
Target {
|
||||
host_alias: "my-vps".into(),
|
||||
user: user.map(str::to_string),
|
||||
port,
|
||||
key_path: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_values_override_ssh_config() {
|
||||
let hc = HostConfig {
|
||||
host_name: Some("10.0.0.5".into()),
|
||||
user: Some("cfg-user".into()),
|
||||
port: Some(2222),
|
||||
identity_file: Some(vec![PathBuf::from("/cfg/key")]),
|
||||
};
|
||||
let t = Target {
|
||||
host_alias: "my-vps".into(),
|
||||
user: Some("toml-user".into()),
|
||||
port: Some(2200),
|
||||
key_path: Some(PathBuf::from("~/custom_key")),
|
||||
};
|
||||
let r = resolve_with(&t, &hc, Some("os-user")).unwrap();
|
||||
assert_eq!(r.user, "toml-user");
|
||||
assert_eq!(r.port, 2200);
|
||||
assert!(r.key_path.starts_with("~"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ssh_config_fills_gaps() {
|
||||
let hc = HostConfig {
|
||||
host_name: Some("10.0.0.5".into()),
|
||||
user: Some("cfg-user".into()),
|
||||
port: Some(2222),
|
||||
identity_file: Some(vec![PathBuf::from("/cfg/key")]),
|
||||
};
|
||||
let r = resolve_with(&target(None, None), &hc, Some("os-user")).unwrap();
|
||||
assert_eq!(r.host, "10.0.0.5");
|
||||
assert_eq!(r.user, "cfg-user");
|
||||
assert_eq!(r.port, 2222);
|
||||
assert_eq!(r.key_path, PathBuf::from("/cfg/key"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_when_nothing_configured() {
|
||||
let r = resolve_with(&target(None, None), &HostConfig::default(), Some("os-user")).unwrap();
|
||||
assert_eq!(r.host, "my-vps");
|
||||
assert_eq!(r.user, "os-user");
|
||||
assert_eq!(r.port, 22);
|
||||
assert!(r.key_path.ends_with(".ssh/id_rsa"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_user_everywhere_is_clear_error() {
|
||||
let err = resolve_with(&target(None, None), &HostConfig::default(), None).unwrap_err();
|
||||
assert!(err.to_string().contains("set 'user'"), "{err:#}");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user