use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
use anyhow::{Context, Result};
use wait_timeout::ChildExt;
pub fn env_allowed(key: &str) -> bool {
if key.starts_with("SELFWARE_") || key.starts_with("LC_") {
return true;
}
matches!(
key,
"PATH"
| "HOME"
| "USER"
| "LOGNAME"
| "SHELL"
| "LANG"
| "TERM"
| "TMPDIR"
| "TZ"
| "CARGO_HOME"
| "RUSTUP_HOME"
| "RUST_BACKTRACE"
| "GOPATH"
| "GOROOT"
| "PYENV_ROOT"
| "NODE_PATH"
| "JAVA_HOME"
| "OPENROUTER_API_KEY"
| "LLM_API_KEY"
| "LLAMA_SERVER_BIN"
| "SWEBENCH_MODELS_DIR"
)
}
pub fn apply_env_allowlist(cmd: &mut Command) {
cmd.env_clear();
for (k, v) in std::env::vars() {
if env_allowed(&k) {
cmd.env(k, v);
}
}
}
#[derive(Debug, Clone)]
pub struct ProcOutcome {
pub stdout: String,
pub exit_code: i32,
pub timed_out: bool,
pub wall_secs: f64,
}
pub fn run_with_group_timeout(mut cmd: Command, timeout: Duration) -> Result<ProcOutcome> {
cmd.stdout(Stdio::piped());
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
cmd.process_group(0);
}
let started = Instant::now();
let mut child = cmd.spawn().context("spawning subprocess")?;
let stdout = child.stdout.take().expect("stdout piped above");
let stdout_handle = std::thread::spawn(move || {
let mut buf = String::new();
let _ = std::io::Read::read_to_string(&mut std::io::BufReader::new(stdout), &mut buf);
buf
});
let deadline = started + timeout;
let poll_interval = Duration::from_millis(500);
loop {
match child.wait_timeout(poll_interval)? {
Some(status) => {
let stdout = stdout_handle.join().unwrap_or_default();
return Ok(ProcOutcome {
stdout,
exit_code: status.code().unwrap_or(-1),
timed_out: false,
wall_secs: started.elapsed().as_secs_f64(),
});
}
None => {
if Instant::now() >= deadline {
kill_group(&mut child);
let _ = child.wait();
let stdout = stdout_handle.join().unwrap_or_default();
return Ok(ProcOutcome {
stdout,
exit_code: -1,
timed_out: true,
wall_secs: started.elapsed().as_secs_f64(),
});
}
}
}
}
}
fn kill_group(child: &mut std::process::Child) {
#[cfg(unix)]
{
use nix::sys::signal::{killpg, Signal};
use nix::unistd::Pid;
let pgid = Pid::from_raw(child.id() as i32);
let _ = killpg(pgid, Signal::SIGTERM);
std::thread::sleep(Duration::from_millis(500));
let _ = killpg(pgid, Signal::SIGKILL);
}
#[cfg(not(unix))]
{
let _ = child.kill();
}
}
#[cfg(test)]
#[path = "../../tests/unit/bench_harness/subprocess/subprocess_test.rs"]
mod tests;