feat: add support for forwarding the entire env

libshpool/src/attach.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     env, fmt, io,
     os::fd::AsFd,
     path::PathBuf,
@@ -253,13 +253,28 @@ impl Attach {
         };
 
         let forward_env = self.config_manager.get().forward_env.clone();
-        let mut local_env_keys = vec!["TERM", "DISPLAY", "LANG", "SSH_AUTH_SOCK"];
-        if let Some(fenv) = &forward_env {
+        let mut local_env_keys = HashSet::with_capacity(4);
+        for key in vec!["TERM", "DISPLAY", "LANG", "SSH_AUTH_SOCK"].into_iter() {
+            local_env_keys.insert(key);
+        }
+        if let Some(config::ForwardEnv::List(fenv)) = &forward_env {
             for var in fenv.iter() {
-                local_env_keys.push(var);
+                local_env_keys.insert(var);
             }
-        }
-        info!("local env keys: {local_env_keys:?}");
+        };
+        let full_env: Vec<(String, String)> = env::vars().collect();
+        let local_env: Vec<(String, String)> = match forward_env {
+            Some(config::ForwardEnv::All(false)) => vec![],
+            Some(config::ForwardEnv::All(true)) => full_env,
+            None | Some(config::ForwardEnv::List(_)) => full_env
+                .into_iter()
+                .filter(|(var, _)| local_env_keys.contains(var.as_str()))
+                .collect(),
+        };
+        info!(
+            "local env keys: {:?}",
+            local_env.iter().map(|(k, _)| k.as_str()).collect::<Vec<_>>()
+        );
 
         let cwd = String::from(env::current_dir().context("getting cwd")?.to_string_lossy());
         let default_dir =
@@ -276,13 +291,7 @@ impl Attach {
             .write_connect_header(ConnectHeader::Attach(AttachHeader {
                 name: resolved.session_name.clone(),
                 local_tty_size: tty_size,
-                local_env: local_env_keys
-                    .into_iter()
-                    .filter_map(|var| {
-                        let val = env::var(var).context("resolving var").ok()?;
-                        Some((String::from(var), val))
-                    })
-                    .collect::<Vec<_>>(),
+                local_env,
                 ttl_secs: self.ttl.map(|d| d.as_secs()),
                 cmd: resolved.cmd.clone(),
                 dir: start_dir,

libshpool/src/config.rs

@@ -230,11 +230,17 @@ pub struct Config {
     /// If it is '.', it will start wherever `shpool attach` is invoked.
     pub default_dir: Option<String>,
 
-    /// A list of environment variables to forward from the environment
+    /// Controls environment variables to forward from the environment
     /// of the initial shell that invoked `shpool attach` to the newly
-    /// launched shell. Note that this config option has no impact when
-    /// reattaching to an existing shell.
-    pub forward_env: Option<Vec<String>>,
+    /// launched shell. This can be a list of environment variable names
+    /// to forward, or a boolean (`true` to forward all environment
+    /// variables from the client environment, `false` to forward none at
+    /// all, not even the defaults). Note that this config
+    /// option has no impact when reattaching to an existing shell.
+    ///
+    /// By default the variables TERM, DISPLAY, LANG, and SSH_AUTH_SOCK
+    /// are forwarded.
+    pub forward_env: Option<ForwardEnv>,
 
     /// The initial path to spawn shell processes with. By default
     /// `/usr/bin:/bin:/usr/sbin:/sbin` (copying openssh). This
@@ -423,6 +429,25 @@ pub struct VarSetting {
     pub value: String,
 }
 
+#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
+#[serde(untagged)]
+pub enum ForwardEnv {
+    All(bool),
+    List(Vec<String>),
+}
+
+impl From<Vec<String>> for ForwardEnv {
+    fn from(list: Vec<String>) -> Self {
+        ForwardEnv::List(list)
+    }
+}
+
+impl From<bool> for ForwardEnv {
+    fn from(all: bool) -> Self {
+        ForwardEnv::All(all)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -457,12 +482,36 @@ mod test {
             var = "foo"
             value = "bar"
             "#,
+            r#"
+            forward_env = true
+            "#,
+            r#"
+            forward_env = false
+            "#,
+            r#"
+            forward_env = ["abc", "efg"]
+            "#,
         ];
 
         for case in cases.into_iter() {
             let _: Config = toml::from_str(case)?;
         }
 
+        let c_true: Config = toml::from_str("forward_env = true")?;
+        assert_eq!(c_true.forward_env, Some(ForwardEnv::All(true)));
+
+        let c_false: Config = toml::from_str("forward_env = false")?;
+        assert_eq!(c_false.forward_env, Some(ForwardEnv::All(false)));
+
+        let c_list: Config = toml::from_str(r#"forward_env = ["abc", "efg"]"#)?;
+        assert_eq!(
+            c_list.forward_env,
+            Some(ForwardEnv::List(vec!["abc".to_string(), "efg".to_string()]))
+        );
+
+        let c_none: Config = toml::from_str("")?;
+        assert_eq!(c_none.forward_env, None);
+
         Ok(())
     }
 
@@ -504,7 +553,7 @@ mod test {
         #[timeout(30000)]
         fn vec_value() -> Result<()> {
             let higher = Config {
-                forward_env: Some(vec!["abc".to_string(), "efg".to_string()]),
+                forward_env: Some(ForwardEnv::List(vec!["abc".to_string(), "efg".to_string()])),
                 motd_args: None,
                 ..Default::default()
             };
@@ -515,11 +564,28 @@ mod test {
             };
 
             let actual = higher.merge(lower);
-            assert_eq!(actual.forward_env, Some(vec!["abc".to_string(), "efg".to_string()]));
+            assert_eq!(
+                actual.forward_env,
+                Some(ForwardEnv::List(vec!["abc".to_string(), "efg".to_string()]))
+            );
             assert_eq!(actual.motd_args, Some(vec!["hij".to_string(), "klm".to_string()]));
             Ok(())
         }
 
+        #[test]
+        #[timeout(30000)]
+        fn forward_env_bool_merge() -> Result<()> {
+            let higher = Config { forward_env: Some(ForwardEnv::All(true)), ..Default::default() };
+            let lower = Config {
+                forward_env: Some(ForwardEnv::List(vec!["abc".to_string()])),
+                ..Default::default()
+            };
+
+            let actual = higher.merge(lower);
+            assert_eq!(actual.forward_env, Some(ForwardEnv::All(true)));
+            Ok(())
+        }
+
         #[test]
         #[timeout(30000)]
         fn map_value() -> Result<()> {

shpool/tests/attach.rs

@@ -181,6 +181,119 @@ fn forward_env() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[test]
+#[timeout(30000)]
+fn forward_env_all() -> anyhow::Result<()> {
+    let mut daemon_proc = support::daemon::Proc::new("forward_env_all.toml", DaemonArgs::default())
+        .context("starting daemon proc")?;
+
+    let bidi_done_w = daemon_proc.events.take().unwrap().waiter(["daemon-bidi-stream-done"]);
+    {
+        let mut attach_proc = daemon_proc
+            .attach(
+                "sh1",
+                AttachArgs {
+                    extra_env: vec![
+                        (String::from("FOO"), String::from("foo")),
+                        (String::from("BAR"), String::from("bar")),
+                        (String::from("BAZ"), String::from("baz")),
+                    ],
+                    ..Default::default()
+                },
+            )
+            .context("starting attach proc")?;
+
+        let mut line_matcher = attach_proc.line_matcher()?;
+
+        attach_proc.run_cmd(r#"echo "$FOO:$BAR:$BAZ" "#)?;
+        line_matcher.scan_until_re("foo:bar:baz$")?;
+    }
+
+    // wait until the daemon has noticed that the connection
+    // has dropped before we attempt to open the connection again
+    daemon_proc.events = Some(bidi_done_w.wait_final_event("daemon-bidi-stream-done")?);
+
+    {
+        let mut attach_proc = daemon_proc
+            .attach(
+                "sh1",
+                AttachArgs {
+                    extra_env: vec![
+                        (String::from("FOO"), String::from("foonew")),
+                        (String::from("BAR"), String::from("barnew")),
+                        (String::from("BAZ"), String::from("baznew")),
+                    ],
+                    ..Default::default()
+                },
+            )
+            .context("starting attach proc")?;
+
+        let mut line_matcher = attach_proc.line_matcher()?;
+
+        attach_proc.run_cmd(r#"source $SHPOOL_SESSION_DIR/forward.env "#)?;
+        attach_proc.run_cmd(r#"echo "$FOO:$BAR:$BAZ" "#)?;
+        line_matcher.scan_until_re("foonew:barnew:baznew$")?;
+    }
+
+    Ok(())
+}
+
+#[test]
+#[timeout(30000)]
+fn forward_env_none() -> anyhow::Result<()> {
+    let mut daemon_proc =
+        support::daemon::Proc::new("forward_env_none.toml", DaemonArgs::default())
+            .context("starting daemon proc")?;
+
+    let bidi_done_w = daemon_proc.events.take().unwrap().waiter(["daemon-bidi-stream-done"]);
+    {
+        let mut attach_proc = daemon_proc
+            .attach(
+                "sh1",
+                AttachArgs {
+                    extra_env: vec![
+                        (String::from("FOO"), String::from("foo")),
+                        (String::from("DISPLAY"), String::from("fake-display")),
+                    ],
+                    ..Default::default()
+                },
+            )
+            .context("starting attach proc")?;
+
+        let mut line_matcher = attach_proc.line_matcher()?;
+
+        attach_proc.run_cmd(r#"echo "$FOO:$DISPLAY" "#)?;
+        line_matcher.scan_until_re(":$")?;
+    }
+
+    // wait until the daemon has noticed that the connection
+    // has dropped before we attempt to open the connection again
+    daemon_proc.events = Some(bidi_done_w.wait_final_event("daemon-bidi-stream-done")?);
+
+    {
+        let mut attach_proc = daemon_proc
+            .attach(
+                "sh1",
+                AttachArgs {
+                    extra_env: vec![
+                        (String::from("FOO"), String::from("foonew")),
+                        (String::from("DISPLAY"), String::from("fake-display-new")),
+                    ],
+                    ..Default::default()
+                },
+            )
+            .context("starting attach proc")?;
+
+        let mut line_matcher = attach_proc.line_matcher()?;
+
+        attach_proc.run_cmd(r#"source $SHPOOL_SESSION_DIR/forward.env "#)?;
+        attach_proc.run_cmd(r#"echo "$FOO:$DISPLAY" "#)?;
+        line_matcher.scan_until_re(":$")?;
+    }
+
+    Ok(())
+}
+
 // Regression test: a high byte (0xFF) in the raw input stream must not
 // kill the session. The keybinding scanner used to index out of bounds
 // on 0xFF, panicking the client->shell thread and disconnecting the

shpool/tests/data/forward_env_all.toml

@@ -0,0 +1,11 @@
+norc = true
+noecho = true
+shell = "/bin/bash"
+session_restore_mode = "simple"
+prompt_prefix = ""
+
+forward_env = true
+
+[env]
+PS1 = "prompt> "
+TERM = ""

shpool/tests/data/forward_env_none.toml

@@ -0,0 +1,11 @@
+norc = true
+noecho = true
+shell = "/bin/bash"
+session_restore_mode = "simple"
+prompt_prefix = ""
+
+forward_env = false
+
+[env]
+PS1 = "prompt> "
+TERM = ""