From dd7151bca3914e30c833038b80ead230f74a3a49 Mon Sep 17 00:00:00 2001 From: Ku6epXBOCTuK Date: Sun, 23 Aug 2026 06:10:15 +0500 Subject: [PATCH] feat: add ssh-agent flow --- src/local.rs | 106 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/ssh.rs | 56 +++++++++++++++++++++++++++ 2 files changed, 162 insertions(+) diff --git a/src/local.rs b/src/local.rs index ef4a796..a6bb930 100644 --- a/src/local.rs +++ b/src/local.rs @@ -2,10 +2,16 @@ use std::path::Path; use std::process::Command; use anyhow::{Context, Result, bail}; +use std::fs; pub fn execute(cmd: &str, workdir: &Path) -> Result<()> { let mut command = shell_command(cmd); command.current_dir(workdir); + for (key, value) in load_dotenv(workdir)? { + if std::env::var_os(&key).is_none() { + command.env(key, value); + } + } let status = command .status() @@ -21,6 +27,49 @@ pub fn execute(cmd: &str, workdir: &Path) -> Result<()> { Ok(()) } +fn load_dotenv(workdir: &Path) -> Result> { + let path = workdir.join(".env"); + if !path.is_file() { + return Ok(Vec::new()); + } + let raw = + fs::read_to_string(&path).with_context(|| format!("cannot read {}", path.display()))?; + let mut vars = Vec::new(); + for (num, line) in raw.lines().enumerate() { + let line = line.trim(); + if line.is_empty() || line.starts_with('#') { + continue; + } + let line = line.strip_prefix("export ").unwrap_or(line).trim_start(); + let Some((key, value)) = line.split_once('=') else { + bail!(".env line {}: expected KEY=VALUE", num + 1); + }; + validate_var_name(key, num)?; + vars.push((key.to_string(), unquote(value.trim()))); + } + Ok(vars) +} + +fn validate_var_name(key: &str, line: usize) -> Result<()> { + let valid = !key.is_empty() + && !key.chars().next().unwrap().is_ascii_digit() + && key.chars().all(|c| c.is_ascii_alphanumeric() || c == '_'); + if !valid { + bail!(".env line {line}: invalid variable name '{key}'"); + } + Ok(()) +} + +fn unquote(mut value: &str) -> String { + for quote in ['"', '\''] { + if value.len() >= 2 && value.starts_with(quote) && value.ends_with(quote) { + value = &value[1..value.len() - 1]; + break; + } + } + value.to_string() +} + fn shell_command(cmd: &str) -> Command { #[cfg(target_os = "windows")] { @@ -118,4 +167,61 @@ mod tests { let err = execute(&ok_cmd(), &missing).unwrap_err(); assert!(err.to_string().contains("failed to spawn"), "{err:#}"); } + + #[test] + fn dotenv_parses_comments_quotes_and_export() { + let workdir = tmpdir("dotenv"); + fs::write( + workdir.join(".env"), + "# comment\nDB_HOST=localhost\nexport API_KEY=\"abc=123\"\nSINGLE='v v'\n\n", + ) + .unwrap(); + let vars = load_dotenv(&workdir).unwrap(); + assert_eq!( + vars, + vec![ + ("DB_HOST".to_string(), "localhost".to_string()), + ("API_KEY".to_string(), "abc=123".to_string()), + ("SINGLE".to_string(), "v v".to_string()), + ] + ); + let _ = fs::remove_dir_all(workdir); + } + + #[test] + fn dotenv_invalid_line_is_clear_error() { + let workdir = tmpdir("dotenv-bad"); + fs::write(workdir.join(".env"), "GOOD=1\nNO_EQUALS_HERE\n").unwrap(); + let err = load_dotenv(&workdir).unwrap_err(); + assert!(err.to_string().contains("line 2"), "{err:#}"); + let _ = fs::remove_dir_all(workdir); + } + + #[test] + fn dotenv_missing_file_is_empty() { + let workdir = tmpdir("dotenv-none"); + assert!(load_dotenv(&workdir).unwrap().is_empty()); + let _ = fs::remove_dir_all(workdir); + } + + #[test] + fn env_vars_visible_inside_local_commands() { + let workdir = tmpdir("dotenv-exec"); + let echo_var = if cfg!(windows) { + "echo %XBP_TEST_VAR%" + } else { + "echo $XBP_TEST_VAR" + }; + fs::write(workdir.join(".env"), "XBP_TEST_VAR=hello-env\n").unwrap(); + let mut command = shell_command(echo_var); + command.current_dir(&workdir); + for (k, v) in load_dotenv(&workdir).unwrap() { + if std::env::var_os(&k).is_none() { + command.env(k, v); + } + } + let out = command.output().unwrap(); + assert!(String::from_utf8_lossy(&out.stdout).contains("hello-env")); + let _ = fs::remove_dir_all(workdir); + } } diff --git a/src/ssh.rs b/src/ssh.rs index b399af8..d662e69 100644 --- a/src/ssh.rs +++ b/src/ssh.rs @@ -5,6 +5,7 @@ use std::sync::{Arc, OnceLock}; use anyhow::{Context, Result, bail}; use russh::ChannelMsg; use russh::client::{self, AuthResult, Handle}; +use russh::keys::agent::client::{AgentClient, AgentStream}; use russh::keys::known_hosts; use russh::keys::{PrivateKeyWithHashAlg, PublicKey, load_secret_key}; use ssh2_config::{HostParams, ParseRule, SshConfig}; @@ -214,6 +215,10 @@ pub fn connect(resolved: &Resolved) -> Result { } async fn authenticate(handle: &mut Handle, r: &Resolved) -> Result<()> { + if try_agent_auth(handle, &r.user).await? { + return Ok(()); + } + 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)", @@ -241,6 +246,57 @@ async fn authenticate(handle: &mut Handle, r: &Resolved) -> Resul } } +async fn try_agent_auth(handle: &mut Handle, user: &str) -> Result { + let Some(mut agent) = connect_agent().await else { + return Ok(false); + }; + let identities = match agent.request_identities().await { + Ok(identities) => identities, + Err(_) => return Ok(false), + }; + for identity in identities { + let hash = handle + .best_supported_rsa_hash() + .await + .ok() + .flatten() + .flatten(); + let public_key = identity.public_key().into_owned(); + match handle + .authenticate_publickey_with(user, public_key, hash, &mut agent) + .await + { + Ok(AuthResult::Success) => return Ok(true), + _ => continue, + } + } + Ok(false) +} + +type DynAgentClient = AgentClient>; + +async fn connect_agent() -> Option { + #[cfg(unix)] + if let Ok(sock) = std::env::var("SSH_AUTH_SOCK") { + if let Ok(client) = AgentClient::connect_uds(sock).await { + return Some(client.dynamic()); + } + } + + #[cfg(windows)] + { + const OPENSSH_PIPE: &str = "\\\\.\\pipe\\openssh-ssh-agent"; + if let Ok(client) = AgentClient::connect_named_pipe(OPENSSH_PIPE).await { + return Some(client.dynamic()); + } + if let Ok(client) = AgentClient::connect_pageant().await { + return Some(client.dynamic()); + } + } + + None +} + impl Session { pub fn exec(&mut self, cmd: &str) -> Result<()> { runtime().block_on(async {