Repeat the process tree walk and spare processes started by signal handlers

src/errors.rs

@@ -34,6 +34,7 @@ pub enum SandboxError {
     Seccomp(String),
     SeccompUnsupportedArch(String),
     PersistentTmpDir(PathBuf, String),
+    ProcChildrenMissing,
 }
 
 impl std::fmt::Display for SandboxError {
@@ -111,6 +112,10 @@ impl std::fmt::Display for SandboxError {
             Self::PersistentTmpDir(path, reason) => {
                 write!(f, "persistent-tmp directory '{}': {reason}", path.display())
             }
+            Self::ProcChildrenMissing => write!(
+                f,
+                "the kernel does not provide /proc/<pid>/task/<tid>/children (CONFIG_PROC_CHILDREN); agent-sandbox needs it to find the processes it must clean up"
+            ),
         }
     }
 }

src/preflight.rs

@@ -1,13 +1,17 @@
 use std::os::unix::fs::PermissionsExt;
+use std::path::Path;
 use std::process::Command;
 
 use crate::SandboxConfig;
 use crate::errors::SandboxError;
 
+const PROC_CHILDREN: &str = "/proc/thread-self/children";
+
 pub fn check(config: &SandboxConfig) -> Result<(), SandboxError> {
     check_bwrap()?;
     check_command(config)?;
     check_chdir(config)?;
+    check_proc_children(Path::new(PROC_CHILDREN))?;
     Ok(())
 }
 
@@ -36,3 +40,14 @@ fn check_command(config: &SandboxConfig) -> Result<(), SandboxError> {
     }
     Ok(())
 }
+
+fn check_proc_children(path: &Path) -> Result<(), SandboxError> {
+    if !path.exists() {
+        return Err(SandboxError::ProcChildrenMissing);
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+#[path = "../tests/unit/preflight.rs"]
+mod tests;

src/process_tree.rs

@@ -1,5 +1,7 @@
+use std::collections::HashSet;
 use std::fs;
 use std::io;
+use std::os::fd::OwnedFd;
 use std::thread;
 use std::time::{Duration, Instant};
 
@@ -8,13 +10,14 @@ use rustix::io::Errno;
 use rustix::process::{self, Pid, PidfdFlags, Signal, WaitOptions};
 
 const REAP_POLL_INTERVAL: Duration = Duration::from_millis(10);
+const MAX_WALK_PASSES: usize = 16;
 
 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]);
+    signal_all_children(&[Signal::TERM, Signal::CONT]);
 
     if reap_until(Instant::now() + grace_period)? {
         return Ok(());
@@ -23,10 +26,14 @@ pub fn terminate_and_reap_all(grace_period: Duration) -> io::Result<()> {
     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);
+// Not exact: a process that forks during the last pass can miss the signals.
+// That is good enough, because every survivor receives SIGKILL after the grace period.
+fn signal_all_children(signals: &[Signal]) {
+    let mut signaled = HashSet::new();
+    for _ in 0..MAX_WALK_PASSES {
+        let root = process::getpid();
+        if !signal_new_children(root, read_children(root), signals, &mut signaled) {
+            return;
         }
     }
 }
@@ -48,15 +55,36 @@ fn parse_pids(text: &str) -> Vec<Pid> {
         .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;
-    };
+fn signal_new_children(
+    parent: Pid,
+    children: Vec<Pid>,
+    signals: &[Signal],
+    signaled: &mut HashSet<Pid>,
+) -> bool {
+    let mut has_signaled_new = false;
 
-    read_parent(pid) == Some(parent)
-        && signals
-            .iter()
-            .all(|&signal| process::pidfd_send_signal(&pidfd, signal).is_ok())
+    for child in children {
+        if signaled.contains(&child) {
+            continue;
+        }
+        let Some(pidfd) = open_if_child_of(parent, child) else {
+            continue;
+        };
+
+        let children_before_signal = read_children(child);
+        send_signals(&pidfd, signals);
+        signaled.insert(child);
+        has_signaled_new = true;
+
+        signal_new_children(child, children_before_signal, signals, signaled);
+    }
+
+    has_signaled_new
+}
+
+fn open_if_child_of(parent: Pid, pid: Pid) -> Option<OwnedFd> {
+    let pidfd = process::pidfd_open(pid, PidfdFlags::empty()).ok()?;
+    (read_parent(pid) == Some(parent)).then_some(pidfd)
 }
 
 fn read_parent(pid: Pid) -> Option<Pid> {
@@ -66,6 +94,12 @@ fn read_parent(pid: Pid) -> Option<Pid> {
     Pid::from_raw(ppid.parse().ok()?)
 }
 
+fn send_signals(pidfd: &OwnedFd, signals: &[Signal]) {
+    for &signal in signals {
+        let _ = process::pidfd_send_signal(pidfd, signal);
+    }
+}
+
 fn reap_until(deadline: Instant) -> io::Result<bool> {
     loop {
         match process::wait(WaitOptions::NOHANG) {
@@ -80,7 +114,7 @@ fn reap_until(deadline: Instant) -> io::Result<bool> {
 
 fn kill_and_reap_all() -> io::Result<()> {
     loop {
-        signal_tree(process::getpid(), &[Signal::KILL]);
+        signal_all_children(&[Signal::KILL]);
 
         match process::wait(WaitOptions::empty()) {
             Ok(_) => {}