Terminate leftover child processes of the entrypoint after bwrap exits

Cargo.lock

@@ -11,6 +11,7 @@ dependencies = [
  "glob",
  "libc",
  "nix",
+ "rustix",
  "seccompiler",
  "serde",
  "shlex",

Cargo.toml

@@ -23,6 +23,7 @@ clap = { version = "4", features = ["derive"] }
 glob = "0.3"
 libc = "0.2"
 nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
+rustix = { version = "1", features = ["process"] }
 seccompiler = "0.5"
 serde = { version = "1", features = ["derive"] }
 shlex = "2.0.1"

src/lib.rs

@@ -7,6 +7,7 @@ mod errors;
 mod memfd;
 mod persistent_tmp;
 mod preflight;
+mod process_tree;
 mod sandbox;
 mod sandbox_info;
 mod seccomp;
@@ -20,9 +21,12 @@ use std::ffi::OsString;
 use std::fs;
 use std::path::PathBuf;
 use std::process::ExitStatus;
+use std::time::Duration;
 
 use crate::sandbox::Bwrap;
 
+const CLEANUP_GRACE_PERIOD: Duration = Duration::from_secs(5);
+
 pub enum SandboxMode {
     Blacklist,
     Whitelist,
@@ -103,5 +107,8 @@ pub fn run(config: SandboxConfig) -> Result<ExitStatus, SandboxError> {
         return Ok(ExitStatus::default());
     }
 
-    Ok(bwrap.run()?)
+    process_tree::become_subreaper()?;
+    let status = bwrap.run()?;
+    process_tree::terminate_and_reap_all(CLEANUP_GRACE_PERIOD)?;
+    Ok(status)
 }

src/process_tree.rs

@@ -0,0 +1,91 @@
+use std::fs;
+use std::io;
+use std::thread;
+use std::time::{Duration, Instant};
+
+use nix::sys::prctl;
+use rustix::io::Errno;
+use rustix::process::{self, Pid, PidfdFlags, Signal, WaitOptions};
+
+const REAP_POLL_INTERVAL: Duration = Duration::from_millis(10);
+
+pub fn become_subreaper() -> io::Result<()> {
+    Ok(prctl::set_child_subreaper(true)?)
+}
+
+pub fn terminate_and_reap_all(grace_period: Duration) -> io::Result<()> {
+    signal_tree(process::getpid(), &[Signal::TERM, Signal::CONT]);
+
+    if reap_until(Instant::now() + grace_period)? {
+        return Ok(());
+    }
+
+    kill_and_reap_all()
+}
+
+fn signal_tree(parent: Pid, signals: &[Signal]) {
+    for child in read_children(parent) {
+        if signal_if_child_of(parent, child, signals) {
+            signal_tree(child, signals);
+        }
+    }
+}
+
+fn read_children(pid: Pid) -> Vec<Pid> {
+    let Ok(tasks) = fs::read_dir(format!("/proc/{}/task", pid.as_raw_pid())) else {
+        return Vec::new();
+    };
+
+    tasks
+        .filter_map(|task| fs::read_to_string(task.ok()?.path().join("children")).ok())
+        .flat_map(|children| parse_pids(&children))
+        .collect()
+}
+
+fn parse_pids(text: &str) -> Vec<Pid> {
+    text.split_whitespace()
+        .filter_map(|n| Pid::from_raw(n.parse().ok()?))
+        .collect()
+}
+
+fn signal_if_child_of(parent: Pid, pid: Pid, signals: &[Signal]) -> bool {
+    let Ok(pidfd) = process::pidfd_open(pid, PidfdFlags::empty()) else {
+        return false;
+    };
+
+    read_parent(pid) == Some(parent)
+        && signals
+            .iter()
+            .all(|&signal| process::pidfd_send_signal(&pidfd, signal).is_ok())
+}
+
+fn read_parent(pid: Pid) -> Option<Pid> {
+    let stat = fs::read_to_string(format!("/proc/{}/stat", pid.as_raw_pid())).ok()?;
+    let fields_after_comm = &stat[stat.rfind(')')? + 1..];
+    let ppid = fields_after_comm.split_whitespace().nth(1)?;
+    Pid::from_raw(ppid.parse().ok()?)
+}
+
+fn reap_until(deadline: Instant) -> io::Result<bool> {
+    loop {
+        match process::wait(WaitOptions::NOHANG) {
+            Ok(Some(_)) => {}
+            Ok(None) if Instant::now() < deadline => thread::sleep(REAP_POLL_INTERVAL),
+            Ok(None) => return Ok(false),
+            Err(Errno::CHILD) => return Ok(true),
+            Err(e) => return Err(e.into()),
+        }
+    }
+}
+
+fn kill_and_reap_all() -> io::Result<()> {
+    loop {
+        signal_tree(process::getpid(), &[Signal::KILL]);
+
+        match process::wait(WaitOptions::empty()) {
+            Ok(_) => {}
+            Err(Errno::CHILD) => return Ok(()),
+            Err(e) => return Err(e.into()),
+        }
+    }
+}

tests/e2e/signals.rs

@@ -1,4 +1,10 @@
+use std::fs;
 use std::os::unix::process::ExitStatusExt;
+use std::path::Path;
+use std::process;
+use std::sync::atomic::{AtomicU32, Ordering};
+
+use tempfile::TempDir;
 
 use crate::common::*;
 
@@ -31,3 +37,132 @@ fn signal_death_of_bwrap_is_mirrored() {
 
     assert_eq!(status.signal(), Some(libc::SIGKILL));
 }
+
+#[test]
+fn detached_child_process_is_killed_after_entrypoint_exits_in_blacklist_mode() {
+    assert_detached_child_process_is_killed("--blacklist", "");
+}
+
+#[test]
+fn detached_child_process_is_killed_after_entrypoint_exits_in_whitelist_mode() {
+    assert_detached_child_process_is_killed("--whitelist", "");
+}
+
+#[test]
+fn detached_child_process_that_ignores_sigterm_is_killed() {
+    assert_detached_child_process_is_killed("--blacklist", "trap \"\" TERM;");
+}
+
+#[test]
+fn detached_child_process_receives_sigterm_before_sigkill() {
+    let dir = TempDir::new().expect("failed to create temp dir");
+    let flag = dir.path().join("terminated");
+    let setup = format!("trap \"touch {}; exit\" TERM;", flag.display());
+
+    let marker = create_unique_sleep_duration();
+    run_with_detached_sleeper(
+        &["--blacklist", "--rw", dir.path().to_str().unwrap()],
+        &setup,
+        &marker,
+    );
+
+    assert!(
+        flag.exists(),
+        "detached child process did not receive SIGTERM"
+    );
+}
+
+#[test]
+fn every_process_in_a_large_tree_receives_sigterm() {
+    let dir = TempDir::new().expect("failed to create temp dir");
+    let marker = create_unique_sleep_duration();
+    let node_script = dir.path().join("node.sh");
+    fs::write(&node_script, TREE_NODE_SCRIPT.replace("{marker}", &marker)).unwrap();
+    let script = format!(
+        "sh {node} root 0 &
+         until [ $(ls {dir} | grep -c '^ready\\.') -eq {TREE_SIZE} ]; do sleep 0.01; done",
+        node = node_script.display(),
+        dir = dir.path().display(),
+    );
+
+    let status = Sandbox::new(&["--blacklist", "--rw", dir.path().to_str().unwrap()])
+        .args(["--", "sh", "-c", &script])
+        .status()
+        .expect("agent-sandbox binary failed to execute");
+    assert!(status.success());
+
+    let terminated = count_files_with_prefix(dir.path(), "terminated.");
+    let leftovers = find_processes_with_arg(&marker);
+    for pid in &leftovers {
+        unsafe { libc::kill(*pid, libc::SIGKILL) };
+    }
+    assert_eq!(terminated, TREE_SIZE, "processes that received SIGTERM");
+    assert!(leftovers.is_empty(), "leftover processes: {leftovers:?}");
+}
+
+fn assert_detached_child_process_is_killed(mode: &str, setup: &str) {
+    let marker = create_unique_sleep_duration();
+    run_with_detached_sleeper(&[mode], setup, &marker);
+
+    let leftovers = find_processes_with_arg(&marker);
+    for pid in &leftovers {
+        unsafe { libc::kill(*pid, libc::SIGKILL) };
+    }
+    assert!(leftovers.is_empty(), "leftover processes: {leftovers:?}");
+}
+
+fn create_unique_sleep_duration() -> String {
+    static NEXT: AtomicU32 = AtomicU32::new(0);
+    let sequence = NEXT.fetch_add(1, Ordering::Relaxed);
+    format!("3600.{}{sequence}", process::id())
+}
+
+fn run_with_detached_sleeper(sandbox_args: &[&str], setup: &str, marker: &str) {
+    let script = format!(
+        "setsid -f sh -c '{setup} sleep {marker} & wait' </dev/null >/dev/null 2>&1
+         until pgrep -fx 'sleep {marker}' >/dev/null; do sleep 0.01; done"
+    );
+
+    let status = Sandbox::new(sandbox_args)
+        .args(["--", "sh", "-c", &script])
+        .status()
+        .expect("agent-sandbox binary failed to execute");
+    assert!(status.success());
+}
+
+fn find_processes_with_arg(arg: &str) -> Vec<i32> {
+    fs::read_dir("/proc")
+        .unwrap()
+        .filter_map(|entry| entry.ok()?.file_name().to_str()?.parse().ok())
+        .filter(|pid: &i32| {
+            fs::read(format!("/proc/{pid}/cmdline"))
+                .is_ok_and(|cmdline| cmdline.split(|&b| b == 0).any(|a| a == arg.as_bytes()))
+        })
+        .collect()
+}
+
+const TREE_SIZE: usize = 1 + 3 + 9 + 27 + 81;
+
+const TREE_NODE_SCRIPT: &str = r#"
+name=$1 depth=$2 dir=$(dirname "$0")
+trap 'touch "$dir/terminated.$name"; exit' TERM
+if [ "$depth" -lt 4 ]; then
+    setsid -f sh "$0" "$name.1" $((depth + 1)) </dev/null >/dev/null 2>&1
+    sh "$0" "$name.2" $((depth + 1)) &
+    sh "$0" "$name.3" $((depth + 1)) &
+fi
+sleep {marker} &
+touch "$dir/ready.$name"
+wait
+"#;
+
+fn count_files_with_prefix(dir: &Path, prefix: &str) -> usize {
+    fs::read_dir(dir)
+        .unwrap()
+        .filter(|entry| {
+            entry
+                .as_ref()
+                .is_ok_and(|e| e.file_name().to_string_lossy().starts_with(prefix))
+        })
+        .count()
+}