hugo/DS4Server
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects Releases Wiki Activity

Load shell environment once at startup

Browse Source
This commit is contained in:
Georg Bauer
2026-08-02 10:10:17 +02:00
parent 118a9d41e0
commit 0b4e4fcc4a
2 changed files with 236 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

284
src/agent.rs
View File

@@ -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<Vec<(OsString, OsString)>> = 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<Vec<(OsString, OsString)>> {
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<u8>) -> Option<Vec<(OsString, OsString)>> {
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::<Vec<_>>();
(!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<usize, String>
}
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::<HashMap<_, _>>();
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();
}

1
src/main.rs
View File

@@ -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}");