Terminate leftover child processes of the entrypoint after bwrap exits

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()),
+        }
+    }
+}