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') {