diff --git a/Cargo.lock b/Cargo.lock index 0727f20..a5a1750 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,7 @@ dependencies = [ "glob", "libc", "nix", + "rustix", "seccompiler", "serde", "shlex", diff --git a/Cargo.toml b/Cargo.toml index 53e5c7f..e91db06 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/lib.rs b/src/lib.rs index 934b211..8dfee5f 100644 --- a/src/lib.rs +++ b/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 { 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) } diff --git a/src/process_tree.rs b/src/process_tree.rs new file mode 100644 index 0000000..033df33 --- /dev/null +++ b/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 { + 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 { + 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 { + 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 { + 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()), + } + } +} diff --git a/tests/e2e/signals.rs b/tests/e2e/signals.rs index bd64f14..104a810 100644 --- a/tests/e2e/signals.rs +++ b/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 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 { + 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 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() +}