From 0b4e4fcc4aa87eb2035da9b5fc61bb86da2b26a5 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 2 Aug 2026 10:10:17 +0200 Subject: [PATCH] Load shell environment once at startup --- src/agent.rs | 284 ++++++++++++++++++++++++++++++++++++++++++--------- src/main.rs | 1 + 2 files changed, 236 insertions(+), 49 deletions(-) diff --git a/src/agent.rs b/src/agent.rs index c9d9004..452334b 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -13,44 +13,168 @@ use crate::settings::{EngineSettings, ReasoningMode, TurnSettings}; use serde::Deserialize; use serde_json::{Map, Value}; use std::collections::HashMap; -use std::ffi::OsString; +use std::ffi::{OsStr, OsString}; use std::fs::{self, File}; +use std::io::Read; +use std::os::unix::ffi::OsStringExt; use std::os::unix::process::CommandExt; use std::path::{Path, PathBuf}; use std::process::{Child, Command, Stdio}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender, TryRecvError}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; const MAX_FILE_BYTES: u64 = 16 * 1024 * 1024; -const SHELL_ENV_ALLOWLIST: &[&str] = &[ - "PATH", - "HOME", - "USER", - "LOGNAME", - "SHELL", - "TMPDIR", - "LANG", - "LC_ALL", - "TERM", - "DEVELOPER_DIR", - "SDKROOT", - "MACOSX_DEPLOYMENT_TARGET", - "RUSTUP_TOOLCHAIN", -]; +const SHELL_ENV_TIMEOUT: Duration = Duration::from_secs(5); +const SHELL_ENV_SENTINEL: &[u8] = b"\0DS4_ENV\0"; +static USER_SHELL_ENVIRONMENT: OnceLock> = OnceLock::new(); fn user_shell() -> OsString { std::env::var_os("SHELL").unwrap_or_else(|| OsString::from("/bin/sh")) } -fn shell_process(shell: &std::ffi::OsStr, command: &str) -> Command { +enum ShellEnvironmentProbe { + Loaded(Vec<(OsString, OsString)>), + Timeout, + Unavailable, +} + +pub(crate) fn initialize_shell_environment() { + let shell = user_shell(); + let environment = match load_shell_environment( + &shell, + std::env::var_os("HOME").as_deref(), + std::env::var_os("ZDOTDIR").as_deref(), + ) { + Some(environment) => environment, + None => { + eprintln!( + "DS4Server: could not load the environment from {}; using the app environment", + shell.to_string_lossy() + ); + std::env::vars_os().collect() + } + }; + let _ = USER_SHELL_ENVIRONMENT.set(environment); +} + +fn load_shell_environment( + shell: &OsStr, + home: Option<&OsStr>, + zdotdir: Option<&OsStr>, +) -> Option> { + match probe_shell_environment(shell, "-il", home, zdotdir, SHELL_ENV_TIMEOUT) { + ShellEnvironmentProbe::Loaded(environment) => Some(environment), + ShellEnvironmentProbe::Timeout => None, + ShellEnvironmentProbe::Unavailable => { + match probe_shell_environment(shell, "-l", home, zdotdir, SHELL_ENV_TIMEOUT) { + ShellEnvironmentProbe::Loaded(environment) => Some(environment), + ShellEnvironmentProbe::Timeout | ShellEnvironmentProbe::Unavailable => None, + } + } + } +} + +fn probe_shell_environment( + shell: &OsStr, + mode: &str, + home: Option<&OsStr>, + zdotdir: Option<&OsStr>, + timeout: Duration, +) -> ShellEnvironmentProbe { + let mut process = Command::new(shell); + process + .arg(mode) + .arg("-c") + .arg("printf '\\0DS4_ENV\\0'; /usr/bin/env -0") + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .process_group(0); + if let Some(home) = home { + process.env("HOME", home); + } + if let Some(zdotdir) = zdotdir { + process.env("ZDOTDIR", zdotdir); + } + let Ok(mut child) = process.spawn() else { + return ShellEnvironmentProbe::Unavailable; + }; + let Some(mut stdout) = child.stdout.take() else { + stop_process_group(&mut child); + return ShellEnvironmentProbe::Unavailable; + }; + let (output_sender, output_receiver) = mpsc::channel(); + thread::spawn(move || { + let mut output = Vec::new(); + let _ = stdout.read_to_end(&mut output); + let _ = output_sender.send(output); + }); + let deadline = Instant::now() + timeout; + loop { + match child.try_wait() { + Ok(Some(status)) => { + let Ok(output) = output_receiver + .recv_timeout(deadline.saturating_duration_since(Instant::now())) + else { + signal_process_group(child.id(), "-KILL"); + return ShellEnvironmentProbe::Timeout; + }; + return if status.success() { + parse_shell_environment(output).map_or( + ShellEnvironmentProbe::Unavailable, + ShellEnvironmentProbe::Loaded, + ) + } else { + ShellEnvironmentProbe::Unavailable + }; + } + Err(_) => { + stop_process_group(&mut child); + return ShellEnvironmentProbe::Unavailable; + } + Ok(None) if Instant::now() < deadline => thread::sleep(Duration::from_millis(10)), + Ok(None) => { + stop_process_group(&mut child); + return ShellEnvironmentProbe::Timeout; + } + } + } +} + +fn parse_shell_environment(output: Vec) -> Option> { + let start = output + .windows(SHELL_ENV_SENTINEL.len()) + .rposition(|window| window == SHELL_ENV_SENTINEL)? + + SHELL_ENV_SENTINEL.len(); + let environment = output[start..] + .split(|byte| *byte == 0) + .filter_map(|entry| { + let equals = entry.iter().position(|byte| *byte == b'=')?; + Some(( + OsString::from_vec(entry[..equals].to_vec()), + OsString::from_vec(entry[equals + 1..].to_vec()), + )) + }) + .collect::>(); + (!environment.is_empty()).then_some(environment) +} + +fn shell_environment() -> Vec<(OsString, OsString)> { + USER_SHELL_ENVIRONMENT + .get() + .cloned() + .unwrap_or_else(|| std::env::vars_os().collect()) +} + +fn shell_process(shell: &OsStr, command: &str) -> Command { let mut process = Command::new("/bin/sh"); process .arg("-c") .arg(format!( - "ulimit -f {}; exec \"$1\" -l -i -c \"$2\"", + "ulimit -f {}; exec \"$1\" -c \"$2\"", MAX_FILE_BYTES / 512 )) .arg("ds4-agent") @@ -829,12 +953,8 @@ impl Tools { .stdout(stdout) .stderr(stderr) .process_group(0) - .env_clear(); - for name in SHELL_ENV_ALLOWLIST { - if let Some(value) = std::env::var_os(name) { - process.env(name, value); - } - } + .env_clear() + .envs(shell_environment()); process.env("PWD", &self.root); let child = process .spawn() @@ -1743,30 +1863,33 @@ fn unique_match(data: &str, needle: &str, label: &str) -> Result } fn stop_job(job: &mut BashJob) { - if job.child.try_wait().ok().flatten().is_some() { + stop_process_group(&mut job.child); +} + +fn stop_process_group(child: &mut Child) { + if child.try_wait().ok().flatten().is_some() { return; } - let pid = job.child.id(); - let _ = Command::new("/bin/kill") - .arg("-TERM") - .arg(format!("-{pid}")) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status(); + let pid = child.id(); + signal_process_group(pid, "-TERM"); let deadline = Instant::now() + Duration::from_secs(1); while Instant::now() < deadline { - if job.child.try_wait().ok().flatten().is_some() { + if child.try_wait().ok().flatten().is_some() { return; } thread::sleep(Duration::from_millis(20)); } + signal_process_group(pid, "-KILL"); + let _ = child.wait(); +} + +fn signal_process_group(pid: u32, signal: &str) { let _ = Command::new("/bin/kill") - .arg("-KILL") + .arg(signal) .arg(format!("-{pid}")) .stdout(Stdio::null()) .stderr(Stdio::null()) .status(); - let _ = job.child.wait(); } fn string<'a>(call: &'a ToolCall, name: &str) -> Option<&'a str> { @@ -2007,14 +2130,11 @@ mod tests { ] { assert!(risky_shell_reason(command, root).is_some(), "{command}"); } - assert!(!SHELL_ENV_ALLOWLIST.contains(&"GITHUB_TOKEN")); - assert!(!SHELL_ENV_ALLOWLIST.contains(&"AWS_SECRET_ACCESS_KEY")); - assert!(!SHELL_ENV_ALLOWLIST.contains(&"SSH_AUTH_SOCK")); } #[cfg(target_os = "macos")] #[test] - fn shell_commands_load_the_user_login_and_interactive_environment() { + fn shell_environment_is_loaded_once_from_login_and_interactive_startup_files() { let directory = std::env::temp_dir().join(format!( "ds4-agent-shell-{}", SystemTime::now() @@ -2025,18 +2145,88 @@ mod tests { fs::create_dir_all(&directory).unwrap(); fs::write( directory.join(".zprofile"), - "export PATH=\"$HOME/homebrew/bin:$PATH\"\nexport DS4_LOGIN_PROFILE=loaded\n", + "export PATH=\"$HOME/login-bin:$PATH\"\nexport DS4_LOGIN=loaded\n", ) .unwrap(); fs::write( directory.join(".zshrc"), - "export DS4_INTERACTIVE_PROFILE=loaded\n", + "export PATH=\"$HOME/interactive-bin:$PATH\"\nexport DS4_INTERACTIVE=loaded\n", ) .unwrap(); + let environment = load_shell_environment( + OsStr::new("/bin/zsh"), + Some(directory.as_os_str()), + Some(directory.as_os_str()), + ) + .unwrap() + .into_iter() + .collect::>(); + assert_eq!( + environment.get(OsStr::new("DS4_LOGIN")), + Some(&OsString::from("loaded")) + ); + assert_eq!( + environment.get(OsStr::new("DS4_INTERACTIVE")), + Some(&OsString::from("loaded")) + ); + let path = environment + .get(OsStr::new("PATH")) + .unwrap() + .to_string_lossy(); + assert!(path.starts_with(&format!( + "{}/interactive-bin:{}/login-bin:", + directory.display(), + directory.display() + ))); + fs::remove_dir_all(directory).unwrap(); + } + + #[cfg(target_os = "macos")] + #[test] + fn shell_environment_probe_stops_blocked_startup_files() { + let directory = std::env::temp_dir().join(format!( + "ds4-agent-shell-timeout-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + fs::create_dir_all(&directory).unwrap(); + fs::write(directory.join(".zshrc"), "sleep 10\n").unwrap(); + + let started = Instant::now(); + assert!(matches!( + probe_shell_environment( + OsStr::new("/bin/zsh"), + "-il", + Some(directory.as_os_str()), + Some(directory.as_os_str()), + Duration::from_millis(50), + ), + ShellEnvironmentProbe::Timeout + )); + assert!(started.elapsed() < Duration::from_secs(2)); + fs::remove_dir_all(directory).unwrap(); + } + + #[cfg(target_os = "macos")] + #[test] + fn shell_commands_skip_user_startup_files() { + let directory = std::env::temp_dir().join(format!( + "ds4-agent-shell-startup-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + fs::create_dir_all(&directory).unwrap(); + fs::write(directory.join(".zprofile"), "export DS4_PROFILE=login\n").unwrap(); + fs::write(directory.join(".zshrc"), "export DS4_RC=interactive\n").unwrap(); + let mut process = shell_process( - std::ffi::OsStr::new("/bin/zsh"), - "printf '%s|%s|%s' \"$DS4_LOGIN_PROFILE\" \"$DS4_INTERACTIVE_PROFILE\" \"$PATH\"", + OsStr::new("/bin/zsh"), + "printf '%s|%s' \"$DS4_PROFILE\" \"$DS4_RC\"", ); let output = process .env_clear() @@ -2048,11 +2238,7 @@ mod tests { .unwrap(); assert!(output.status.success()); let stdout = String::from_utf8(output.stdout).unwrap(); - assert!(stdout.starts_with("loaded|loaded|"), "{stdout}"); - assert!( - stdout.contains(&format!("{}/homebrew/bin", directory.display())), - "{stdout}" - ); + assert_eq!(stdout, "|"); fs::remove_dir_all(directory).unwrap(); } diff --git a/src/main.rs b/src/main.rs index 32cf0f1..8f677f9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,6 +36,7 @@ fn main() -> iced::Result { } return Ok(()); } + agent::initialize_shell_environment(); #[cfg(target_os = "macos")] if let Err(error) = engine::configure_metal_sources() { eprintln!("DS4Server: {error}");