//! Persistent shell state across command invocations. //! //! After each command, the shell's environment //! (env vars, cwd, functions, aliases, options) is serialized via a dump script, //! and replayed before the next command. Each invocation is still a fresh process, //! but the user observes a single continuous session. //! //! State is transported via extra file descriptors (fd 3 for input, fd 4 for output) //! so that dump traffic never pollutes stdout/stderr. use std::os::unix::io::{AsRawFd, FromRawFd, OwnedFd}; use std::path::{Path, PathBuf}; use std::process::Stdio; use std::collections::HashMap; use std::time::Duration; use command_fds::FdMapping; use nix::libc; use tokio::io::AsyncReadExt; // ============================================================================ // Marker constants // ============================================================================ const BASH_STATE_START_MARKER: &str = "__GROK_BASH_STATE_START__"; const BASH_STATE_END_MARKER: &str = "__GROK_BASH_STATE_END__"; const ZSH_STATE_START_MARKER: &str = "__GROK_ZSH_STATE_START__"; const ZSH_STATE_END_MARKER: &str = "__GROK_ZSH_STATE_END__"; /// Marker emitted by the init path to separate login-shell noise (MOTD, etc.) /// from the actual state dump on stdout. const INIT_STATE_MARKER: &str = "__GROK_INIT_STATE_MARKER__"; /// Maximum time to wait for a shell state init (login shell + rc files). const INIT_TIMEOUT: Duration = Duration::from_secs(15); /// Maximum time to wait for the dump reader task after the child process exits. /// Uses a 5s close timeout. If a background process inherits fd 4, /// the reader would hang forever without this. const DUMP_READ_TIMEOUT: Duration = Duration::from_secs(5); /// Environment overrides applied to every agent terminal / persistent shell spawn. /// Prevents color/pager noise in captured output and marks the process as agent-driven. pub fn shell_env_overrides() -> HashMap { HashMap::from([ ("TERM".to_string(), "dumb".to_string()), ("NO_COLOR".to_string(), "1".to_string()), ("FORCE_COLOR".to_string(), "0".to_string()), // Re-applied last via [`crate::util::apply_grok_agent_marker`] so request // env cannot clear it. ( crate::util::KIGI_AGENT_ENV.to_string(), crate::util::KIGI_AGENT_ENV_VALUE.to_string(), ), ]) } /// Returns the sudo alias injection string if `SUDO_ASKPASS` is configured. /// When set, `alias sudo='sudo -A'` makes any `sudo` in the user's command /// use the askpass helper instead of blocking on tty input. fn sudo_alias_injection() -> String { match std::env::var("SUDO_ASKPASS") { Ok(val) if !val.is_empty() => "alias sudo='sudo -A'; ".to_string(), _ => String::new(), } } // ============================================================================ // Dump scripts (embedded as const strings) // ============================================================================ /// Bash state dump script. Captures env vars, POSIX options, bash options, /// functions, and aliases as base64-encoded replayable shell snippets. const DUMP_BASH_STATE_SCRIPT: &str = r##" dump_bash_state() { set -euo pipefail if ! command -v base64 >/dev/null 2>&1; then echo "Error: base64 command is required" >&2 return 1 fi _emit() { builtin printf '%s\n' "$1" } _emit_encoded() { local content="$1" local var_name="$2" if [[ -n "$content" ]]; then builtin printf 'grok_snap_%s=$(command base64 -d <<'"'"'KIGI_SNAP_EOF_%s'"'"'\n' "$var_name" "$var_name" command base64 <<<"$content" | command tr -d '\n' builtin printf '\nKIGI_SNAP_EOF_%s\n' "$var_name" builtin printf ')\n' builtin printf 'eval "$grok_snap_%s"\n' "$var_name" fi } _emit "__GROK_BASH_STATE_START__" _emit "$PWD" local env_vars env_vars=$(builtin export -p 2>/dev/null | command grep -viE '_proxy=|KIGI_SANDBOX|KIGI_AGENT=|SUDO_ASKPASS|KIGI_ASKPASS|ELECTRON_RUN_AS_NODE|SSH_AUTH_SOCK|DBUS_SESSION_BUS_ADDRESS|XDG_RUNTIME_DIR|WAYLAND_DISPLAY|GPG_TTY' || true) _emit_encoded "$env_vars" "ENV_VARS_B64" local posix_opts posix_opts=$(builtin shopt -po 2>/dev/null | command grep -v '^set -o nounset$' | command grep -v '^set +o nounset$' || true) _emit_encoded "$posix_opts" "POSIX_OPTS_B64" local bash_opts bash_opts=$(builtin shopt -p 2>/dev/null || true) _emit_encoded "$bash_opts" "BASH_OPTS_B64" local all_functions all_functions=$(builtin declare -f 2>/dev/null || true) _emit_encoded "$all_functions" "FUNCTIONS_B64" local aliases aliases=$(builtin alias -p 2>/dev/null || true) _emit_encoded "$aliases" "ALIASES_B64" _emit "# end of bash state dump" _emit "__GROK_BASH_STATE_END__" } "##; /// Zsh state dump script. Captures env vars, zsh options, functions, and aliases /// as base64-encoded replayable shell snippets. const DUMP_ZSH_STATE_SCRIPT: &str = r##" function dump_zsh_state() { emulate -L zsh -o errreturn -o pipefail set -u builtin zmodload -F zsh/parameter p:parameters p:options p:functions p:aliases p:galiases p:saliases 2>/dev/null || true _emit() { builtin print -r -- "$1" } _emit_encoded() { local content="$1" local var_name="$2" if [[ -n "$content" ]]; then builtin printf 'grok_snap_%s=$(command base64 -d <<'"'"'KIGI_SNAP_EOF_%s'"'"'\n' "$var_name" "$var_name" command base64 <<<"$content" | command tr -d '\n' builtin printf '\nKIGI_SNAP_EOF_%s\n' "$var_name" builtin printf ')\n' builtin printf 'eval "$grok_snap_%s"\n' "$var_name" fi } _emit "__GROK_ZSH_STATE_START__" _emit "$PWD" local env_vars env_vars=$(builtin typeset -xp 2>/dev/null | command grep -viE '_proxy=|KIGI_SANDBOX|KIGI_AGENT=|SUDO_ASKPASS|KIGI_ASKPASS|ELECTRON_RUN_AS_NODE|SSH_AUTH_SOCK|DBUS_SESSION_BUS_ADDRESS|XDG_RUNTIME_DIR|WAYLAND_DISPLAY|GPG_TTY' || true) _emit_encoded "$env_vars" "ENV_VARS_B64" local zsh_opts zsh_opts=$(setopt 2>/dev/null | command grep -v '^nounset$' | command awk '{printf "builtin setopt %s 2>/dev/null || true\n", $0}' || true) _emit_encoded "$zsh_opts" "ZSH_OPTS_B64" local all_functions all_functions=$(builtin typeset -f 2>/dev/null || true) _emit_encoded "$all_functions" "FUNCTIONS_B64" local aliases aliases=$({ builtin alias -L; builtin alias -gL; builtin alias -sL } 2>/dev/null || true) _emit_encoded "$aliases" "ALIASES_B64" _emit "# end of zsh state dump" _emit "__GROK_ZSH_STATE_END__" } "##; // ============================================================================ // Shell kind // ============================================================================ #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ShellKind { Bash, Zsh, } impl ShellKind { /// Detect the user's shell from `$SHELL`, falling back to bash. pub fn detect() -> Self { match kigi_config::shell::detect_unix_shell_kind() { kigi_config::shell::UnixShellKind::Bash => Self::Bash, kigi_config::shell::UnixShellKind::Zsh => Self::Zsh, } } /// Resolved absolute path to the shell binary. Falls back from `$SHELL` → /// `which` → common dirs → `/bin/`. Result is cached process-wide /// in `kigi_config::shell::unix_shell_path`. See that function for /// the full cascade. Returns `&'static str`. pub fn binary_path(&self) -> &'static str { let kind = match self { Self::Bash => kigi_config::shell::UnixShellKind::Bash, Self::Zsh => kigi_config::shell::UnixShellKind::Zsh, }; kigi_config::shell::unix_shell_path(kind) } /// The user's primary rc file name (relative to `$HOME`). pub fn rc_file_name(&self) -> &str { match self { Self::Bash => ".bashrc", Self::Zsh => ".zshrc", } } fn dump_script(&self) -> &'static str { match self { Self::Bash => DUMP_BASH_STATE_SCRIPT, Self::Zsh => DUMP_ZSH_STATE_SCRIPT, } } fn dump_function_name(&self) -> &str { match self { Self::Bash => "dump_bash_state", Self::Zsh => "dump_zsh_state", } } fn start_marker(&self) -> &str { match self { Self::Bash => BASH_STATE_START_MARKER, Self::Zsh => ZSH_STATE_START_MARKER, } } fn end_marker(&self) -> &str { match self { Self::Bash => BASH_STATE_END_MARKER, Self::Zsh => ZSH_STATE_END_MARKER, } } } // ============================================================================ // ShellState // ============================================================================ /// Persistent shell state: a serialized snapshot that can be replayed to restore /// env vars, cwd, functions, aliases, and shell options in a fresh shell process. #[derive(Debug, Clone)] pub struct ShellState { /// Current working directory (from the dump's first line). pub cwd: PathBuf, /// Replayable shell script (everything after the cwd line, minus markers). pub snapshot: String, /// Which shell produced this state. pub shell: ShellKind, } impl ShellState { /// Initialize shell state by running an interactive login shell that loads /// the user's rc files, then capturing the resulting environment. /// /// This is the expensive path — only called once (lazily on first command). pub async fn init( shell: ShellKind, cwd: &Path, ) -> Result { let dump_script = shell.dump_script(); let dump_fn = shell.dump_function_name(); // Build the one-liner: define the dump function, print a marker (to separate // login noise from our output), then call the dump function. let script = format!("{dump_script} builtin printf '{INIT_STATE_MARKER}\\n'; {dump_fn}"); let args: Vec<&str> = match shell { ShellKind::Bash => vec!["-O", "extglob", "-ilc", &script], ShellKind::Zsh => vec!["-o", "extendedglob", "-ilc", &script], }; // stderr is intentionally discarded (Stdio::null) — we never read it, // and piping it would risk a deadlock if the user's rc files write >64KB // to stderr (fills the pipe buffer, child blocks on write, parent blocks // on stdout read). let mut cmd = tokio::process::Command::new(shell.binary_path()); cmd.args(&args) .current_dir(cwd) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::null()) .kill_on_drop(true); crate::util::detach_command(&mut cmd); cmd.envs(crate::util::pager_env()); let mut child = cmd.spawn().map_err(|e| { crate::computer::types::ComputerError::io(format!( "failed to spawn {shell:?} for shell state init: {e}" )) })?; let mut full_output = String::new(); if let Some(ref mut stdout) = child.stdout { // Apply init timeout to prevent hangs from slow rc files (e.g. network mounts). match tokio::time::timeout(INIT_TIMEOUT, stdout.read_to_string(&mut full_output)).await { Ok(Ok(_)) => {} Ok(Err(e)) => { tracing::warn!("shell state init: read error: {e}"); } Err(_) => { tracing::warn!( "shell state init: timed out after {}s reading login shell output", INIT_TIMEOUT.as_secs() ); child.kill().await.ok(); } } } let _ = child.wait().await; // Extract output after our marker (skip MOTD, bashrc echo, etc.) let snapshot_raw = parse_after_marker(&full_output, INIT_STATE_MARKER); match parse_dump(shell, snapshot_raw) { Some((parsed_cwd, rest)) => Ok(Self { cwd: parsed_cwd, snapshot: rest, shell, }), None => { // Dump failed or was killed — use empty state with the given cwd. tracing::warn!("shell state init: dump markers missing, using empty state"); Ok(Self { cwd: cwd.to_path_buf(), snapshot: String::new(), shell, }) } } } /// Build the wrapper command and fd pipe pair for a persistent shell invocation. /// /// Returns `(shell_binary, args, state_in_fd, state_out_fd)` where: /// - `state_in_fd` is a pipe the caller writes the prior snapshot to (fd 3 in the child) /// - `state_out_fd` is a pipe the caller reads the new dump from (fd 4 in the child) /// /// The caller must: /// 1. Write `self.snapshot` to `state_in_fd` then close it /// 2. Spawn the child with the returned args + fd_mappings /// 3. Read `state_out_fd` after the child exits /// 4. Call `parse_dump()` on the result to update state /// /// `cwd_override`: if provided, the child uses this cwd instead of the /// persistent shell's tracked cwd. Used for per-call `working_directory` /// overrides from `TerminalRunRequest`. pub fn prepare_command( &self, user_command: &str, cwd_override: Option<&Path>, search_shadows: super::SearchShadowConfig, ) -> std::io::Result { let dump_script = self.shell.dump_script(); let dump_fn = self.shell.dump_function_name(); let sudo_inject = sudo_alias_injection(); let search_inject = super::embedded_search_tools::search_injection(search_shadows); // Create two OS pipes: one for state-in (fd 3), one for state-out (fd 4). // os_pipe() creates fds with O_CLOEXEC (atomically on Linux, // best-effort on macOS) so concurrent forks can't leak them. let (state_in_read, state_in_write) = os_pipe()?; let (state_out_read, state_out_write) = os_pipe()?; // Ensure the parent-only ends have CLOEXEC (redundant on Linux // where os_pipe uses pipe2, but needed as a safety net on macOS // where pipe+fcntl has a small race window). set_cloexec(&state_in_write)?; set_cloexec(&state_out_read)?; // The child-bound ends (state_in_read, state_out_write) also have // CLOEXEC from os_pipe(). This is fine: fd_mappings uses dup2() // which clears CLOEXEC on the target fd (3/4), so the child keeps // them across exec. The originals are closed on exec by CLOEXEC. // The wrapper command: // 1. Read prior snapshot from fd 3, eval it (restores env/funcs/aliases/opts) // 2. Run the user's command (passed as $1) // 3. Dump new state to fd 4 // 4. Exit with the user command's exit code let wrapper = match self.shell { ShellKind::Bash => format!( // Merge the user command's stderr into its // stdout via `2>&1` so the captured byte stream preserves // chronological write order. Without this, the bash tool's // separate stdout/stderr pipes (each read in lockstep) // emit all-of-stdout-then-all-of-stderr in a single poll // tick, so a command like `echo X 1>&2 && echo Y` shows // up as `Y\nX\n` instead of the chronological `X\nY\n`. // Shell-level diagnostics (eval syntax errors, etc.) still // land on the outer shell's stderr — those are unaffected. // Re-export KIGI_AGENT=1 after snapshot eval so agent-definition // selectors (or other values) from prior shells cannot clear the // agent sentinel (process env alone is insufficient). "{dump_script} \ snap=$(command cat <&3) && builtin shopt -s extglob && builtin eval -- \"$snap\" && \ {{ builtin set +u 2>/dev/null || true; \ builtin export KIGI_AGENT=1; \ builtin export PWD=\"$(builtin pwd)\"; \ builtin shopt -s expand_aliases 2>/dev/null; {sudo_inject}{search_inject}\ builtin eval \"$1\" 2>&1; }}; \ COMMAND_EXIT_CODE=$?; {dump_fn} >&4; builtin exit $COMMAND_EXIT_CODE" ), // After snapshot restore: force nonomatch so login dumps cannot re-arm NOMATCH for model globs. ShellKind::Zsh => format!( "{dump_script} \ snap=$(command cat <&3); \ builtin unsetopt aliases 2>/dev/null; \ builtin unalias -m '*' 2>/dev/null || true; \ builtin eval \"$snap\" && \ {{ builtin unsetopt nounset 2>/dev/null || true; \ builtin setopt nonomatch 2>/dev/null || true; \ builtin export KIGI_AGENT=1; \ builtin export PWD=\"$(builtin pwd)\"; \ builtin setopt aliases 2>/dev/null; {sudo_inject}{search_inject}\ builtin eval \"$1\" 2>&1; }}; \ COMMAND_EXIT_CODE=$?; {dump_fn} >&4; builtin exit $COMMAND_EXIT_CODE" ), }; let effective_cwd = cwd_override.unwrap_or(&self.cwd); let args: Vec = match self.shell { ShellKind::Bash => vec![ "-O".into(), "extglob".into(), "-c".into(), wrapper, "--".into(), user_command.into(), ], ShellKind::Zsh => vec!["-c".into(), wrapper, "--".into(), user_command.into()], }; let fd_mappings = vec![ FdMapping { parent_fd: state_in_read, child_fd: 3, }, FdMapping { parent_fd: state_out_write, child_fd: 4, }, ]; Ok(PreparedCommand { binary: self.shell.binary_path().to_string(), args, fd_mappings, state_in_write, state_out_read, cwd: effective_cwd.to_path_buf(), }) } /// Update this state from a raw dump string (read from fd 4). /// Returns `true` if the state was successfully updated. pub fn update_from_dump(&mut self, raw: &str) -> bool { match parse_dump(self.shell, raw) { Some((new_cwd, new_snapshot)) => { if new_cwd.is_absolute() { self.cwd = new_cwd; } self.snapshot = new_snapshot; true } None => { tracing::debug!("shell state dump markers missing, keeping previous state"); false } } } } /// Everything needed to spawn a persistent shell command. pub struct PreparedCommand { /// Resolved shell binary path (e.g. `/bin/bash`, `/opt/homebrew/bin/bash`, /// or `/run/current-system/sw/bin/bash` on NixOS). See /// [`ShellKind::binary_path`] for the resolution cascade. pub binary: String, /// Full argument list for the shell. pub args: Vec, /// Fd mappings to pass to `CommandFdExt::fd_mappings()`. pub fd_mappings: Vec, /// Write end of the state-input pipe. Caller writes `snapshot` here, then drops. pub state_in_write: OwnedFd, /// Read end of the state-output pipe. Caller reads the new dump from here after exit. pub state_out_read: OwnedFd, /// Working directory for the child process. pub cwd: PathBuf, } // ============================================================================ // Helpers // ============================================================================ /// Create an OS pipe, returning `(read_end, write_end)` as `OwnedFd`. /// /// On Linux, uses `nix::unistd::pipe2(O_CLOEXEC)` to atomically set /// close-on-exec, eliminating the race window between `pipe()` and /// `fcntl(F_SETFD)` where a concurrent `fork()` could leak fds to an /// unrelated child. /// /// On macOS, `pipe2` is not exposed by `nix` 0.30 (the kernel added it /// in 10.15 but `nix`'s cfg gate hasn't caught up). Falls back to /// `pipe()` + `fcntl(FD_CLOEXEC)` with a best-effort race window. fn os_pipe() -> std::io::Result<(OwnedFd, OwnedFd)> { // Linux: atomic O_CLOEXEC via pipe2. #[cfg(target_os = "linux")] { nix::unistd::pipe2(nix::fcntl::OFlag::O_CLOEXEC) .map_err(|e| std::io::Error::from_raw_os_error(e as i32)) } // macOS (and other non-Linux Unix): pipe() + fcntl best-effort. #[cfg(not(target_os = "linux"))] { let (read_fd, write_fd) = nix::unistd::pipe().map_err(|e| std::io::Error::from_raw_os_error(e as i32))?; // Best-effort CLOEXEC — small race window on macOS between pipe() // and these fcntl calls, but unavoidable without pipe2. let _ = set_cloexec(&read_fd); let _ = set_cloexec(&write_fd); Ok((read_fd, write_fd)) } } /// Set FD_CLOEXEC on a file descriptor so it is NOT inherited by child processes. /// /// This is critical for pipe fds that should stay parent-only: without CLOEXEC, /// the child inherits both ends of a pipe after fork, preventing EOF from being /// signaled when the parent closes its end. fn set_cloexec(fd: &OwnedFd) -> std::io::Result<()> { let raw = fd.as_raw_fd(); let flags = unsafe { libc::fcntl(raw, libc::F_GETFD) }; if flags < 0 { return Err(std::io::Error::last_os_error()); } let ret = unsafe { libc::fcntl(raw, libc::F_SETFD, flags | libc::FD_CLOEXEC) }; if ret < 0 { return Err(std::io::Error::last_os_error()); } Ok(()) } /// Parse a state dump, validating start/end markers. /// Returns `(cwd, snapshot_rest)` or `None` if markers are missing. fn parse_dump(shell: ShellKind, raw: &str) -> Option<(PathBuf, String)> { let start = shell.start_marker(); let end = shell.end_marker(); let start_line = format!("{start}\n"); let end_line = format!("{end}\n"); if !raw.starts_with(&start_line) || !raw.ends_with(&end_line) { return None; } // Strip markers let without_markers = &raw[start_line.len()..raw.len() - end_line.len()]; // First line is $PWD let newline_pos = without_markers.find('\n')?; let cwd = &without_markers[..newline_pos]; let rest = &without_markers[newline_pos..]; // includes the leading \n Some((PathBuf::from(cwd), rest.to_string())) } /// Extract the portion of `output` after the first occurrence of `marker\n`. /// If the marker is not found, returns the full output. fn parse_after_marker<'a>(output: &'a str, marker: &str) -> &'a str { let needle = format!("{marker}\n"); match output.find(&needle) { Some(idx) => &output[idx + needle.len()..], None => output, } } /// Write the snapshot to the state-in pipe, then close the fd. /// Uses blocking I/O on a dedicated thread (pipes are not regular files). pub async fn write_snapshot_to_pipe(snapshot: &str, fd: OwnedFd) -> std::io::Result<()> { let data = snapshot.to_string(); tokio::task::spawn_blocking(move || { use std::io::Write; // Safety: we own the fd. let mut file = unsafe { std::fs::File::from_raw_fd(fd.as_raw_fd()) }; std::mem::forget(fd); file.write_all(data.as_bytes())?; file.flush()?; drop(file); // closes the fd → child sees EOF on its read end Ok(()) }) .await .map_err(std::io::Error::other)? } /// Read the full dump output from the state-out pipe with a timeout. /// /// If a background process inherits fd 4, the pipe never closes and the reader /// hangs. The timeout (5s close timeout) prevents this from /// blocking the actor loop forever. On timeout, returns whatever was read so far /// (which is typically empty, so marker validation will fail and prior state is kept). pub async fn read_dump_from_pipe(fd: OwnedFd) -> std::io::Result { // Read until either of the END markers appears, *not* until EOF. // // When the user's command backgrounds a subprocess (`cmd &`), the bg // shell inherits fd 4 (the dump pipe's write-end) and keeps it open // until *it* exits. The parent shell finishes its dump and exits, but // the kernel doesn't close the read-end's EOF until every write-end // holder closes theirs. Without marker-driven termination we'd block // on `read_to_string` for the entire bg lifetime, hit the 5s safety // timeout, and discard the (perfectly complete) dump — which manifests // as `cd` / function / alias state silently failing to persist after // any command that backgrounds something. (See harness scenarios // "State persistence after backgrounded command" and the cd-roundtrip // tests for shell state persistence parity.) // // We additionally cap on `DUMP_READ_TIMEOUT` so a shell that crashed // before emitting the END marker doesn't wedge the actor. match tokio::time::timeout( DUMP_READ_TIMEOUT, tokio::task::spawn_blocking(move || { use std::io::Read; let mut file = unsafe { std::fs::File::from_raw_fd(fd.as_raw_fd()) }; std::mem::forget(fd); let mut buf = String::new(); let mut chunk = [0u8; 4096]; loop { let n = file.read(&mut chunk)?; if n == 0 { // EOF: every write-end holder closed fd 4 (the // expected path when no bg subprocess was spawned). break; } buf.push_str(&String::from_utf8_lossy(&chunk[..n])); // Either marker suffices; we accept whichever shell the // child happens to be (bash vs zsh). if buf.contains(BASH_STATE_END_MARKER) || buf.contains(ZSH_STATE_END_MARKER) { break; } } drop(file); Ok(buf) }), ) .await { Ok(join_result) => join_result.map_err(std::io::Error::other)?, Err(_timeout) => { tracing::warn!( "shell state dump read timed out after {}s (END marker never arrived)", DUMP_READ_TIMEOUT.as_secs() ); Ok(String::new()) } } } // ============================================================================ // Tests // ============================================================================ #[cfg(test)] mod tests { use super::*; #[test] fn shell_env_overrides_marks_agent_terminal() { let env = shell_env_overrides(); assert_eq!( env.get(crate::util::KIGI_AGENT_ENV).map(String::as_str), Some(crate::util::KIGI_AGENT_ENV_VALUE) ); assert_eq!(env.get("TERM").map(String::as_str), Some("dumb")); assert_eq!(env.get("NO_COLOR").map(String::as_str), Some("1")); } #[test] fn parse_dump_valid_bash() { let raw = "__GROK_BASH_STATE_START__\n\ /home/user/project\n\ export FOO=bar\n\ # end of bash state dump\n\ __GROK_BASH_STATE_END__\n"; let (cwd, rest) = parse_dump(ShellKind::Bash, raw).unwrap(); assert_eq!(cwd, PathBuf::from("/home/user/project")); assert!(rest.contains("export FOO=bar")); } #[test] fn parse_dump_valid_zsh() { let raw = "__GROK_ZSH_STATE_START__\n\ /tmp\n\ typeset -x FOO=bar\n\ # end of zsh state dump\n\ __GROK_ZSH_STATE_END__\n"; let (cwd, rest) = parse_dump(ShellKind::Zsh, raw).unwrap(); assert_eq!(cwd, PathBuf::from("/tmp")); assert!(rest.contains("typeset -x FOO=bar")); } #[test] fn parse_dump_missing_start_marker() { let raw = "/home/user\nexport FOO=bar\n__GROK_BASH_STATE_END__\n"; assert!(parse_dump(ShellKind::Bash, raw).is_none()); } #[test] fn parse_dump_missing_end_marker() { let raw = "__GROK_BASH_STATE_START__\n/home/user\nexport FOO=bar\n"; assert!(parse_dump(ShellKind::Bash, raw).is_none()); } #[test] fn parse_dump_wrong_shell_markers() { let raw = "__GROK_ZSH_STATE_START__\n/tmp\nstuff\n__GROK_ZSH_STATE_END__\n"; assert!(parse_dump(ShellKind::Bash, raw).is_none()); } #[test] fn parse_dump_empty_snapshot() { let raw = "__GROK_BASH_STATE_START__\n\ /home/user\n\ # end of bash state dump\n\ __GROK_BASH_STATE_END__\n"; let (cwd, rest) = parse_dump(ShellKind::Bash, raw).unwrap(); assert_eq!(cwd, PathBuf::from("/home/user")); assert!(rest.contains("# end of bash state dump")); } #[test] fn parse_after_marker_found() { let output = "Welcome to Ubuntu\nMOTD line\n__GROK_INIT_STATE_MARKER__\nactual data\n"; let result = parse_after_marker(output, "__GROK_INIT_STATE_MARKER__"); assert_eq!(result, "actual data\n"); } #[test] fn parse_after_marker_not_found() { let output = "just some output\n"; let result = parse_after_marker(output, "__GROK_INIT_STATE_MARKER__"); assert_eq!(result, output); } /// Returns true iff a usable bash binary exists at the resolved path. /// Used to gate integration tests so they're skipped (rather than failing) /// on systems where bash isn't installed (e.g. minimal containers). /// On NixOS the resolver returns the nix-store / profile path, so this /// guard works there too. fn bash_available() -> bool { std::path::Path::new(ShellKind::Bash.binary_path()).exists() } /// Returns true iff a usable zsh binary exists at the resolved path. /// Mirrors [`bash_available`] so zsh integration tests skip (rather than /// fail) on systems without zsh installed. fn zsh_available() -> bool { std::path::Path::new(ShellKind::Zsh.binary_path()).exists() } #[test] fn shell_kind_binary_path_resolves_to_correct_kind() { // The resolver may pick any absolute path (e.g. `/opt/homebrew/bin/bash` // on macOS+brew, `/run/current-system/sw/bin/bash` on NixOS), but the // file name must match the requested kind. let bash = std::path::Path::new(ShellKind::Bash.binary_path()) .file_name() .and_then(|n| n.to_str()) .unwrap_or_default() .to_string(); assert_eq!(bash, "bash", "Bash kind must resolve to a *bash* binary"); let zsh = std::path::Path::new(ShellKind::Zsh.binary_path()) .file_name() .and_then(|n| n.to_str()) .unwrap_or_default() .to_string(); assert_eq!(zsh, "zsh", "Zsh kind must resolve to a *zsh* binary"); } #[test] fn shell_state_update_from_dump() { let mut state = ShellState { cwd: PathBuf::from("/old"), snapshot: String::new(), shell: ShellKind::Bash, }; let dump = "__GROK_BASH_STATE_START__\n\ /new/dir\n\ export X=1\n\ # end of bash state dump\n\ __GROK_BASH_STATE_END__\n"; assert!(state.update_from_dump(dump)); assert_eq!(state.cwd, PathBuf::from("/new/dir")); assert!(state.snapshot.contains("export X=1")); } #[test] fn shell_state_update_keeps_previous_on_bad_dump() { let mut state = ShellState { cwd: PathBuf::from("/original"), snapshot: "old stuff".into(), shell: ShellKind::Bash, }; assert!(!state.update_from_dump("garbage output\n")); assert_eq!(state.cwd, PathBuf::from("/original")); assert_eq!(state.snapshot, "old stuff"); } #[tokio::test] async fn test_init_bash() { // Integration test: actually runs bash and captures state. // Skip in environments without bash. if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); assert!(state.cwd.is_absolute()); // The snapshot should contain at least some env var exports assert!( state.snapshot.contains("grok_snap_") || state.snapshot.is_empty(), "snapshot should contain encoded blocks or be empty: {:?}", &state.snapshot[..state.snapshot.len().min(200)] ); } #[tokio::test] async fn test_prepare_command_and_roundtrip() { use command_fds::CommandFdExt; // Integration test: prepare a command, spawn it, verify state roundtrip. if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); // Run "export KIGI_TEST_VAR=hello" and capture the new state let prep = state .prepare_command( "export KIGI_TEST_VAR=hello", None, crate::computer::local::SearchShadowConfig::default(), ) .unwrap(); let mut cmd = tokio::process::Command::new(&prep.binary); cmd.args(&prep.args) .current_dir(&prep.cwd) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true); cmd.fd_mappings(prep.fd_mappings).unwrap(); let child = cmd.spawn().unwrap(); // Drop cmd to release the FdMapping OwnedFds held in its pre_exec closure. // Without this, the parent keeps the write-end of the state-out pipe open, // and the read task never sees EOF. drop(cmd); // Write snapshot to fd 3 let snapshot = state.snapshot.clone(); let write_handle = tokio::spawn( async move { write_snapshot_to_pipe(&snapshot, prep.state_in_write).await }, ); // Read new dump from fd 4 let read_handle = tokio::spawn(async move { read_dump_from_pipe(prep.state_out_read).await }); let output = child.wait_with_output().await.unwrap(); assert!(output.status.success(), "command failed: {:?}", output); write_handle.await.unwrap().unwrap(); let dump = read_handle.await.unwrap().unwrap(); assert!( state.update_from_dump(&dump), "dump should have valid markers, got: {:?}", &dump[..dump.len().min(500)] ); // The snapshot contains base64-encoded env vars, so the variable name // won't appear in plaintext. Verify the dump was valid and non-empty. assert!( !state.snapshot.is_empty(), "snapshot should be non-empty after a successful command" ); assert!(state.cwd.is_absolute(), "cwd should be absolute"); } /// Helper: run a command against a ShellState, update state, return (exit_code, stdout). async fn run_command(state: &mut ShellState, command: &str) -> (i32, String) { use command_fds::CommandFdExt; let prep = state .prepare_command( command, None, crate::computer::local::SearchShadowConfig::default(), ) .unwrap(); let mut cmd = tokio::process::Command::new(&prep.binary); cmd.args(&prep.args) .current_dir(&prep.cwd) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true); cmd.fd_mappings(prep.fd_mappings).unwrap(); let child = cmd.spawn().unwrap(); drop(cmd); let snapshot = state.snapshot.clone(); let write_handle = tokio::spawn( async move { write_snapshot_to_pipe(&snapshot, prep.state_in_write).await }, ); let read_handle = tokio::spawn(async move { read_dump_from_pipe(prep.state_out_read).await }); let output = child.wait_with_output().await.unwrap(); write_handle.await.unwrap().unwrap(); let dump = read_handle.await.unwrap().unwrap(); state.update_from_dump(&dump); let code = output.status.code().unwrap_or(-1); let stdout = String::from_utf8_lossy(&output.stdout).to_string(); (code, stdout) } #[tokio::test] async fn test_cd_persists_across_commands() { if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); // cd to /tmp (macOS resolves to /private/tmp via symlink) let (code, _) = run_command(&mut state, "cd /tmp").await; assert_eq!(code, 0); // Next command should see the resolved /tmp as cwd let (code, stdout) = run_command(&mut state, "pwd").await; assert_eq!(code, 0); let actual_pwd = stdout.trim(); assert!( actual_pwd == "/tmp" || actual_pwd == "/private/tmp", "cwd should be /tmp or /private/tmp, got: {actual_pwd}" ); assert_eq!(state.cwd.to_str().unwrap(), actual_pwd); } #[tokio::test] async fn test_env_var_persists_across_commands() { if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); // Export a variable let (code, _) = run_command(&mut state, "export MY_TEST_VAR=persistent_value").await; assert_eq!(code, 0); // Next command should see it let (code, stdout) = run_command(&mut state, "echo $MY_TEST_VAR").await; assert_eq!(code, 0); assert_eq!(stdout.trim(), "persistent_value"); } /// GPG_TTY exported in one command must not be replayed into the next via the snapshot. #[tokio::test] async fn test_gpg_tty_excluded_from_snapshot_bash() { if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); let (code, _) = run_command(&mut state, "export GPG_TTY=/grok-sentinel-tty").await; assert_eq!(code, 0); let (code, stdout) = run_command(&mut state, "echo \"[$GPG_TTY]\"").await; assert_eq!(code, 0); assert!( !stdout.contains("/grok-sentinel-tty"), "GPG_TTY must not persist across commands via the snapshot, got: {stdout:?}" ); } /// Same as the bash case, exercising the zsh `typeset -xp` dump grep. #[tokio::test] async fn test_gpg_tty_excluded_from_snapshot_zsh() { if !zsh_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Zsh, &cwd).await.unwrap(); let (code, _) = run_command(&mut state, "export GPG_TTY=/grok-sentinel-tty").await; assert_eq!(code, 0); let (code, stdout) = run_command(&mut state, "echo \"[$GPG_TTY]\"").await; assert_eq!(code, 0); assert!( !stdout.contains("/grok-sentinel-tty"), "GPG_TTY must not persist across commands via the snapshot, got: {stdout:?}" ); } /// Unmatched globs must not abort with NOMATCH after login snapshot restore. #[tokio::test] async fn test_zsh_unmatched_glob_does_not_abort() { if !zsh_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Zsh, &cwd).await.unwrap(); let prep = state .prepare_command( "true", None, crate::computer::local::SearchShadowConfig::default(), ) .unwrap(); let wrapper = prep .args .iter() .find(|a| a.contains("builtin eval")) .expect("zsh wrapper should be in prepare_command args"); let snap_pos = wrapper .find("eval \"$snap\"") .expect("wrapper must restore snapshot via eval \"$snap\""); let nonomatch_pos = wrapper .find("setopt nonomatch") .expect("wrapper must force nonomatch after snapshot restore"); assert!( nonomatch_pos > snap_pos, "nonomatch must be forced after snapshot eval, wrapper={wrapper:?}" ); let pattern = "--include=*.no_such_ext_xyz_12345"; let (code, stdout) = run_command(&mut state, &format!("printf '%s\\n' {pattern}")).await; assert_eq!(code, 0, "unmatched glob must not fail, stdout={stdout:?}"); assert!( stdout.contains(pattern), "expected literal unmatched glob in output, got: {stdout:?}" ); assert!( !stdout.contains("no matches found"), "NOMATCH error must not appear, got: {stdout:?}" ); } #[tokio::test] async fn test_function_persists_across_commands() { if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); // Define a function let (code, _) = run_command(&mut state, "greet() { echo \"hello $1\"; }").await; assert_eq!(code, 0); // Call it in the next command let (code, stdout) = run_command(&mut state, "greet world").await; assert_eq!(code, 0); assert_eq!(stdout.trim(), "hello world"); } #[tokio::test] async fn test_alias_persists_across_commands() { if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); // Define an alias let (code, _) = run_command(&mut state, "alias ll='ls -la'").await; assert_eq!(code, 0); // Verify the alias survives by checking the dump itself (base64-encoded). // We can't check plaintext in the snapshot since it's base64, but we can // verify the snapshot is valid and non-empty (alias was captured in the dump). assert!( !state.snapshot.is_empty(), "snapshot should be non-empty after alias" ); } #[tokio::test] async fn test_embedded_search_shadows() { use super::super::embedded_search_tools::search_injection; let shadows = crate::computer::local::SearchShadowConfig::default(); if !bash_available() { return; } let inject = search_injection(shadows); if inject.is_empty() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); let prep = state.prepare_command("true", None, shadows).unwrap(); // Shadows enabled → the self-resolving find/grep functions are always // installed (they fall back to the OS binary if bfs/ugrep aren't found). assert!( prep.args .iter() .any(|a| a.contains("find()") && a.contains("grep()")), "prepare_command should install find/grep shadows; args={:?}", prep.args ); let (code, stdout) = run_command(&mut state, "type -t find && type -t grep").await; assert_eq!(code, 0, "type -t find/grep failed: {stdout}"); assert_eq!( stdout.lines().filter(|l| l.trim() == "function").count(), 2, "expected find and grep to both be shell functions, got: {stdout:?}" ); // Only assert the binary actually routes to ugrep when ugrep resolves on // this host; otherwise the shadow correctly falls back to OS grep. if which::which("ugrep").is_ok() { let (code, stdout) = run_command(&mut state, "grep --version | head -3").await; assert_eq!(code, 0, "grep --version failed: {stdout}"); assert!( stdout.to_ascii_lowercase().contains("ugrep"), "expected ugrep in --version output when ugrep resolves: {stdout:?}" ); } } #[tokio::test] async fn test_bad_command_preserves_state() { if !bash_available() { return; } let cwd = std::env::current_dir().unwrap(); let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap(); // Set up some state let (_, _) = run_command(&mut state, "export SURVIVE_TEST=yes").await; let prev_cwd = state.cwd.clone(); // Run a failing command — state should still update (dump runs regardless) let (code, _) = run_command(&mut state, "false").await; assert_ne!(code, 0); // Previous state should still be there assert_eq!(state.cwd, prev_cwd); let (_, stdout) = run_command(&mut state, "echo $SURVIVE_TEST").await; assert_eq!(stdout.trim(), "yes"); } }