Mirror environment config in SDK client

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<String>,
     identity_files: Vec<String>,
     identities_only: bool,
+    send_env: Vec<String>,
+    set_env: Vec<EnvVar>,
 }
 
 fn load_sdk_ssh_config(alias: Option<&str>, ssh_port: Option<u16>) -> Result<SdkSshConfig> {
@@ -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<EnvVar>,
+) -> Vec<EnvVar> {
+    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<EnvVar> {
+    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!(