feat: add ssh-agent flow
This commit is contained in:
+106
@@ -2,10 +2,16 @@ use std::path::Path;
|
|||||||
use std::process::Command;
|
use std::process::Command;
|
||||||
|
|
||||||
use anyhow::{Context, Result, bail};
|
use anyhow::{Context, Result, bail};
|
||||||
|
use std::fs;
|
||||||
|
|
||||||
pub fn execute(cmd: &str, workdir: &Path) -> Result<()> {
|
pub fn execute(cmd: &str, workdir: &Path) -> Result<()> {
|
||||||
let mut command = shell_command(cmd);
|
let mut command = shell_command(cmd);
|
||||||
command.current_dir(workdir);
|
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
|
let status = command
|
||||||
.status()
|
.status()
|
||||||
@@ -21,6 +27,49 @@ pub fn execute(cmd: &str, workdir: &Path) -> Result<()> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn load_dotenv(workdir: &Path) -> Result<Vec<(String, String)>> {
|
||||||
|
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 {
|
fn shell_command(cmd: &str) -> Command {
|
||||||
#[cfg(target_os = "windows")]
|
#[cfg(target_os = "windows")]
|
||||||
{
|
{
|
||||||
@@ -118,4 +167,61 @@ mod tests {
|
|||||||
let err = execute(&ok_cmd(), &missing).unwrap_err();
|
let err = execute(&ok_cmd(), &missing).unwrap_err();
|
||||||
assert!(err.to_string().contains("failed to spawn"), "{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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+56
@@ -5,6 +5,7 @@ use std::sync::{Arc, OnceLock};
|
|||||||
use anyhow::{Context, Result, bail};
|
use anyhow::{Context, Result, bail};
|
||||||
use russh::ChannelMsg;
|
use russh::ChannelMsg;
|
||||||
use russh::client::{self, AuthResult, Handle};
|
use russh::client::{self, AuthResult, Handle};
|
||||||
|
use russh::keys::agent::client::{AgentClient, AgentStream};
|
||||||
use russh::keys::known_hosts;
|
use russh::keys::known_hosts;
|
||||||
use russh::keys::{PrivateKeyWithHashAlg, PublicKey, load_secret_key};
|
use russh::keys::{PrivateKeyWithHashAlg, PublicKey, load_secret_key};
|
||||||
use ssh2_config::{HostParams, ParseRule, SshConfig};
|
use ssh2_config::{HostParams, ParseRule, SshConfig};
|
||||||
@@ -214,6 +215,10 @@ pub fn connect(resolved: &Resolved) -> Result<Session> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn authenticate(handle: &mut Handle<ClientHandler>, r: &Resolved) -> Result<()> {
|
async fn authenticate(handle: &mut Handle<ClientHandler>, 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| {
|
let secret = load_secret_key(&r.key_path, None).map_err(|err| {
|
||||||
anyhow::anyhow!(
|
anyhow::anyhow!(
|
||||||
"cannot load private key {}: {err} (passphrase-protected keys are not supported yet)",
|
"cannot load private key {}: {err} (passphrase-protected keys are not supported yet)",
|
||||||
@@ -241,6 +246,57 @@ async fn authenticate(handle: &mut Handle<ClientHandler>, r: &Resolved) -> Resul
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn try_agent_auth(handle: &mut Handle<ClientHandler>, user: &str) -> Result<bool> {
|
||||||
|
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<Box<dyn AgentStream + Send + Unpin>>;
|
||||||
|
|
||||||
|
async fn connect_agent() -> Option<DynAgentClient> {
|
||||||
|
#[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 {
|
impl Session {
|
||||||
pub fn exec(&mut self, cmd: &str) -> Result<()> {
|
pub fn exec(&mut self, cmd: &str) -> Result<()> {
|
||||||
runtime().block_on(async {
|
runtime().block_on(async {
|
||||||
|
|||||||
Reference in New Issue
Block a user