diff --git a/libshpool/src/attach.rs b/libshpool/src/attach.rs index b1ba5ffc..68da5b6f 100644 --- a/libshpool/src/attach.rs +++ b/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::>() + ); 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::>(), + local_env, ttl_secs: self.ttl.map(|d| d.as_secs()), cmd: resolved.cmd.clone(), dir: start_dir, diff --git a/libshpool/src/config.rs b/libshpool/src/config.rs index 1c68d0e1..e3d8d7b2 100644 --- a/libshpool/src/config.rs +++ b/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, - /// 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>, + /// 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, /// 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), +} + +impl From> for ForwardEnv { + fn from(list: Vec) -> Self { + ForwardEnv::List(list) + } +} + +impl From 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<()> { diff --git a/shpool/tests/attach.rs b/shpool/tests/attach.rs index a31a465d..64fe9c66 100644 --- a/shpool/tests/attach.rs +++ b/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 diff --git a/shpool/tests/data/forward_env_all.toml b/shpool/tests/data/forward_env_all.toml new file mode 100644 index 00000000..4f02f5a5 --- /dev/null +++ b/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 = "" diff --git a/shpool/tests/data/forward_env_none.toml b/shpool/tests/data/forward_env_none.toml new file mode 100644 index 00000000..0ef2baf8 --- /dev/null +++ b/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 = ""