Add explicit environment propagation

src/bin/dosh-client.rs

@@ -9,8 +9,8 @@ use dosh::auth::{
 use dosh::config::{expand_tilde, load_client_config, load_hosts_config, load_server_config};
 use dosh::crypto;
 use dosh::native::{
-    ForwardingKind, ForwardingRequest, KnownHostStatus, TrustResult, derive_native_session_key,
-    generate_native_ephemeral, host_fingerprint, load_ed25519_identity,
+    EnvVar, ForwardingKind, ForwardingRequest, KnownHostStatus, TrustResult,
+    derive_native_session_key, generate_native_ephemeral, host_fingerprint, load_ed25519_identity,
     load_ed25519_identity_with_passphrase, parse_host_public_key_line, remove_trusted_host,
     sign_user_auth, trust_host, verify_known_host, verify_server_hello,
 };
@@ -270,6 +270,7 @@ async fn main() -> Result<()> {
     } else {
         None
     };
+    let cache_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref());
 
     let socket = UdpSocket::bind("0.0.0.0:0").await?;
 
@@ -278,7 +279,8 @@ async fn main() -> Result<()> {
         cached.udp_port = dosh_port;
         log_timing(args.verbose, "credential_lookup_end", started.elapsed());
         if config.cache_attach_tickets && allow_cache {
-            match try_ticket_attach(&socket, &cached, cols, rows).await {
+            match try_ticket_attach(&socket, &cached, cols, rows, cache_requested_env.clone()).await
+            {
                 Ok((frame, cred)) => {
                     log_timing(args.verbose, "udp_ticket_attach_ready", started.elapsed());
                     save_cache(&cache_path, &cred)?;
@@ -347,6 +349,10 @@ async fn main() -> Result<()> {
     }
 
     if !args.local_auth && auth_allows(&auth_preference, "native") {
+        if resolved_ssh_config.is_none() {
+            resolved_ssh_config = Some(ssh_config(&server, ssh_port).unwrap_or_default());
+        }
+        let cold_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref());
         let native_start = Instant::now();
         match try_native_auth(
             &socket,
@@ -363,6 +369,7 @@ async fn main() -> Result<()> {
             rows,
             forwarding_requests(&local_forwards, &remote_forwards, &dynamic_forwards),
             resolved_ssh_config.as_ref(),
+            cold_requested_env,
         )
         .await
         {
@@ -408,6 +415,10 @@ async fn main() -> Result<()> {
     }
 
     let bootstrap_start = Instant::now();
+    if !args.local_auth && resolved_ssh_config.is_none() {
+        resolved_ssh_config = Some(ssh_config(&server, ssh_port).unwrap_or_default());
+    }
+    let cold_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref());
     let bootstrap = if args.local_auth {
         local_bootstrap(&session, &mode, cols, rows, dosh_port, target_udp_host)?
     } else {
@@ -429,7 +440,8 @@ async fn main() -> Result<()> {
     };
     log_timing(args.verbose, "ssh_bootstrap", bootstrap_start.elapsed());
 
-    let (ok, mut cred) = bootstrap_attach(&socket, &server, &bootstrap, cols, rows).await?;
+    let (ok, mut cred) =
+        bootstrap_attach(&socket, &server, &bootstrap, cols, rows, cold_requested_env).await?;
     cred.last_rendered_seq = ok.initial_seq;
     if allow_cache {
         save_cache(&cache_path, &cred)?;
@@ -1181,6 +1193,8 @@ struct SshConfig {
     identities_only: bool,
     proxy_jump: Option<String>,
     proxy_command: Option<String>,
+    send_env: Vec<String>,
+    set_env: Vec<EnvVar>,
 }
 
 fn ssh_config(server: &str, ssh_port: Option<u16>) -> Result<SshConfig> {
@@ -1227,15 +1241,117 @@ fn parse_ssh_config(raw: &str) -> SshConfig {
             config.proxy_jump = Some(value.to_string());
         } else if key.eq_ignore_ascii_case("proxycommand") && !empty_or_none(value) {
             config.proxy_command = Some(value.to_string());
+        } 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 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 empty_or_none(value: &str) -> bool {
     value.is_empty() || value.eq_ignore_ascii_case("none")
 }
 
+fn requested_env(
+    config: &dosh::config::ClientConfig,
+    host: &dosh::config::HostConfig,
+    ssh_config: Option<&SshConfig>,
+) -> Vec<EnvVar> {
+    let mut values = BTreeMap::new();
+    let mut patterns = host
+        .send_env
+        .clone()
+        .unwrap_or_else(|| config.send_env.clone());
+    if let Some(ssh_config) = ssh_config {
+        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());
+        }
+    }
+    if let Some(ssh_config) = ssh_config {
+        for env in &ssh_config.set_env {
+            values.insert(env.name.clone(), env.value.clone());
+        }
+    }
+    values
+        .into_iter()
+        .map(|(name, value)| EnvVar { name, value })
+        .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);
+            p += 1;
+            star_value = v;
+        } else if let Some(star_pos) = star {
+            p = star_pos + 1;
+            star_value += 1;
+            v = star_value;
+        } else {
+            return false;
+        }
+    }
+    while p < pattern.len() && pattern[p] == b'*' {
+        p += 1;
+    }
+    p == pattern.len()
+}
+
 fn resolve_addr(host: &str, port: u16) -> Result<SocketAddr> {
     (host, port)
         .to_socket_addrs()
@@ -1259,6 +1375,7 @@ async fn try_native_auth(
     rows: u16,
     requested_forwardings: Vec<ForwardingRequest>,
     ssh_config_hint: Option<&SshConfig>,
+    requested_env: Vec<EnvVar>,
 ) -> Result<(Frame, CachedCredential)> {
     let addr = resolve_addr(udp_host, udp_port)?;
     let owned_ssh_config;
@@ -1285,6 +1402,7 @@ async fn try_native_auth(
         supported_user_key_algorithms: vec!["ssh-ed25519".to_string()],
         cached_host_key_fingerprint: None,
         attach_ticket_envelope: None,
+        requested_env,
     };
     let packet = protocol::encode_plain(
         PacketKind::NativeClientHello,
@@ -1440,6 +1558,7 @@ async fn try_native_auth_check(
         supported_user_key_algorithms: vec!["ssh-ed25519".to_string()],
         cached_host_key_fingerprint: None,
         attach_ticket_envelope: None,
+        requested_env: Vec::new(),
     };
     let packet = protocol::encode_plain(
         PacketKind::NativeClientHello,
@@ -1651,12 +1770,14 @@ async fn bootstrap_attach(
     bootstrap: &BootstrapResponse,
     cols: u16,
     rows: u16,
+    requested_env: Vec<EnvVar>,
 ) -> Result<(AttachOk, CachedCredential)> {
     let addr = resolve_addr(&bootstrap.udp_host, bootstrap.udp_port)?;
     let req = BootstrapAttachRequest {
         bootstrap: bootstrap.clone(),
         cols,
         rows,
+        requested_env,
     };
     let body = protocol::to_body(&req)?;
     let packet =
@@ -1701,6 +1822,7 @@ async fn try_ticket_attach(
     cached: &CachedCredential,
     cols: u16,
     rows: u16,
+    requested_env: Vec<EnvVar>,
 ) -> Result<(Frame, CachedCredential)> {
     let addr = resolve_addr(&cached.udp_host, cached.udp_port)?;
     let client_nonce = crypto::random_12();
@@ -1714,6 +1836,7 @@ async fn try_ticket_attach(
         mode: cached.mode.clone(),
         cols,
         rows,
+        requested_env,
     })?;
     let ciphertext = crypto::seal(
         &request_key,
@@ -2453,7 +2576,8 @@ async fn reconnect(
         return Ok(Some(frame));
     }
 
-    let Ok((frame, next)) = try_ticket_attach(socket, cred, size.0, size.1).await else {
+    let Ok((frame, next)) = try_ticket_attach(socket, cred, size.0, size.1, Vec::new()).await
+    else {
         return Ok(None);
     };
     *cred = next;
@@ -2840,10 +2964,11 @@ mod tests {
         DynamicForward, FrameBuffer, LocalForward, Predictor, RemoteForward, SshConfig,
         auth_allows, load_first_native_identity_with_prompt, parse_dynamic_forward,
         parse_local_forward, parse_remote_forward, parse_ssh_config, raw_contains_host_table,
-        recv_response_until, ssh_destination_host, ssh_username, startup_command,
+        recv_response_until, requested_env, ssh_destination_host, ssh_username, startup_command,
         toml_bare_key_or_quoted,
     };
-    use dosh::config::ClientConfig;
+    use dosh::config::{ClientConfig, HostConfig};
+    use dosh::native::EnvVar;
     use dosh::protocol::{self, Frame, PacketKind};
     use ed25519_dalek::{SigningKey, VerifyingKey};
 
@@ -2891,7 +3016,9 @@ mod tests {
              identitiesonly yes\n\
              proxyjump bastion\n\
              proxycommand ssh bastion -W %h:%p\n\
-             identityfile none\n",
+             identityfile none\n\
+             sendenv LANG LC_*\n\
+             setenv DOSH_MODE=native DOSH_COLOR=true\n",
         );
 
         assert_eq!(
@@ -2904,9 +3031,61 @@ mod tests {
             parsed.proxy_command,
             Some("ssh bastion -W %h:%p".to_string())
         );
+        assert_eq!(
+            parsed.send_env,
+            vec!["LANG".to_string(), "LC_*".to_string()]
+        );
+        assert_eq!(
+            parsed.set_env,
+            vec![
+                EnvVar {
+                    name: "DOSH_MODE".to_string(),
+                    value: "native".to_string()
+                },
+                EnvVar {
+                    name: "DOSH_COLOR".to_string(),
+                    value: "true".to_string()
+                }
+            ]
+        );
         assert!(parsed.identity_files.is_empty());
     }
 
+    #[test]
+    fn requested_env_merges_config_host_and_ssh_setenv() {
+        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 = SshConfig {
+            set_env: vec![EnvVar {
+                name: "DOSH_SSH".to_string(),
+                value: "true".to_string(),
+            }],
+            ..SshConfig::default()
+        };
+
+        assert_eq!(
+            requested_env(&config, &host, Some(&ssh_config)),
+            vec![
+                EnvVar {
+                    name: "DOSH_KEEP".to_string(),
+                    value: "yes".to_string()
+                },
+                EnvVar {
+                    name: "DOSH_MODE".to_string(),
+                    value: "host".to_string()
+                },
+                EnvVar {
+                    name: "DOSH_SSH".to_string(),
+                    value: "true".to_string()
+                }
+            ]
+        );
+    }
+
     #[test]
     fn detects_existing_imported_host_tables() {
         assert!(raw_contains_host_table(