use std::process::Output;
use std::time::Duration;
use tokio::process::Command as TokioCommand;
#[cfg(unix)]
pub fn sigkill_process_group(pid: u32) {
unsafe {
libc::kill(-(pid as i32), libc::SIGKILL);
}
}
pub fn harden(cmd: &mut TokioCommand) {
cmd.stdin(std::process::Stdio::null()).kill_on_drop(true);
#[cfg(unix)]
cmd.process_group(0);
}
pub async fn run_capture_with_timeout(
mut cmd: TokioCommand,
timeout: Duration,
label: &str,
on_timeout: impl FnOnce() -> anyhow::Error,
) -> anyhow::Result<Output> {
harden(&mut cmd);
cmd.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
let child = cmd
.spawn()
.map_err(|e| anyhow::anyhow!("failed to spawn {label}: {e}"))?;
#[cfg(unix)]
let pid = child.id();
match tokio::time::timeout(timeout, child.wait_with_output()).await {
Ok(result) => result.map_err(|e| anyhow::anyhow!("failed to run {label}: {e}")),
Err(_) => {
#[cfg(unix)]
if let Some(pid) = pid {
sigkill_process_group(pid);
}
Err(on_timeout())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn run_capture_with_timeout_returns_output_on_success() {
let mut cmd = TokioCommand::new("echo");
cmd.arg("hello");
let output = run_capture_with_timeout(cmd, Duration::from_secs(10), "echo", || {
anyhow::anyhow!("unexpected timeout")
})
.await
.expect("echo should succeed");
assert!(output.status.success());
assert_eq!(String::from_utf8_lossy(&output.stdout).trim(), "hello");
}
#[tokio::test]
async fn run_capture_with_timeout_errors_via_on_timeout() {
let mut cmd = TokioCommand::new("sleep");
cmd.arg("30");
let err = run_capture_with_timeout(cmd, Duration::from_secs(1), "sleep", || {
anyhow::anyhow!("slept too long")
})
.await
.expect_err("sleep 30 with a 1s budget must time out");
assert!(err.to_string().contains("slept too long"));
}
#[tokio::test]
async fn run_capture_with_timeout_detaches_stdin() {
let cmd = TokioCommand::new("cat");
let output = run_capture_with_timeout(cmd, Duration::from_secs(5), "cat", || {
anyhow::anyhow!("cat hung")
})
.await
.expect("cat with null stdin exits immediately");
assert!(output.status.success());
}
#[cfg(unix)]
#[tokio::test]
async fn run_capture_with_timeout_kills_process_group() {
use std::io::Read;
let tmp = tempfile::tempdir().expect("tempdir");
let pidfile = tmp.path().join("grandchild.pid");
let script = format!("sleep 30 & echo $! > {} ; wait", pidfile.display());
let mut cmd = TokioCommand::new("sh");
cmd.arg("-c").arg(&script);
let err = run_capture_with_timeout(cmd, Duration::from_secs(1), "sh-tree", || {
anyhow::anyhow!("sh-tree timed out")
})
.await
.expect_err("sh tree with a 1s budget must time out");
assert!(err.to_string().contains("sh-tree timed out"));
let mut s = String::new();
std::fs::File::open(&pidfile)
.expect("sh should have recorded the grandchild pid")
.read_to_string(&mut s)
.expect("read pidfile");
let grandchild: i32 = s.trim().parse().expect("grandchild pid");
assert_grandchild_reaped(grandchild).await;
}
#[cfg(unix)]
#[test]
fn sigkill_process_group_reaps_grandchild() {
use std::io::Read;
use std::os::unix::process::CommandExt;
use std::thread::sleep;
let tmp = tempfile::tempdir().unwrap();
let pidfile = tmp.path().join("grandchild.pid");
let script = format!("sleep 30 & echo $! > {} ; wait", pidfile.display());
let mut cmd = std::process::Command::new("sh");
cmd.arg("-c").arg(&script);
cmd.process_group(0);
let mut child = cmd.spawn().expect("spawn sh tree");
let leader = child.id();
let mut grandchild = None;
for _ in 0..100 {
if let Ok(mut f) = std::fs::File::open(&pidfile) {
let mut s = String::new();
f.read_to_string(&mut s).ok();
if let Ok(pid) = s.trim().parse::<i32>() {
grandchild = Some(pid);
break;
}
}
sleep(Duration::from_millis(20));
}
let grandchild = grandchild.expect("sh should record the grandchild pid");
sigkill_process_group(leader);
let _ = child.wait();
let mut gone = false;
for _ in 0..100 {
if unsafe { libc::kill(grandchild, 0) } == -1 {
let errno = std::io::Error::last_os_error().raw_os_error();
if errno == Some(libc::ESRCH) || errno == Some(libc::EPERM) {
gone = true;
break;
}
}
sleep(Duration::from_millis(20));
}
assert!(gone, "grandchild {grandchild} survived the group kill");
}
#[cfg(unix)]
async fn assert_grandchild_reaped(grandchild: i32) {
let mut gone = false;
for _ in 0..60 {
if unsafe { libc::kill(grandchild, 0) } == -1 {
let errno = std::io::Error::last_os_error().raw_os_error();
if errno == Some(libc::ESRCH) || errno == Some(libc::EPERM) {
gone = true;
break;
}
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
assert!(
gone,
"grandchild {grandchild} survived; group kill did not reach it"
);
}
}