acyclic_native_runtime/
process_tree.rs

1#![allow(unsafe_code)]
2
3use std::io;
4use std::process::{Child, Command, ExitStatus, Output};
5
6/// One child and descendants that remain in its operating-system containment.
7///
8/// Windows uses a Job object. Unix uses a process group, which cannot contain
9/// a descendant that deliberately creates a new process group or session.
10pub struct ProcessTree {
11    child: Option<Child>,
12    guard: platform::Guard,
13}
14
15impl ProcessTree {
16    pub(crate) fn spawn(command: &mut Command) -> io::Result<Self> {
17        let (child, guard) = platform::spawn(command)?;
18        Ok(Self {
19            child: Some(child),
20            guard,
21        })
22    }
23
24    /// Polls the direct child without releasing ownership of its descendants.
25    pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
26        self.child
27            .as_mut()
28            .ok_or_else(|| io::Error::other("process output was already collected"))?
29            .try_wait()
30    }
31
32    /// Waits for the direct child while retaining ownership of its descendants.
33    pub fn wait(&mut self) -> io::Result<ExitStatus> {
34        self.child
35            .as_mut()
36            .ok_or_else(|| io::Error::other("process output was already collected"))?
37            .wait()
38    }
39
40    /// Collects the direct child's output while retaining the descendant guard.
41    pub fn wait_with_output(&mut self) -> io::Result<Output> {
42        self.child
43            .take()
44            .ok_or_else(|| io::Error::other("process output was already collected"))?
45            .wait_with_output()
46    }
47
48    /// Terminates every process in the tree and reaps the direct child.
49    pub fn terminate(&mut self) -> io::Result<()> {
50        self.terminate_descendants()?;
51        if let Some(child) = self.child.as_mut() {
52            child.wait()?;
53            self.child.take();
54        }
55        Ok(())
56    }
57
58    /// Signals termination to the entire tree while retaining the direct child
59    /// so its captured output can still be collected.
60    pub fn terminate_descendants(&mut self) -> io::Result<()> {
61        self.guard.terminate()
62    }
63}
64
65impl Drop for ProcessTree {
66    fn drop(&mut self) {
67        let _ = self.terminate();
68    }
69}
70
71#[cfg(unix)]
72mod platform {
73    use std::io;
74    use std::os::unix::process::CommandExt as _;
75    use std::process::{Child, Command};
76
77    pub(super) struct Guard {
78        process_group: libc::pid_t,
79        active: bool,
80    }
81
82    pub(super) fn spawn(command: &mut Command) -> io::Result<(Child, Guard)> {
83        command.process_group(0);
84        let child = command.spawn()?;
85        let process_group = libc::pid_t::try_from(child.id())
86            .map_err(|_| io::Error::other("child process id does not fit pid_t"))?;
87        Ok((
88            child,
89            Guard {
90                process_group,
91                active: true,
92            },
93        ))
94    }
95
96    impl Guard {
97        pub(super) fn terminate(&mut self) -> io::Result<()> {
98            if !self.active {
99                return Ok(());
100            }
101            // SAFETY: a negative, nonzero pid addresses exactly this owned
102            // process group; SIGKILL requires no shared memory or signal data.
103            if unsafe { libc::kill(-self.process_group, libc::SIGKILL) } == 0 {
104                self.active = false;
105                return Ok(());
106            }
107            let error = io::Error::last_os_error();
108            if error.raw_os_error() == Some(libc::ESRCH) {
109                self.active = false;
110                Ok(())
111            } else {
112                Err(error)
113            }
114        }
115    }
116}
117
118#[cfg(test)]
119#[allow(clippy::expect_used)]
120mod tests {
121    use super::ProcessTree;
122    use std::fs;
123    use std::process::{Command, Stdio};
124    use std::thread;
125    use std::time::{Duration, Instant};
126
127    const MODE: &str = "ACYCLIC_PROCESS_TREE_TEST_MODE";
128    const ROOT: &str = "ACYCLIC_PROCESS_TREE_TEST_ROOT";
129
130    #[test]
131    fn process_tree_helper() {
132        let Ok(mode) = std::env::var(MODE) else {
133            return;
134        };
135        let root = std::path::PathBuf::from(std::env::var_os(ROOT).expect("helper root"));
136        if mode == "grandchild" {
137            fs::write(root.join("grandchild-ready"), b"ready").expect("grandchild ready");
138            thread::sleep(Duration::from_millis(750));
139            fs::write(root.join("escaped"), b"descendant survived").expect("escaped marker");
140            return;
141        }
142        assert_eq!(mode, "child");
143        let mut grandchild = Command::new(std::env::current_exe().expect("test executable"));
144        grandchild
145            .args(["--exact", "process_tree::tests::process_tree_helper"])
146            .env(MODE, "grandchild")
147            .env(ROOT, &root)
148            .stdin(Stdio::null())
149            .stdout(Stdio::null())
150            .stderr(Stdio::null());
151        let mut grandchild = grandchild.spawn().expect("spawn grandchild");
152        let deadline = Instant::now() + Duration::from_secs(5);
153        while !root.join("grandchild-ready").exists() && Instant::now() < deadline {
154            thread::sleep(Duration::from_millis(10));
155        }
156        assert!(root.join("grandchild-ready").exists());
157        fs::write(root.join("tree-ready"), b"ready").expect("tree ready");
158        let _ = grandchild.wait();
159    }
160
161    #[test]
162    fn termination_contains_descendants() {
163        let temporary = tempfile::tempdir().expect("temporary process-tree directory");
164        let mut command = Command::new(std::env::current_exe().expect("test executable"));
165        command
166            .args(["--exact", "process_tree::tests::process_tree_helper"])
167            .env(MODE, "child")
168            .env(ROOT, temporary.path())
169            .stdin(Stdio::null())
170            .stdout(Stdio::null())
171            .stderr(Stdio::null());
172        let mut tree = ProcessTree::spawn(&mut command).expect("spawn process tree");
173        let deadline = Instant::now() + Duration::from_secs(5);
174        while !temporary.path().join("tree-ready").exists() && Instant::now() < deadline {
175            thread::sleep(Duration::from_millis(10));
176        }
177        assert!(temporary.path().join("tree-ready").exists());
178        tree.terminate().expect("terminate process tree");
179        thread::sleep(Duration::from_secs(1));
180        assert!(!temporary.path().join("escaped").exists());
181    }
182}
183
184#[cfg(windows)]
185mod platform {
186    use std::io;
187    use std::mem::{size_of, zeroed};
188    use std::os::windows::io::AsRawHandle as _;
189    use std::os::windows::process::CommandExt as _;
190    use std::process::{Child, Command};
191    use windows_sys::Win32::Foundation::{CloseHandle, HANDLE, INVALID_HANDLE_VALUE};
192    use windows_sys::Win32::System::Diagnostics::ToolHelp::{
193        CreateToolhelp32Snapshot, TH32CS_SNAPTHREAD, THREADENTRY32, Thread32First, Thread32Next,
194    };
195    use windows_sys::Win32::System::JobObjects::{
196        AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
197        JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation,
198        SetInformationJobObject, TerminateJobObject,
199    };
200    use windows_sys::Win32::System::Threading::{
201        CREATE_NEW_PROCESS_GROUP, CREATE_SUSPENDED, OpenThread, ResumeThread, THREAD_SUSPEND_RESUME,
202    };
203
204    pub(super) struct Guard {
205        job: HANDLE,
206        active: bool,
207    }
208
209    pub(super) fn spawn(command: &mut Command) -> io::Result<(Child, Guard)> {
210        let mut guard = Guard::new()?;
211        command.creation_flags(CREATE_SUSPENDED | CREATE_NEW_PROCESS_GROUP);
212        let mut child = command.spawn()?;
213        // SAFETY: the Job and Child each own live handles for the duration of
214        // this call. The suspended child cannot create descendants before it
215        // has been assigned to the kill-on-close Job.
216        if unsafe { AssignProcessToJobObject(guard.job, child.as_raw_handle().cast()) } == 0 {
217            let error = io::Error::last_os_error();
218            let _ = guard.terminate();
219            let _ = child.kill();
220            let _ = child.wait();
221            return Err(error);
222        }
223        if let Err(error) = resume_process(child.id()) {
224            let _ = guard.terminate();
225            let _ = child.wait();
226            return Err(error);
227        }
228        Ok((child, guard))
229    }
230
231    impl Guard {
232        fn new() -> io::Result<Self> {
233            // SAFETY: null security attributes and name request an unnamed Job
234            // with default access controls.
235            let job = unsafe { CreateJobObjectW(std::ptr::null(), std::ptr::null()) };
236            if job.is_null() {
237                return Err(io::Error::last_os_error());
238            }
239            // SAFETY: the structure is plain Win32 data and zero is its
240            // documented initial state before setting the requested limit.
241            let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { zeroed() };
242            limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
243            // SAFETY: `limits` is initialized and its exact size is provided.
244            let configured = unsafe {
245                SetInformationJobObject(
246                    job,
247                    JobObjectExtendedLimitInformation,
248                    (&raw const limits).cast(),
249                    u32::try_from(size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>())
250                        .map_err(|_| io::Error::other("Job limit structure is too large"))?,
251                )
252            };
253            if configured == 0 {
254                let error = io::Error::last_os_error();
255                // SAFETY: `job` is an owned live handle.
256                unsafe { CloseHandle(job) };
257                return Err(error);
258            }
259            Ok(Self { job, active: true })
260        }
261
262        pub(super) fn terminate(&mut self) -> io::Result<()> {
263            if !self.active {
264                return Ok(());
265            }
266            // SAFETY: `job` remains owned until Drop closes it.
267            if unsafe { TerminateJobObject(self.job, 1) } == 0 {
268                Err(io::Error::last_os_error())
269            } else {
270                self.active = false;
271                Ok(())
272            }
273        }
274    }
275
276    impl Drop for Guard {
277        fn drop(&mut self) {
278            let _ = self.terminate();
279            // SAFETY: `job` is owned by this guard and closed exactly once.
280            unsafe { CloseHandle(self.job) };
281        }
282    }
283
284    fn resume_process(process_id: u32) -> io::Result<()> {
285        // SAFETY: flags request a system-owned snapshot handle.
286        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0) };
287        if snapshot == INVALID_HANDLE_VALUE {
288            return Err(io::Error::last_os_error());
289        }
290        // SAFETY: zero is the documented initialization; dwSize is set below.
291        let mut entry: THREADENTRY32 = unsafe { zeroed() };
292        entry.dwSize = u32::try_from(size_of::<THREADENTRY32>())
293            .map_err(|_| io::Error::other("thread entry structure is too large"))?;
294        // SAFETY: snapshot and entry are live for the iteration.
295        let mut found = unsafe { Thread32First(snapshot, &raw mut entry) } != 0;
296        let mut result = Err(io::Error::new(
297            io::ErrorKind::NotFound,
298            "suspended child thread was not found",
299        ));
300        while found {
301            if entry.th32OwnerProcessID == process_id {
302                // SAFETY: the thread id came from the live system snapshot.
303                let thread = unsafe { OpenThread(THREAD_SUSPEND_RESUME, 0, entry.th32ThreadID) };
304                if thread.is_null() {
305                    result = Err(io::Error::last_os_error());
306                } else {
307                    // SAFETY: `thread` is a live handle with resume access.
308                    result = if unsafe { ResumeThread(thread) } == u32::MAX {
309                        Err(io::Error::last_os_error())
310                    } else {
311                        Ok(())
312                    };
313                    // SAFETY: this function owns the thread handle.
314                    unsafe { CloseHandle(thread) };
315                }
316                break;
317            }
318            // SAFETY: snapshot and entry remain live for the next iteration.
319            found = unsafe { Thread32Next(snapshot, &raw mut entry) } != 0;
320        }
321        // SAFETY: this function owns the snapshot handle.
322        unsafe { CloseHandle(snapshot) };
323        result
324    }
325}
acyclic_native_runtime/process_tree.rs

acyclic_native_runtime/
process_tree.rs
