use std::collections::BTreeSet;
use std::io::{self, Read, Write};
use std::path::PathBuf;
use std::process::{Child, Command, ExitStatus, Stdio};
use std::time::{Duration, Instant};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EnvPolicy {
InheritExcept { deny: Vec<String> },
ClearExcept { allow: Vec<String> },
}
impl EnvPolicy {
pub fn minimal_allow() -> Vec<String> {
let base: &[&str] = if cfg!(windows) {
&["PATH", "PATHEXT", "SYSTEMROOT", "SYSTEMDRIVE", "COMSPEC", "TEMP", "TMP", "USERPROFILE"]
} else {
&["PATH", "HOME", "TMPDIR", "LANG", "LC_ALL", "TZ"]
};
base.iter().map(|s| s.to_string()).collect()
}
}
impl Default for EnvPolicy {
fn default() -> Self {
EnvPolicy::InheritExcept { deny: secret_env_vars() }
}
}
static SECRET_ENV: std::sync::Mutex<Option<BTreeSet<String>>> = std::sync::Mutex::new(None);
pub fn deny_env_var(name: &str) {
if name.trim().is_empty() {
return;
}
let mut guard = SECRET_ENV.lock().unwrap_or_else(|e| e.into_inner());
guard.get_or_insert_with(BTreeSet::new).insert(name.to_string());
}
pub fn secret_env_vars() -> Vec<String> {
let guard = SECRET_ENV.lock().unwrap_or_else(|e| e.into_inner());
guard.as_ref().map(|s| s.iter().cloned().collect()).unwrap_or_default()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum StderrMode {
#[default]
Pipe,
Inherit,
}
#[derive(Debug, Clone)]
pub struct SpawnPolicy {
pub timeout: Option<Duration>,
pub max_output_bytes: usize,
pub env: EnvPolicy,
pub stderr: StderrMode,
pub current_dir: Option<PathBuf>,
}
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(300);
pub const DEFAULT_MAX_OUTPUT: usize = 64 * 1024 * 1024;
impl Default for SpawnPolicy {
fn default() -> Self {
SpawnPolicy {
timeout: Some(DEFAULT_TIMEOUT),
max_output_bytes: DEFAULT_MAX_OUTPUT,
env: EnvPolicy::default(),
stderr: StderrMode::default(),
current_dir: None,
}
}
}
impl SpawnPolicy {
pub fn deny_vars<I: IntoIterator<Item = String>>(mut self, vars: I) -> Self {
let mut denied: BTreeSet<String> = match self.env {
EnvPolicy::InheritExcept { deny } => deny.into_iter().collect(),
EnvPolicy::ClearExcept { allow } => {
self.env = EnvPolicy::ClearExcept { allow };
return self;
}
};
denied.extend(vars);
self.env = EnvPolicy::InheritExcept { deny: denied.into_iter().collect() };
self
}
pub fn timeout(mut self, timeout: Option<Duration>) -> Self {
self.timeout = timeout;
self
}
pub fn stderr(mut self, mode: StderrMode) -> Self {
self.stderr = mode;
self
}
}
#[derive(Debug)]
pub struct SpawnOutput {
pub status: ExitStatus,
pub stdout: Vec<u8>,
pub stderr: Vec<u8>,
pub timed_out: bool,
pub stdout_truncated: bool,
pub stderr_truncated: bool,
}
impl SpawnOutput {
pub fn stderr_text(&self) -> String {
String::from_utf8_lossy(&self.stderr).trim().to_string()
}
pub fn failure(&self, what: &str) -> Option<String> {
if self.timed_out {
return Some(format!("{what} timed out and was killed"));
}
if !self.status.success() {
let err = self.stderr_text();
return Some(if err.is_empty() {
format!("{what} exited with {}", self.status)
} else {
format!("{what} exited with {}: {err}", self.status)
});
}
None
}
}
fn spawn_past_busy_exec(cmd: &mut Command) -> io::Result<Child> {
let mut attempt = 0u32;
loop {
match cmd.spawn() {
Err(e) if e.kind() == io::ErrorKind::ExecutableFileBusy && attempt < 20 => {
attempt += 1;
std::thread::sleep(Duration::from_millis(5 * u64::from(attempt)));
}
other => return other,
}
}
}
pub fn run(
mut cmd: Command,
stdin: Option<&[u8]>,
extra_env: &[(&str, &str)],
policy: &SpawnPolicy,
) -> io::Result<SpawnOutput> {
match &policy.env {
EnvPolicy::InheritExcept { deny } => {
for var in deny {
cmd.env_remove(var);
}
}
EnvPolicy::ClearExcept { allow } => {
cmd.env_clear();
for var in allow {
if let Ok(val) = std::env::var(var) {
cmd.env(var, val);
}
}
}
}
for (k, v) in extra_env {
cmd.env(k, v);
}
if let Some(dir) = &policy.current_dir {
cmd.current_dir(dir);
}
cmd.stdin(if stdin.is_some() { Stdio::piped() } else { Stdio::null() })
.stdout(Stdio::piped())
.stderr(match policy.stderr {
StderrMode::Pipe => Stdio::piped(),
StderrMode::Inherit => Stdio::inherit(),
});
let mut child = spawn_past_busy_exec(&mut cmd)?;
let stdin_thread = child.stdin.take().map(|mut pipe| {
let payload = stdin.unwrap_or_default().to_vec();
std::thread::spawn(move || {
let _ = pipe.write_all(&payload);
})
});
let cap = policy.max_output_bytes;
let out_thread = child.stdout.take().map(|pipe| std::thread::spawn(move || drain(pipe, cap)));
let err_thread = child.stderr.take().map(|pipe| std::thread::spawn(move || drain(pipe, cap)));
let (status, timed_out) = wait_bounded(&mut child, policy.timeout)?;
if let Some(t) = stdin_thread {
let _ = t.join();
}
let (stdout, stdout_truncated) = out_thread.and_then(|t| t.join().ok()).unwrap_or((Vec::new(), false));
let (stderr, stderr_truncated) = err_thread.and_then(|t| t.join().ok()).unwrap_or((Vec::new(), false));
Ok(SpawnOutput { status, stdout, stderr, timed_out, stdout_truncated, stderr_truncated })
}
fn drain<R: Read>(mut src: R, cap: usize) -> (Vec<u8>, bool) {
let mut kept = Vec::new();
let mut buf = [0u8; 16 * 1024];
let mut truncated = false;
loop {
match src.read(&mut buf) {
Ok(0) => break,
Ok(n) => {
if kept.len() < cap {
let room = cap - kept.len();
let take = room.min(n);
kept.extend_from_slice(&buf[..take]);
if take < n {
truncated = true;
}
} else {
truncated = true;
}
}
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(_) => break,
}
}
(kept, truncated)
}
fn wait_bounded(child: &mut Child, timeout: Option<Duration>) -> io::Result<(ExitStatus, bool)> {
let Some(limit) = timeout else {
return Ok((child.wait()?, false));
};
let deadline = Instant::now() + limit;
let mut nap = Duration::from_millis(1);
loop {
if let Some(status) = child.try_wait()? {
return Ok((status, false));
}
if Instant::now() >= deadline {
let _ = child.kill();
let status = child.wait()?;
return Ok((status, true));
}
std::thread::sleep(nap);
nap = (nap * 2).min(Duration::from_millis(50));
}
}
#[cfg(test)]
mod tests {
use super::*;
fn sh(script: &str) -> Command {
let mut c = Command::new("/bin/sh");
c.arg("-c").arg(script);
c
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn captures_stdout_and_exit_status() {
let out = run(sh("printf hello"), None, &[], &SpawnPolicy::default()).unwrap();
assert!(out.status.success());
assert_eq!(out.stdout, b"hello");
assert!(!out.timed_out);
assert!(!out.stdout_truncated);
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn stdin_reaches_the_child() {
let out = run(sh("cat"), Some(b"payload"), &[], &SpawnPolicy::default()).unwrap();
assert_eq!(out.stdout, b"payload");
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn timeout_kills_a_hung_child() {
let policy = SpawnPolicy::default().timeout(Some(Duration::from_millis(150)));
let out = run(sh("sleep 30"), None, &[], &policy).unwrap();
assert!(out.timed_out, "expected the child to be killed");
assert!(!out.status.success());
assert!(out.failure("tool").unwrap().contains("timed out"));
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn output_cap_truncates_without_hanging() {
let policy = SpawnPolicy { max_output_bytes: 1024, ..SpawnPolicy::default() };
let out = run(sh("head -c 200000 /dev/zero"), None, &[], &policy).unwrap();
assert_eq!(out.stdout.len(), 1024);
assert!(out.stdout_truncated);
assert!(!out.timed_out, "draining past the cap must not stall the child");
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn large_stdin_does_not_deadlock() {
let big = vec![b'x'; 4 * 1024 * 1024];
let policy = SpawnPolicy::default().timeout(Some(Duration::from_secs(20)));
let out = run(sh("cat"), Some(&big), &[], &policy).unwrap();
assert!(!out.timed_out, "write-then-read deadlock");
assert_eq!(out.stdout.len(), big.len());
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn denied_vars_do_not_reach_the_child() {
std::env::set_var("AREEV_TEST_SECRET", "hunter2");
let policy = SpawnPolicy::default().deny_vars(["AREEV_TEST_SECRET".to_string()]);
let out = run(sh("printf %s \"${AREEV_TEST_SECRET:-absent}\""), None, &[], &policy).unwrap();
std::env::remove_var("AREEV_TEST_SECRET");
assert_eq!(String::from_utf8_lossy(&out.stdout), "absent");
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn inherited_vars_still_reach_the_child() {
std::env::set_var("AREEV_TEST_KEEP", "kept");
std::env::set_var("AREEV_TEST_DROP", "dropped");
let policy = SpawnPolicy::default().deny_vars(["AREEV_TEST_DROP".to_string()]);
let out = run(sh("printf %s \"${AREEV_TEST_KEEP:-absent}\""), None, &[], &policy).unwrap();
std::env::remove_var("AREEV_TEST_KEEP");
std::env::remove_var("AREEV_TEST_DROP");
assert_eq!(String::from_utf8_lossy(&out.stdout), "kept");
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn clear_except_drops_everything_unlisted_but_keeps_extras() {
std::env::set_var("AREEV_TEST_AMBIENT", "ambient");
let policy = SpawnPolicy {
env: EnvPolicy::ClearExcept { allow: EnvPolicy::minimal_allow() },
..SpawnPolicy::default()
};
let out = run(
sh("printf %s \"${AREEV_TEST_AMBIENT:-absent}/${AREEV_EXTRA:-none}\""),
None,
&[("AREEV_EXTRA", "set")],
&policy,
)
.unwrap();
std::env::remove_var("AREEV_TEST_AMBIENT");
assert_eq!(String::from_utf8_lossy(&out.stdout), "absent/set");
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn nonzero_exit_reports_stderr() {
let out = run(sh("echo boom >&2; exit 3"), None, &[], &SpawnPolicy::default()).unwrap();
let msg = out.failure("embed command").unwrap();
assert!(msg.contains("embed command"), "{msg}");
assert!(msg.contains("boom"), "{msg}");
}
#[test]
#[cfg_attr(windows, ignore = "uses /bin/sh")]
fn registered_secrets_are_scrubbed_without_the_seam_asking() {
std::env::set_var("AREEV_TEST_REGISTERED", "hunter2");
deny_env_var("AREEV_TEST_REGISTERED");
let out = run(
sh("printf %s \"${AREEV_TEST_REGISTERED:-absent}\""),
None,
&[],
&SpawnPolicy::default(),
)
.unwrap();
std::env::remove_var("AREEV_TEST_REGISTERED");
assert_eq!(String::from_utf8_lossy(&out.stdout), "absent");
}
#[test]
fn deny_env_var_ignores_blanks() {
deny_env_var(" ");
deny_env_var("");
assert!(!secret_env_vars().iter().any(|v| v.trim().is_empty()));
}
#[test]
#[cfg(target_os = "linux")]
fn an_executable_still_open_for_writing_is_retried_not_failed() {
use std::os::unix::fs::PermissionsExt;
let dir = std::env::temp_dir().join(format!("areev-etxtbsy-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("busy.sh");
let mut held = std::fs::File::create(&path).unwrap();
held.write_all(b"#!/bin/sh\nprintf ok\n").unwrap();
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o700)).unwrap();
let bare = Command::new(&path).spawn().map(|mut c| c.wait());
assert_eq!(bare.err().map(|e| e.kind()), Some(io::ErrorKind::ExecutableFileBusy));
let release = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(40));
drop(held);
});
let out = run(Command::new(&path), None, &[], &SpawnPolicy::default()).unwrap();
release.join().unwrap();
std::fs::remove_dir_all(&dir).ok();
assert_eq!(out.stdout, b"ok");
}
#[test]
fn spawn_failure_is_an_error_not_a_panic() {
let cmd = Command::new("areev-no-such-binary-eaf1");
assert!(run(cmd, None, &[], &SpawnPolicy::default()).is_err());
}
}