diff --git a/src/client.rs b/src/client.rs index d26f84c..c1809e1 100644 --- a/src/client.rs +++ b/src/client.rs @@ -15,6 +15,7 @@ use crate::ssh_agent; use crate::transport::{DoshTransport, SessionRole, SessionTransportConfig, TransportConfig}; use crate::udp::{bind_udp_for_peer, recv_udp_retrying_transient, send_udp_retrying_transient}; use anyhow::{Context, Result, anyhow, bail}; +use std::collections::BTreeMap; use std::net::{SocketAddr, ToSocketAddrs}; use std::path::PathBuf; use std::process::Command; @@ -179,6 +180,8 @@ impl DoshClientBuilder { Duration::from_millis(self.client.config.native_auth_timeout_ms.max(1)) }); let session = self.session.unwrap_or_else(default_sdk_session); + let requested_env = + sdk_requested_env(&self.client.config, &host_config, &ssh_config, self.env); let requested_forwardings = self .services .iter() @@ -208,7 +211,7 @@ impl DoshClientBuilder { &self.identity_files, self.use_ssh_agent, self.trust_on_first_use, - &self.env, + &requested_env, timeout, ) .await @@ -481,6 +484,8 @@ struct SdkSshConfig { user: Option, identity_files: Vec, identities_only: bool, + send_env: Vec, + set_env: Vec, } fn load_sdk_ssh_config(alias: Option<&str>, ssh_port: Option) -> Result { @@ -520,11 +525,117 @@ fn parse_sdk_ssh_config(raw: &str) -> SdkSshConfig { config.identity_files.push(value.to_string()); } else if key.eq_ignore_ascii_case("identitiesonly") { config.identities_only = value.eq_ignore_ascii_case("yes"); + } else if key.eq_ignore_ascii_case("sendenv") && !empty_or_none(value) { + config + .send_env + .extend(value.split_whitespace().map(ToString::to_string)); + } else if key.eq_ignore_ascii_case("setenv") && !empty_or_none(value) { + config.set_env.extend(parse_set_env_values(value)); } } config } +fn sdk_requested_env( + config: &ClientConfig, + host: &HostConfig, + ssh_config: &SdkSshConfig, + explicit_env: Vec, +) -> Vec { + let mut values = BTreeMap::new(); + let mut patterns = host + .send_env + .clone() + .unwrap_or_else(|| config.send_env.clone()); + patterns.extend(ssh_config.send_env.clone()); + + for (name, value) in std::env::vars() { + if valid_env_name(&name) && patterns.iter().any(|pattern| glob_matches(pattern, &name)) { + values.insert(name, value); + } + } + for (name, value) in &config.set_env { + if valid_env_name(name) && !value.as_bytes().contains(&0) { + values.insert(name.clone(), value.clone()); + } + } + for (name, value) in &host.set_env { + if valid_env_name(name) && !value.as_bytes().contains(&0) { + values.insert(name.clone(), value.clone()); + } + } + for env in &ssh_config.set_env { + if valid_env_name(&env.name) && !env.value.as_bytes().contains(&0) { + values.insert(env.name.clone(), env.value.clone()); + } + } + for env in explicit_env { + if valid_env_name(&env.name) && !env.value.as_bytes().contains(&0) { + values.insert(env.name, env.value); + } + } + values + .into_iter() + .map(|(name, value)| EnvVar { name, value }) + .collect() +} + +fn parse_set_env_values(value: &str) -> Vec { + value + .split_whitespace() + .filter_map(|entry| { + let (name, value) = entry.split_once('=')?; + if valid_env_name(name) && !value.as_bytes().contains(&0) { + Some(EnvVar { + name: name.to_string(), + value: value.to_string(), + }) + } else { + None + } + }) + .collect() +} + +fn valid_env_name(name: &str) -> bool { + let mut chars = name.chars(); + let Some(first) = chars.next() else { + return false; + }; + if !(first == '_' || first.is_ascii_alphabetic()) { + return false; + } + chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric()) +} + +fn glob_matches(pattern: &str, value: &str) -> bool { + let pattern = pattern.as_bytes(); + let value = value.as_bytes(); + let (mut p, mut v) = (0, 0); + let mut star = None; + let mut star_value = 0; + while v < value.len() { + if p < pattern.len() && (pattern[p] == b'?' || pattern[p] == value[v]) { + p += 1; + v += 1; + } else if p < pattern.len() && pattern[p] == b'*' { + star = Some(p); + star_value = v; + p += 1; + } else if let Some(s) = star { + p = s + 1; + star_value += 1; + v = star_value; + } else { + return false; + } + } + while p < pattern.len() && pattern[p] == b'*' { + p += 1; + } + p == pattern.len() +} + fn empty_or_none(value: &str) -> bool { value.is_empty() || value.eq_ignore_ascii_case("none") } @@ -724,7 +835,9 @@ mod tests { port 2222\n\ identityfile ~/.ssh/work\n\ identityfile none\n\ - identitiesonly yes\n", + identitiesonly yes\n\ + sendenv LANG LC_*\n\ + setenv DOSH_MODE=sdk DOSH_COLOR=true BAD-NAME=nope\n", ); assert_eq!(parsed.hostname.as_deref(), Some("10.0.0.5")); @@ -732,6 +845,20 @@ mod tests { assert_eq!(parsed.port, Some(2222)); assert_eq!(parsed.identity_files, vec!["~/.ssh/work"]); assert!(parsed.identities_only); + assert_eq!(parsed.send_env, vec!["LANG", "LC_*"]); + assert_eq!( + parsed.set_env, + vec![ + EnvVar { + name: "DOSH_MODE".to_string(), + value: "sdk".to_string() + }, + EnvVar { + name: "DOSH_COLOR".to_string(), + value: "true".to_string() + } + ] + ); } #[test] @@ -753,6 +880,7 @@ mod tests { user: Some("deploy".to_string()), identity_files: vec![format!("{}/%r_%h_%p", dir.path().display())], identities_only: true, + ..SdkSshConfig::default() }; let config = ClientConfig { identity_files: vec![config_identity.display().to_string()], @@ -789,6 +917,62 @@ mod tests { assert_eq!(paths, vec![dir.path().join("ssh"), config_identity]); } + #[test] + fn sdk_requested_env_merges_host_ssh_and_explicit_overrides() { + let mut config = ClientConfig::default(); + config.send_env.clear(); + config.set_env.insert("DOSH_MODE".into(), "client".into()); + config.set_env.insert("DOSH_KEEP".into(), "yes".into()); + let mut host = HostConfig::default(); + host.set_env.insert("DOSH_MODE".into(), "host".into()); + let ssh_config = SdkSshConfig { + set_env: vec![EnvVar { + name: "DOSH_SSH".to_string(), + value: "true".to_string(), + }], + ..SdkSshConfig::default() + }; + + assert_eq!( + sdk_requested_env( + &config, + &host, + &ssh_config, + vec![ + EnvVar { + name: "DOSH_MODE".to_string(), + value: "explicit".to_string(), + }, + EnvVar { + name: "BAD-NAME".to_string(), + value: "ignored".to_string(), + } + ], + ), + vec![ + EnvVar { + name: "DOSH_KEEP".to_string(), + value: "yes".to_string() + }, + EnvVar { + name: "DOSH_MODE".to_string(), + value: "explicit".to_string() + }, + EnvVar { + name: "DOSH_SSH".to_string(), + value: "true".to_string() + } + ] + ); + } + + #[test] + fn sdk_env_globs_match_send_env_patterns() { + assert!(glob_matches("LC_*", "LC_ALL")); + assert!(glob_matches("TERM", "TERM")); + assert!(!glob_matches("LC_*", "LANG")); + } + #[test] fn default_identity_paths_are_expanded() { assert!(