use std::collections::BTreeSet;
use std::io::{self, Read, Write};
use std::process::{Child, Command, ExitStatus, Stdio};
use std::time::{Duration, Instant};
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(300);
pub const DEFAULT_MAX_OUTPUT: usize = 64 * 1024 * 1024;
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)]
pub struct SpawnOutput {
pub status: ExitStatus,
pub stdout: Vec<u8>,
pub timed_out: bool,
pub stdout_truncated: bool,
}
impl SpawnOutput {
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() {
return Some(format!("{what} exited with {}", self.status));
}
None
}
}
pub fn run_argv(argv: &[String], stdin: &str, timeout: Option<Duration>) -> io::Result<SpawnOutput> {
let mut cmd = Command::new(&argv[0]);
cmd.args(&argv[1..]);
for var in secret_env_vars() {
cmd.env_remove(var);
}
cmd.stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::inherit());
let mut child = cmd.spawn()?;
let payload = stdin.as_bytes().to_vec();
let stdin_thread = child.stdin.take().map(|mut pipe| {
std::thread::spawn(move || {
let _ = pipe.write_all(&payload);
})
});
let out_thread = child.stdout.take().map(|pipe| std::thread::spawn(move || drain(pipe)));
let (status, timed_out) = wait_bounded(&mut child, 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));
Ok(SpawnOutput { status, stdout, timed_out, stdout_truncated })
}
fn drain<R: Read>(mut src: R) -> (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() < DEFAULT_MAX_OUTPUT {
let take = (DEFAULT_MAX_OUTPUT - kept.len()).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));
}
}