From 96fba5b6cc8e3cf31e18c011bbbab0b6af7399b9 Mon Sep 17 00:00:00 2001 From: Ethan Pailes Date: Thu, 3 Sep 2026 16:19:06 +0000 Subject: [PATCH] feat: add automatic env var reload to prompt hooks This patch adds automatic env var reloading to the default prompt hook shell shims. Previously, you had to set this up yourself manually if you wanted it. This should provide a much more seemless experience. Closes #111 --- libshpool/src/daemon/server.rs | 81 +++++- libshpool/src/daemon/shell_inject.rs | 130 +++++++-- shpool/tests/attach.rs | 338 +++++++++++++++++++++++- shpool/tests/data/forward_env_bash.toml | 10 + shpool/tests/data/forward_env_fish.toml | 9 + shpool/tests/data/forward_env_zsh.toml | 9 + shpool/tests/support/line_matcher.rs | 8 +- 7 files changed, 547 insertions(+), 38 deletions(-) create mode 100644 shpool/tests/data/forward_env_bash.toml create mode 100644 shpool/tests/data/forward_env_fish.toml create mode 100644 shpool/tests/data/forward_env_zsh.toml diff --git a/libshpool/src/daemon/server.rs b/libshpool/src/daemon/server.rs index fc49bd4b..fb23fb86 100644 --- a/libshpool/src/daemon/server.rs +++ b/libshpool/src/daemon/server.rs @@ -623,11 +623,14 @@ impl Server { let session_env_file = self.session_env_file(session_name); info!("populating {:?}", session_env_file); - fs::write( - session_env_file, - header.local_env.iter().map(|(k, v)| format!("{k}={v}")).collect::>().join("\n"), - ) - .context("writing session env")?; + let content = format_forward_env(&header.local_env); + fs::write(&session_env_file, content).context("writing session env")?; + + // Remove the stamp file if present so shells with 1-second timestamp + // resolution (e.g. bash 3.2 on macOS) reload unconditionally on + // reattach without needing a full second to elapse. + let stamp_file = format!("{}.stamp", session_env_file.display()); + let _ = fs::remove_file(stamp_file); Ok(()) } @@ -1361,6 +1364,31 @@ impl Server { } } +fn format_forward_env<'a, I>(env_vars: I) -> String +where + I: IntoIterator, +{ + let mut content = String::new(); + for (k, v) in env_vars { + if is_valid_env_key(k) { + content.push_str(&format!("export {k}='{}'\n", v.replace('\'', "'\\''"))); + } else { + warn!("skipping invalid environment variable key: {k}"); + } + } + content +} + +fn is_valid_env_key(key: &str) -> bool { + let mut chars = key.chars(); + match chars.next() { + Some(c) if c.is_ascii_alphabetic() || c == '_' => { + chars.all(|c| c.is_ascii_alphanumeric() || c == '_') + } + _ => false, + } +} + // HACK: this is not a good way to detect shells that don't support our // sentinel injection approach, but it is better than just hanging when a // user tries to start one. @@ -1497,3 +1525,46 @@ impl std::fmt::Display for ShellSelectionError { } impl std::error::Error for ShellSelectionError {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_format_forward_env_basic() { + let vars = vec![ + (String::from("FOO"), String::from("bar")), + (String::from("SPACES"), String::from("hello world")), + (String::from("QUOTES"), String::from("don't stop")), + (String::from("SPECIAL"), String::from("$(evil) `evil` $VAR")), + (String::from("MULTILINE"), String::from("line1\nline2")), + ]; + let res = format_forward_env(&vars); + assert_eq!( + res, + "export FOO='bar'\n\ + export SPACES='hello world'\n\ + export QUOTES='don'\\''t stop'\n\ + export SPECIAL='$(evil) `evil` $VAR'\n\ + export MULTILINE='line1\nline2'\n" + ); + } + + #[test] + fn test_format_forward_env_invalid_keys() { + let vars = vec![ + (String::from("GOOD_KEY_1"), String::from("val")), + (String::from("_ALSO_GOOD"), String::from("val")), + (String::from("1BAD_KEY"), String::from("val")), + (String::from("BAD-DASH"), String::from("val")), + (String::from("BAD KEY"), String::from("val")), + (String::from(""), String::from("val")), + ]; + let res = format_forward_env(&vars); + assert_eq!( + res, + "export GOOD_KEY_1='val'\n\ + export _ALSO_GOOD='val'\n" + ); + } +} diff --git a/libshpool/src/daemon/shell_inject.rs b/libshpool/src/daemon/shell_inject.rs index c67adc60..eaf89576 100644 --- a/libshpool/src/daemon/shell_inject.rs +++ b/libshpool/src/daemon/shell_inject.rs @@ -72,45 +72,127 @@ pub fn maybe_setup( let prompt_prefix = prompt_prefix.replace("$SHPOOL_SESSION_NAME", session_name); let mut script = match (prompt_prefix.as_str(), shell_type) { - (_, Ok(KnownShell::Bash)) => format!( - r#" - if [[ -z "${{PROMPT_COMMAND+x}}" ]]; then - PS1="{prompt_prefix}${{PS1}}" - else - SHPOOL__OLD_PROMPT_COMMAND=("${{PROMPT_COMMAND[@]}}") - SHPOOL__OLD_PS1="${{PS1}}" - function __shpool__prompt_command() {{ - PS1="${{SHPOOL__OLD_PS1}}" - for prompt_hook in "${{SHPOOL__OLD_PROMPT_COMMAND[@]}}" - do - eval "${{prompt_hook}}" - done - PS1="{prompt_prefix}${{PS1}}" - }} - PROMPT_COMMAND=__shpool__prompt_command - fi + (_, Ok(KnownShell::Bash)) => { + // In Bash 5.1+, PROMPT_COMMAND supports arrays. However, older + // versions of Bash (such as Bash 3.2, the default system shell on + // macOS) only execute PROMPT_COMMAND if it is a scalar string; + // assigning an array causes Bash 3.2 to silently ignore it. + // We capture any existing hooks (array or scalar) into + // SHPOOL__OLD_PROMPT_COMMAND, unset PROMPT_COMMAND, and assign + // PROMPT_COMMAND as a scalar string to ensure universal + // compatibility. + format!( + r#" + SHPOOL__OLD_PROMPT_COMMAND=("${{PROMPT_COMMAND[@]}}") + SHPOOL__OLD_PS1="${{PS1}}" + function __shpool__prompt_command() {{ + local ret=$? + local env_file="${{SHPOOL_SESSION_DIR}}/forward.env" + local stamp_file="${{env_file}}.stamp" + if [ -n "${{SHPOOL_SESSION_DIR}}" ] && [ -f "${{env_file}}" ]; then + if [ ! -f "${{stamp_file}}" ] || [ "${{env_file}}" -nt "${{stamp_file}}" ]; then + touch -r "${{env_file}}" "${{stamp_file}}" 2>/dev/null + + local allexport_was_set=0 + case "$-" in + *a*) allexport_was_set=1 ;; + esac + set -a + . "${{env_file}}" + if [ "$allexport_was_set" -eq 0 ] ; then + set +a + fi + fi + fi + + PS1="${{SHPOOL__OLD_PS1}}" + (exit $ret) + for prompt_hook in "${{SHPOOL__OLD_PROMPT_COMMAND[@]}}" + do + eval "${{prompt_hook}}" + ret=$? + done + PS1="{prompt_prefix}${{PS1}}" + return $ret + }} + unset PROMPT_COMMAND + PROMPT_COMMAND=__shpool__prompt_command "# - ), + ) + } (_, Ok(KnownShell::Zsh)) => format!( r#" typeset -a precmd_functions SHPOOL__OLD_PROMPT="${{PROMPT}}" function __shpool__reset_rprompt() {{ + local ret=$? + local env_file="${{SHPOOL_SESSION_DIR:-}}/forward.env" + local stamp_file="${{env_file}}.stamp" + if [ -n "${{SHPOOL_SESSION_DIR:-}}" ] && [ -f "${{env_file}}" ]; then + if [ ! -f "${{stamp_file}}" ] || [ "${{env_file}}" -nt "${{stamp_file}}" ]; then + touch -r "${{env_file}}" "${{stamp_file}}" 2>/dev/null + + local allexport_was_set=0 + case "$-" in + *a*) allexport_was_set=1 ;; + esac + set -a + . "${{env_file}}" + if [ "$allexport_was_set" -eq 0 ] ; then + set +a + fi + fi + fi + PROMPT="${{SHPOOL__OLD_PROMPT}}" + return $ret }} precmd_functions[1,0]=(__shpool__reset_rprompt) function __shpool__prompt_command() {{ + local ret=$? PROMPT="{prompt_prefix}${{PROMPT}}" + return $ret }} precmd_functions+=(__shpool__prompt_command) "# ), - (_, Ok(KnownShell::Fish)) => format!( - r#" - functions --copy fish_prompt shpool__old_prompt - function fish_prompt; echo -n "{prompt_prefix}"; shpool__old_prompt; end - "# - ), + (_, Ok(KnownShell::Fish)) => { + // Fish only added the `-nt` (newer-than) binary operator to its + // builtin `test` in fish 4.0b1. In older fish versions (such as + // fish 3.x), calling builtin `test -nt` errors with "unexpected + // argument". To maintain zero-fork prompt evaluation on + // fish 4+ while remaining compatible with fish 3, we + // probe for `-nt` support once at injection + // time and define `__shpool_is_newer` to use the builtin if + // available, falling back to `command test` (coreutils) + // otherwise. + format!( + r#" + functions --copy fish_prompt shpool__old_prompt + function __shpool_set_status; return $argv[1]; end + set -l __shpool_nt_err (test /dev/null -nt /dev/null 2>&1) + if test -z "$__shpool_nt_err" + function __shpool_is_newer; test $argv[1] -nt $argv[2]; end + else + function __shpool_is_newer; command test $argv[1] -nt $argv[2]; end + end + function fish_prompt + set -l last_status $status + set -l env_file "$SHPOOL_SESSION_DIR/forward.env" + set -l stamp_file "$env_file.stamp" + if test -n "$SHPOOL_SESSION_DIR"; and test -f "$env_file" + if test ! -f "$stamp_file"; or __shpool_is_newer "$env_file" "$stamp_file" + touch -r "$env_file" "$stamp_file" 2>/dev/null + source "$env_file" + end + end + echo -n "{prompt_prefix}" + __shpool_set_status $last_status + shpool__old_prompt + end + "# + ) + } (_, Err(e)) => { warn!("could not sniff shell: {}", e); diff --git a/shpool/tests/attach.rs b/shpool/tests/attach.rs index a31a465d..ad883685 100644 --- a/shpool/tests/attach.rs +++ b/shpool/tests/attach.rs @@ -5,6 +5,7 @@ use std::{ fs, io::BufRead, io::{Read, Write}, + path::PathBuf, process::{Command, Stdio}, thread, time, }; @@ -181,6 +182,300 @@ fn forward_env() -> anyhow::Result<()> { Ok(()) } +#[test] +#[timeout(30000)] +fn forward_env_live_reload_bash() -> anyhow::Result<()> { + let mut daemon_proc = + support::daemon::Proc::new("forward_env_bash.toml", DaemonArgs::default()) + .context("starting daemon proc")?; + + let mut waiter = daemon_proc + .events + .take() + .unwrap() + .waiter(["daemon-wrote-s2c-chunk", "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")), + ], + ..Default::default() + }, + ) + .context("starting attach proc")?; + + let mut lm = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=foo:bar$")?; + } + + daemon_proc.events = Some(waiter.wait_final_event("daemon-bidi-stream-done")?); + + // Reattach with updated environment variables + let bidi_done_w2 = 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("foonew")), + (String::from("BAR"), String::from("bar with 'quotes'")), + ], + ..Default::default() + }, + ) + .context("reattaching proc")?; + + let mut lm = attach_proc.line_matcher()?; + + // Give the reattach resize protocol (REATTACH_RESIZE_DELAY 50ms) time + // to complete so SIGWINCH handling doesn't interfere with command + // input. + thread::sleep(time::Duration::from_millis(500)); + + // Trigger a prompt cycle so PROMPT_COMMAND runs with the new + // forward.env + attach_proc.run_cmd("echo prompt_cycle_reattach")?; + lm.scan_until_re("prompt_cycle_reattach$")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=foonew:bar with 'quotes'$")?; + + // Test external modification of forward.env + attach_proc.run_cmd(r#"echo "session_dir=$SHPOOL_SESSION_DIR" "#)?; + let caps = lm.scan_until_re_captures(r#"session_dir=(/[^\s\r\n"]+)"#)?; + let session_dir = PathBuf::from(caps[1].as_ref().unwrap()); + let forward_env_path = session_dir.join("forward.env"); + + // Sleep for at least 1.1s so that older shells with 1-second timestamp + // resolution on `test -nt` (like bash 3.2 on macOS) see a distinct + // mtime. + thread::sleep(time::Duration::from_millis(1100)); + fs::write(&forward_env_path, "export FOO='external_foo'\nexport BAR='external_bar'\n")?; + + attach_proc.run_cmd("echo prompt_cycle_external")?; + lm.scan_until_re("prompt_cycle_external$")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=external_foo:external_bar$")?; + + // Test deletion of forward.env: shell continues functioning normally + fs::remove_file(&forward_env_path)?; + attach_proc.run_cmd("echo prompt_cycle_deleted")?; + lm.scan_until_re("prompt_cycle_deleted$")?; + attach_proc.run_cmd(r#"echo "deleted_ok" "#)?; + lm.scan_until_re("deleted_ok$")?; + + // Recreating forward.env reloads again + thread::sleep(time::Duration::from_millis(1100)); + fs::write(&forward_env_path, "export FOO='recreated_foo'\nexport BAR='recreated_bar'\n")?; + attach_proc.run_cmd("echo prompt_cycle_recreated")?; + lm.scan_until_re("prompt_cycle_recreated$")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=recreated_foo:recreated_bar$")?; + } + + daemon_proc.events = Some(bidi_done_w2.wait_final_event("daemon-bidi-stream-done")?); + + Ok(()) +} + +#[test] +#[timeout(30000)] +#[cfg_attr(target_os = "macos", ignore)] +fn forward_env_live_reload_zsh() -> anyhow::Result<()> { + let mut daemon_proc = support::daemon::Proc::new("forward_env_zsh.toml", DaemonArgs::default()) + .context("starting daemon proc")?; + + let mut waiter = daemon_proc + .events + .take() + .unwrap() + .waiter(["daemon-wrote-s2c-chunk", "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")), + ], + ..Default::default() + }, + ) + .context("starting attach proc")?; + + let mut lm = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=foo:bar$")?; + } + + daemon_proc.events = Some(waiter.wait_final_event("daemon-bidi-stream-done")?); + + // Reattach with updated environment variables + let bidi_done_w2 = 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("foonew")), + (String::from("BAR"), String::from("bar with 'quotes'")), + ], + ..Default::default() + }, + ) + .context("reattaching proc")?; + + let mut lm = attach_proc.line_matcher()?; + + // Give the reattach resize protocol (REATTACH_RESIZE_DELAY 50ms) time + // to complete so SIGWINCH handling doesn't interfere with command + // input. + thread::sleep(time::Duration::from_millis(500)); + + // Trigger a prompt cycle so precmd_functions runs with the new + // forward.env + attach_proc.run_cmd("echo prompt_cycle_reattach")?; + lm.scan_until_re("prompt_cycle_reattach$")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=foonew:bar with 'quotes'$")?; + + // Test external modification of forward.env + attach_proc.run_cmd(r#"echo "session_dir=$SHPOOL_SESSION_DIR" "#)?; + let caps = lm.scan_until_re_captures(r#"session_dir=(/[^\s\r\n"]+)"#)?; + let session_dir = PathBuf::from(caps[1].as_ref().unwrap()); + let forward_env_path = session_dir.join("forward.env"); + + thread::sleep(time::Duration::from_millis(50)); + fs::write(&forward_env_path, "export FOO='external_foo'\nexport BAR='external_bar'\n")?; + + attach_proc.run_cmd("echo prompt_cycle_external")?; + lm.scan_until_re("prompt_cycle_external$")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=external_foo:external_bar$")?; + + // Test deletion of forward.env: shell continues functioning normally + fs::remove_file(&forward_env_path)?; + attach_proc.run_cmd("echo prompt_cycle_deleted")?; + lm.scan_until_re("prompt_cycle_deleted$")?; + attach_proc.run_cmd(r#"echo "deleted_ok" "#)?; + lm.scan_until_re("deleted_ok$")?; + + // Recreating forward.env reloads again + thread::sleep(time::Duration::from_millis(50)); + fs::write(&forward_env_path, "export FOO='recreated_foo'\nexport BAR='recreated_bar'\n")?; + attach_proc.run_cmd("echo prompt_cycle_recreated")?; + lm.scan_until_re("prompt_cycle_recreated$")?; + attach_proc.run_cmd(r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=recreated_foo:recreated_bar$")?; + } + + daemon_proc.events = Some(bidi_done_w2.wait_final_event("daemon-bidi-stream-done")?); + + Ok(()) +} + +#[test] +#[timeout(30000)] +#[cfg_attr(target_os = "macos", ignore)] +fn forward_env_live_reload_fish() -> anyhow::Result<()> { + let mut daemon_proc = + support::daemon::Proc::new("forward_env_fish.toml", DaemonArgs::default()) + .context("starting daemon proc")?; + + let run_fish = |proc: &mut support::attach::Proc, cmd: &str| -> anyhow::Result<()> { + proc.run_raw(format!("{cmd}\r").into_bytes()) + }; + + 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")), + ], + ..Default::default() + }, + ) + .context("starting attach proc")?; + + let mut lm = attach_proc.line_matcher()?; + thread::sleep(time::Duration::from_millis(800)); + run_fish(&mut attach_proc, r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=foo:bar$")?; + } + + daemon_proc.events = Some(bidi_done_w.wait_final_event("daemon-bidi-stream-done")?); + + // Reattach with updated environment variables + let bidi_done_w2 = 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("foonew")), + (String::from("BAR"), String::from("bar with 'quotes'")), + ], + ..Default::default() + }, + ) + .context("reattaching proc")?; + + let mut lm = attach_proc.line_matcher()?; + thread::sleep(time::Duration::from_millis(500)); + + // Trigger a prompt cycle so fish_prompt runs with the new forward.env + run_fish(&mut attach_proc, "echo prompt_cycle_reattach")?; + lm.scan_until_re("prompt_cycle_reattach$")?; + run_fish(&mut attach_proc, r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=foonew:bar with 'quotes'$")?; + + // Test external modification of forward.env + run_fish(&mut attach_proc, r#"echo "session_dir=$SHPOOL_SESSION_DIR" "#)?; + let caps = lm.scan_until_re_captures(r#"session_dir=(/[^\s\r\n"]+)"#)?; + let session_dir = PathBuf::from(caps[1].as_ref().unwrap()); + let forward_env_path = session_dir.join("forward.env"); + + thread::sleep(time::Duration::from_millis(50)); + fs::write(&forward_env_path, "export FOO='external_foo'\nexport BAR='external_bar'\n")?; + + run_fish(&mut attach_proc, "echo prompt_cycle_ext")?; + lm.scan_until_re("prompt_cycle_ext$")?; + run_fish(&mut attach_proc, r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=external_foo:external_bar$")?; + + // Test deletion of forward.env: shell continues functioning normally + fs::remove_file(&forward_env_path)?; + run_fish(&mut attach_proc, "echo deleted_ok")?; + lm.scan_until_re("deleted_ok$")?; + + // Recreating forward.env reloads again + thread::sleep(time::Duration::from_millis(50)); + fs::write(&forward_env_path, "export FOO='recreated_foo'\nexport BAR='recreated_bar'\n")?; + run_fish(&mut attach_proc, "echo prompt_cycle_recreate")?; + lm.scan_until_re("prompt_cycle_recreate$")?; + run_fish(&mut attach_proc, r#"echo "val=$FOO:$BAR" "#)?; + lm.scan_until_re("val=recreated_foo:recreated_bar$")?; + } + + daemon_proc.events = Some(bidi_done_w2.wait_final_event("daemon-bidi-stream-done")?); + + 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 @@ -942,20 +1237,26 @@ fn has_right_default_path() -> anyhow::Result<()> { fn screen_restore() -> anyhow::Result<()> { let mut daemon_proc = support::daemon::Proc::new("restore_screen.toml", DaemonArgs::default()) .context("starting daemon proc")?; - let bidi_done_w = daemon_proc.events.take().unwrap().waiter(["daemon-bidi-stream-done"]); + let mut waiter = daemon_proc + .events + .take() + .unwrap() + .waiter(["daemon-wrote-s2c-chunk", "daemon-bidi-stream-done"]); { let mut attach_proc = daemon_proc.attach("sh1", Default::default()).context("starting attach proc")?; let mut line_matcher = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; + attach_proc.run_cmd("echo foo")?; line_matcher.scan_until_re("foo$")?; } // 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")?); + daemon_proc.events = Some(waiter.wait_final_event("daemon-bidi-stream-done")?); { let mut attach_proc = @@ -978,20 +1279,26 @@ fn screen_restore() -> anyhow::Result<()> { fn screen_wide_restore() -> anyhow::Result<()> { let mut daemon_proc = support::daemon::Proc::new("restore_screen.toml", DaemonArgs::default()) .context("starting daemon proc")?; - let bidi_done_w = daemon_proc.events.take().unwrap().waiter(["daemon-bidi-stream-done"]); + let mut waiter = daemon_proc + .events + .take() + .unwrap() + .waiter(["daemon-wrote-s2c-chunk", "daemon-bidi-stream-done"]); { let mut attach_proc = daemon_proc.attach("sh1", Default::default()).context("starting attach proc")?; let mut line_matcher = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; + attach_proc.run_cmd("echo ooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooy")?; line_matcher.scan_until_re("ooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooyooooxooooy$")?; } // 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")?); + daemon_proc.events = Some(waiter.wait_final_event("daemon-bidi-stream-done")?); { let mut attach_proc = @@ -1051,26 +1358,37 @@ fn lines_restore() -> anyhow::Result<()> { fn screen_restore_input_modes() -> anyhow::Result<()> { let mut daemon_proc = support::daemon::Proc::new("restore_screen.toml", DaemonArgs::default()) .context("starting daemon proc")?; - let bidi_done_w = daemon_proc.events.take().unwrap().waiter(["daemon-bidi-stream-done"]); + let mut waiter = daemon_proc + .events + .take() + .unwrap() + .waiter(["daemon-wrote-s2c-chunk", "daemon-bidi-stream-done"]); { let mut attach_proc = daemon_proc.attach("sh1", Default::default()).context("starting attach proc")?; let mut line_matcher = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; + attach_proc.run_cmd("printf '\\033[?1000h'; echo modes-on")?; line_matcher.scan_until_re("modes-on$")?; } // 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")?); + daemon_proc.events = Some(waiter.wait_final_event("daemon-bidi-stream-done")?); { let mut attach_proc = daemon_proc.attach("sh1", Default::default()).context("starting attach proc")?; let mut line_matcher = attach_proc.line_matcher()?; + // Give the reattach resize protocol (REATTACH_RESIZE_DELAY 50ms) time + // to complete so SIGWINCH handling doesn't interfere with command + // input. + thread::sleep(time::Duration::from_millis(500)); + // The restore buffer does not end in a newline, so give the line // matcher something that does before asserting on it. Nothing in the // shell emits mouse reporting, so a match can only come from the @@ -1213,9 +1531,11 @@ fn prompt_prefix_bash() -> anyhow::Result<()> { let mut daemon_proc = support::daemon::Proc::new("prompt_prefix_bash.toml", DaemonArgs::default()) .context("starting daemon proc")?; + let mut waiter = daemon_proc.events.take().unwrap().waiter(["daemon-wrote-s2c-chunk"]); let mut attach_proc = daemon_proc.attach("sh1", AttachArgs::default())?; let mut lm = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; attach_proc.run_cmd("echo")?; lm.scan_until_re(".*session_name=sh1 prompt>.*")?; @@ -1229,9 +1549,11 @@ fn prompt_prefix_zsh() -> anyhow::Result<()> { let mut daemon_proc = support::daemon::Proc::new("prompt_prefix_zsh.toml", DaemonArgs::default()) .context("starting daemon proc")?; + let mut waiter = daemon_proc.events.take().unwrap().waiter(["daemon-wrote-s2c-chunk"]); let mut attach_proc = daemon_proc.attach("sh1", AttachArgs::default())?; let mut lm = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; attach_proc.run_cmd("echo")?; lm.scan_until_re(".*session_name=sh1.*")?; @@ -1892,6 +2214,7 @@ fn templated_session_name_no_switch_on_unrelated_var() -> anyhow::Result<()> { fn start_cmd() -> anyhow::Result<()> { let mut daemon_proc = support::daemon::Proc::new("norc.toml", DaemonArgs::default()) .context("starting daemon proc")?; + let mut waiter = daemon_proc.events.take().unwrap().waiter(["daemon-wrote-s2c-chunk"]); let mut attach_proc = daemon_proc .attach( "sh1", @@ -1903,6 +2226,7 @@ fn start_cmd() -> anyhow::Result<()> { .context("starting attach proc")?; let mut line_matcher = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; attach_proc.run_cmd("echo $STARTUP_CMD_RAN")?; line_matcher.scan_until_re("true$")?; @@ -1915,6 +2239,7 @@ fn start_cmd() -> anyhow::Result<()> { fn start_cmd_template() -> anyhow::Result<()> { let mut daemon_proc = support::daemon::Proc::new("norc.toml", DaemonArgs::default()) .context("starting daemon proc")?; + let mut waiter = daemon_proc.events.take().unwrap().waiter(["daemon-wrote-s2c-chunk"]); daemon_proc.var_set("myvar", "myval")?; @@ -1929,6 +2254,7 @@ fn start_cmd_template() -> anyhow::Result<()> { .context("starting attach proc")?; let mut line_matcher = attach_proc.line_matcher()?; + waiter.wait_event("daemon-wrote-s2c-chunk")?; attach_proc.run_cmd("echo $STARTUP_CMD_VAR")?; line_matcher.scan_until_re("myval$")?; diff --git a/shpool/tests/data/forward_env_bash.toml b/shpool/tests/data/forward_env_bash.toml new file mode 100644 index 00000000..0c8e139e --- /dev/null +++ b/shpool/tests/data/forward_env_bash.toml @@ -0,0 +1,10 @@ +norc = true +noecho = true +shell = "/bin/bash" +session_restore_mode = "simple" +prompt_prefix = "shpool> " +forward_env = ["FOO", "BAR"] + +[env] +PS1 = "prompt> " +TERM = "" diff --git a/shpool/tests/data/forward_env_fish.toml b/shpool/tests/data/forward_env_fish.toml new file mode 100644 index 00000000..e748643f --- /dev/null +++ b/shpool/tests/data/forward_env_fish.toml @@ -0,0 +1,9 @@ +norc = true +noecho = true +shell = "/usr/bin/fish" +session_restore_mode = "simple" +prompt_prefix = "shpool> " +forward_env = ["FOO", "BAR"] + +[env] +TERM = "" diff --git a/shpool/tests/data/forward_env_zsh.toml b/shpool/tests/data/forward_env_zsh.toml new file mode 100644 index 00000000..53b5eb9f --- /dev/null +++ b/shpool/tests/data/forward_env_zsh.toml @@ -0,0 +1,9 @@ +norc = true +noecho = true +shell = "/usr/bin/zsh" +session_restore_mode = "simple" +prompt_prefix = "shpool> " +forward_env = ["FOO", "BAR"] + +[env] +TERM = "" diff --git a/shpool/tests/support/line_matcher.rs b/shpool/tests/support/line_matcher.rs index 45b72a8a..307d547a 100644 --- a/shpool/tests/support/line_matcher.rs +++ b/shpool/tests/support/line_matcher.rs @@ -3,7 +3,7 @@ use std::{io, io::BufRead, time}; use anyhow::{anyhow, Context}; use regex::Regex; -const CMD_READ_TIMEOUT: time::Duration = time::Duration::from_secs(3); +const CMD_READ_TIMEOUT: time::Duration = time::Duration::from_secs(5); const CMD_READ_SLEEP_DUR: time::Duration = time::Duration::from_millis(20); pub struct LineMatcher { @@ -37,7 +37,7 @@ where /// Scan lines until one matches the given regex, returning its captures. pub fn scan_until_re_captures(&mut self, re: &str) -> anyhow::Result>> { let compiled_re = Regex::new(re)?; - let start = time::Instant::now(); + let mut start = time::Instant::now(); let mut line = String::new(); loop { match self.out.read_line(&mut line) { @@ -57,6 +57,7 @@ where return Err(e).context("reading line from shell output")?; } Ok(_) => { + start = time::Instant::now(); if line.ends_with('\n') { line.pop(); if line.ends_with('\r') { @@ -136,7 +137,7 @@ where /// Scan through all the remaining lines and ensure that no persistant /// assertions fail (the never match regex). pub fn drain(&mut self) -> anyhow::Result<()> { - let start = time::Instant::now(); + let mut start = time::Instant::now(); let mut line = String::new(); loop { match self.out.read_line(&mut line) { @@ -156,6 +157,7 @@ where return Err(e).context("reading line from shell output")?; } Ok(_) => { + start = time::Instant::now(); if line.ends_with('\n') { line.pop(); if line.ends_with('\r') {