rightkit_process/
owned.rs

1use std::{
2    ffi::OsStr,
3    io,
4    num::NonZeroU32,
5    process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, ExitStatus},
6    sync::{
7        atomic::{AtomicU32, Ordering},
8        Arc,
9    },
10    thread,
11    time::{Duration, Instant},
12};
13
14use process_wrap::std::{StdChildWrapper, StdCommandWrap, StdCommandWrapper};
15
16#[derive(Debug)]
17pub struct OwnedCommand {
18    command: Command,
19    windows_hidden: bool,
20}
21
22impl OwnedCommand {
23    pub fn new(program: impl AsRef<OsStr>) -> Self {
24        Self::from_command(Command::new(program))
25    }
26
27    pub fn from_command(command: Command) -> Self {
28        Self {
29            command,
30            windows_hidden: false,
31        }
32    }
33
34    pub fn command_mut(&mut self) -> &mut Command {
35        &mut self.command
36    }
37
38    pub fn windows_hide(&mut self) -> &mut Self {
39        self.windows_hidden = true;
40        self
41    }
42
43    pub fn spawn(self) -> io::Result<OwnedChild> {
44        spawn_owned(self.command, self.windows_hidden, false, None)
45    }
46}
47
48#[derive(Debug)]
49pub struct OwnedChild {
50    child: Option<Box<dyn StdChildWrapper>>,
51}
52
53#[derive(Clone, Copy, Debug, Eq, PartialEq)]
54pub enum WaitOutcome {
55    Exited(ExitStatus),
56    Terminated(ExitStatus),
57}
58
59impl OwnedChild {
60    pub fn id(&self) -> u32 {
61        self.child().id()
62    }
63
64    pub fn take_stdin(&mut self) -> Option<ChildStdin> {
65        self.child_mut().stdin().take()
66    }
67
68    pub fn take_stdout(&mut self) -> Option<ChildStdout> {
69        self.child_mut().stdout().take()
70    }
71
72    pub fn take_stderr(&mut self) -> Option<ChildStderr> {
73        self.child_mut().stderr().take()
74    }
75
76    pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
77        self.child_mut().try_wait()
78    }
79
80    pub fn wait(&mut self) -> io::Result<ExitStatus> {
81        self.child_mut().wait()
82    }
83
84    pub fn wait_timeout(&mut self, timeout: Duration) -> io::Result<Option<ExitStatus>> {
85        let started = Instant::now();
86        loop {
87            if let Some(status) = self.try_wait()? {
88                return Ok(Some(status));
89            }
90            if started.elapsed() >= timeout {
91                return Ok(None);
92            }
93            thread::sleep(Duration::from_millis(5).min(timeout.saturating_sub(started.elapsed())));
94        }
95    }
96
97    pub fn wait_or_kill(&mut self, grace: Duration) -> io::Result<WaitOutcome> {
98        #[cfg(unix)]
99        self.child().signal(nix::libc::SIGTERM)?;
100
101        if let Some(status) = self.wait_timeout(grace)? {
102            return Ok(WaitOutcome::Exited(status));
103        }
104        self.terminate_tree().map(WaitOutcome::Terminated)
105    }
106
107    pub fn terminate_tree(&mut self) -> io::Result<ExitStatus> {
108        if let Some(status) = self.try_wait()? {
109            return Ok(status);
110        }
111        self.child_mut().start_kill()?;
112        self.child_mut().wait()
113    }
114
115    fn child(&self) -> &dyn StdChildWrapper {
116        self.child
117            .as_deref()
118            .expect("owned child is always present")
119    }
120
121    fn child_mut(&mut self) -> &mut dyn StdChildWrapper {
122        self.child
123            .as_deref_mut()
124            .expect("owned child is always present")
125    }
126}
127
128impl Drop for OwnedChild {
129    fn drop(&mut self) {
130        let _ = self.terminate_tree();
131    }
132}
133
134#[derive(Clone, Copy, Debug, Eq, PartialEq)]
135pub struct AdoptedProcess {
136    pid: NonZeroU32,
137}
138
139impl AdoptedProcess {
140    pub fn new(pid: NonZeroU32) -> Self {
141        Self { pid }
142    }
143
144    pub fn id(&self) -> u32 {
145        self.pid.get()
146    }
147
148    pub fn is_running(&self) -> io::Result<bool> {
149        process_is_running(self.pid)
150    }
151}
152
153#[derive(Debug)]
154struct CleanupWrapper {
155    captured_pid: Option<Arc<AtomicU32>>,
156}
157
158impl StdCommandWrapper for CleanupWrapper {
159    fn post_spawn(&mut self, child: &mut Child, _core: &StdCommandWrap) -> io::Result<()> {
160        if let Some(pid) = &self.captured_pid {
161            pid.store(child.id(), Ordering::SeqCst);
162        }
163        Ok(())
164    }
165
166    fn wrap_child(
167        &mut self,
168        child: Box<dyn StdChildWrapper>,
169        _core: &StdCommandWrap,
170    ) -> io::Result<Box<dyn StdChildWrapper>> {
171        Ok(Box::new(CleanupChild { child: Some(child) }))
172    }
173}
174
175#[derive(Debug)]
176struct CleanupChild {
177    child: Option<Box<dyn StdChildWrapper>>,
178}
179
180impl CleanupChild {
181    fn child(&self) -> &dyn StdChildWrapper {
182        self.child.as_deref().expect("cleanup child is present")
183    }
184
185    fn child_mut(&mut self) -> &mut dyn StdChildWrapper {
186        self.child.as_deref_mut().expect("cleanup child is present")
187    }
188}
189
190impl StdChildWrapper for CleanupChild {
191    fn inner(&self) -> &Child {
192        self.child().inner()
193    }
194
195    fn inner_mut(&mut self) -> &mut Child {
196        self.child_mut().inner_mut()
197    }
198
199    fn into_inner(mut self: Box<Self>) -> Child {
200        self.child
201            .take()
202            .expect("cleanup child is present")
203            .into_inner()
204    }
205
206    fn stdin(&mut self) -> &mut Option<ChildStdin> {
207        self.child_mut().stdin()
208    }
209
210    fn stdout(&mut self) -> &mut Option<ChildStdout> {
211        self.child_mut().stdout()
212    }
213
214    fn stderr(&mut self) -> &mut Option<ChildStderr> {
215        self.child_mut().stderr()
216    }
217
218    fn id(&self) -> u32 {
219        self.child().id()
220    }
221
222    fn start_kill(&mut self) -> io::Result<()> {
223        self.child_mut().start_kill()
224    }
225
226    fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
227        self.child_mut().try_wait()
228    }
229
230    fn wait(&mut self) -> io::Result<ExitStatus> {
231        self.child_mut().wait()
232    }
233
234    #[cfg(unix)]
235    fn signal(&self, signal: i32) -> io::Result<()> {
236        self.child().signal(signal)
237    }
238}
239
240impl Drop for CleanupChild {
241    fn drop(&mut self) {
242        let Some(child) = self.child.as_deref_mut() else {
243            return;
244        };
245        if !matches!(child.try_wait(), Ok(Some(_))) {
246            let _ = child.start_kill();
247            let _ = child.wait();
248        }
249    }
250}
251
252#[derive(Debug)]
253struct FailAfterSpawn;
254
255impl StdCommandWrapper for FailAfterSpawn {
256    fn wrap_child(
257        &mut self,
258        _child: Box<dyn StdChildWrapper>,
259        _core: &StdCommandWrap,
260    ) -> io::Result<Box<dyn StdChildWrapper>> {
261        Err(io::Error::other("injected post-spawn wrapping failure"))
262    }
263}
264
265fn spawn_owned(
266    command: Command,
267    windows_hidden: bool,
268    fail_after_spawn: bool,
269    captured_pid: Option<Arc<AtomicU32>>,
270) -> io::Result<OwnedChild> {
271    let mut command = StdCommandWrap::from(command);
272    command.wrap(CleanupWrapper { captured_pid });
273
274    #[cfg(windows)]
275    {
276        use process_wrap::std::{CreationFlags, JobObject};
277        use windows::Win32::System::Threading::CREATE_NO_WINDOW;
278
279        if windows_hidden {
280            command.wrap(CreationFlags(CREATE_NO_WINDOW));
281        }
282        command.wrap(JobObject);
283    }
284
285    #[cfg(unix)]
286    {
287        use process_wrap::std::ProcessGroup;
288        let _ = windows_hidden;
289        command.wrap(ProcessGroup::leader());
290    }
291
292    if fail_after_spawn {
293        command.wrap(FailAfterSpawn);
294    }
295
296    command
297        .spawn()
298        .map(|child| OwnedChild { child: Some(child) })
299}
300
301#[cfg(test)]
302fn spawn_for_test(command: Command, captured_pid: Arc<AtomicU32>) -> io::Result<OwnedChild> {
303    spawn_owned(command, true, true, Some(captured_pid))
304}
305
306#[cfg(windows)]
307fn process_is_running(pid: NonZeroU32) -> io::Result<bool> {
308    use windows::Win32::{
309        Foundation::{CloseHandle, ERROR_INVALID_PARAMETER, WAIT_TIMEOUT},
310        System::Threading::{
311            OpenProcess, WaitForSingleObject, PROCESS_QUERY_LIMITED_INFORMATION,
312            PROCESS_SYNCHRONIZE,
313        },
314    };
315
316    let handle = unsafe {
317        OpenProcess(
318            PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_SYNCHRONIZE,
319            false,
320            pid.get(),
321        )
322    };
323    let handle = match handle {
324        Ok(handle) => handle,
325        Err(error) => {
326            if error.code() == ERROR_INVALID_PARAMETER.to_hresult() {
327                return Ok(false);
328            }
329            return Err(io::Error::other(error.to_string()));
330        }
331    };
332    let wait = unsafe { WaitForSingleObject(handle, 0) };
333    let close = unsafe { CloseHandle(handle) };
334    close.map_err(|error| io::Error::other(error.to_string()))?;
335    match wait.0 {
336        0 => Ok(false),
337        value if value == WAIT_TIMEOUT.0 => Ok(true),
338        _ => Err(io::Error::last_os_error()),
339    }
340}
341
342#[cfg(unix)]
343fn process_is_running(pid: NonZeroU32) -> io::Result<bool> {
344    use nix::{errno::Errno, sys::signal::kill, unistd::Pid};
345
346    let pid = i32::try_from(pid.get())
347        .map(Pid::from_raw)
348        .map_err(io::Error::other)?;
349    match kill(pid, None) {
350        Ok(()) | Err(Errno::EPERM) => Ok(true),
351        Err(Errno::ESRCH) => Ok(false),
352        Err(error) => Err(io::Error::from(error)),
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    #[test]
359    fn a_failure_after_spawn_cleans_up_the_partially_wrapped_child() {
360        use std::{
361            process::Command,
362            sync::{
363                atomic::{AtomicU32, Ordering},
364                Arc,
365            },
366            thread,
367            time::{Duration, Instant},
368        };
369
370        // A long-lived child that survives on its own if the partial-spawn cleanup fails.
371        let command = if cfg!(windows) {
372            let mut command = Command::new("powershell.exe");
373            command.args([
374                "-NoLogo",
375                "-NoProfile",
376                "-NonInteractive",
377                "-Command",
378                "Start-Sleep -Seconds 60",
379            ]);
380            command
381        } else {
382            let mut command = Command::new("sleep");
383            command.arg("60");
384            command
385        };
386        let spawned_pid = Arc::new(AtomicU32::new(0));
387
388        assert!(super::spawn_for_test(command, Arc::clone(&spawned_pid)).is_err());
389        let pid = spawned_pid.load(Ordering::SeqCst);
390        assert_ne!(pid, 0, "test must fail after the OS process was spawned");
391        let process = super::AdoptedProcess::new(std::num::NonZeroU32::new(pid).unwrap());
392        let deadline = Instant::now() + Duration::from_secs(5);
393        while process.is_running().unwrap_or(false) {
394            assert!(
395                Instant::now() < deadline,
396                "partially spawned process leaked"
397            );
398            thread::sleep(Duration::from_millis(20));
399        }
400    }
401}
rightkit_process/owned.rs

rightkit_process/
owned.rs
