if available use a Child's pidfd for kill/wait

the8472 · the8472 · commit 12efa53b19aa · 2023-11-16T02:05:37.000+01:00
diff --git a/library/std/src/os/linux/process.rs b/library/std/src/os/linux/process.rs
@@ -152,6 +152,12 @@ pub trait CommandExt: Sealed {
     /// in a guaranteed race-free manner (e.g. if the `clone3` system call
     /// is supported). Otherwise, [`pidfd`] will return an error.
     ///
+    /// If a pidfd has been successfully created and not been taken from the `Child`
+    /// then calls to `kill()`, `wait()` and `try_wait()` will use the pidfd
+    /// instead of the pid. This can prevent pid recycling races, e.g.
+    /// those  caused by rogue libraries in the same process prematurely reaping
+    /// zombie children via `waitpid(-1, ...)` calls.
+    ///
     /// [`Command`]: process::Command
     /// [`Child`]: process::Child
     /// [`pidfd`]: fn@ChildExt::pidfd
diff --git a/library/std/src/sys/unix/process/process_unix.rs b/library/std/src/sys/unix/process/process_unix.rs
@@ -9,6 +9,8 @@ use core::ffi::NonZero_c_int;
 
 #[cfg(target_os = "linux")]
 use crate::os::linux::process::PidFd;
+#[cfg(target_os = "linux")]
+use crate::os::unix::io::AsRawFd;
 
 #[cfg(any(
     target_os = "macos",
@@ -816,17 +818,41 @@ impl Process {
         // and used for another process, and we probably shouldn't be killing
         // random processes, so return Ok because the process has exited already.
         if self.status.is_some() {
-            Ok(())
-        } else {
-            cvt(unsafe { libc::kill(self.pid, libc::SIGKILL) }).map(drop)
+            return Ok(());
+        }
+        #[cfg(target_os = "linux")]
+        if let Some(pid_fd) = self.pidfd.as_ref() {
+            // pidfd_send_signal predates pidfd_open. so if we were able to get an fd then sending signals will work too
+            return cvt(unsafe {
+                libc::syscall(
+                    libc::SYS_pidfd_send_signal,
+                    pid_fd.as_raw_fd(),
+                    libc::SIGKILL,
+                    crate::ptr::null::<()>(),
+                    0,
+                )
+            })
+            .map(drop);
         }
+        cvt(unsafe { libc::kill(self.pid, libc::SIGKILL) }).map(drop)
     }
 
     pub fn wait(&mut self) -> io::Result<ExitStatus> {
         use crate::sys::cvt_r;
         if let Some(status) = self.status {
             return Ok(status);
         }
+        #[cfg(target_os = "linux")]
+        if let Some(pid_fd) = self.pidfd.as_ref() {
+            let mut siginfo: libc::siginfo_t = unsafe { crate::mem::zeroed() };
+
+            cvt_r(|| unsafe {
+                libc::waitid(libc::P_PIDFD, pid_fd.as_raw_fd() as u32, &mut siginfo, libc::WEXITED)
+            })?;
+            let status = ExitStatus::from_waitid_siginfo(siginfo);
+            self.status = Some(status);
+            return Ok(status);
+        }
         let mut status = 0 as c_int;
         cvt_r(|| unsafe { libc::waitpid(self.pid, &mut status, 0) })?;
         self.status = Some(ExitStatus::new(status));
@@ -837,6 +863,25 @@ impl Process {
         if let Some(status) = self.status {
             return Ok(Some(status));
         }
+        #[cfg(target_os = "linux")]
+        if let Some(pid_fd) = self.pidfd.as_ref() {
+            let mut siginfo: libc::siginfo_t = unsafe { crate::mem::zeroed() };
+
+            cvt(unsafe {
+                libc::waitid(
+                    libc::P_PIDFD,
+                    pid_fd.as_raw_fd() as u32,
+                    &mut siginfo,
+                    libc::WEXITED | libc::WNOHANG,
+                )
+            })?;
+            if unsafe { siginfo.si_pid() } == 0 {
+                return Ok(None);
+            }
+            let status = ExitStatus::from_waitid_siginfo(siginfo);
+            self.status = Some(status);
+            return Ok(Some(status));
+        }
         let mut status = 0 as c_int;
         let pid = cvt(unsafe { libc::waitpid(self.pid, &mut status, libc::WNOHANG) })?;
         if pid == 0 {
@@ -866,6 +911,20 @@ impl ExitStatus {
         ExitStatus(status)
     }
 
+    #[cfg(target_os = "linux")]
+    pub fn from_waitid_siginfo(siginfo: libc::siginfo_t) -> ExitStatus {
+        let status = unsafe { siginfo.si_status() };
+
+        match siginfo.si_code {
+            libc::CLD_EXITED => ExitStatus((status & 0xff) << 8),
+            libc::CLD_KILLED => ExitStatus(status),
+            libc::CLD_DUMPED => ExitStatus(status | 0x80),
+            libc::CLD_CONTINUED => ExitStatus(0xffff),
+            libc::CLD_STOPPED | libc::CLD_TRAPPED => ExitStatus(((status & 0xff) << 8) | 0x7f),
+            _ => unreachable!("waitid() should only return the above codes"),
+        }
+    }
+
     fn exited(&self) -> bool {
         libc::WIFEXITED(self.0)
     }
diff --git a/library/std/src/sys/unix/process/process_unix/tests.rs b/library/std/src/sys/unix/process/process_unix/tests.rs
@@ -64,6 +64,7 @@ fn test_command_fork_no_unwind() {
 #[test]
 #[cfg(target_os = "linux")]
 fn test_command_pidfd() {
+    use crate::assert_matches::assert_matches;
     use crate::os::fd::{AsRawFd, RawFd};
     use crate::os::linux::process::{ChildExt, CommandExt};
     use crate::process::Command;
@@ -78,7 +79,7 @@ fn test_command_pidfd() {
     };
 
     // always exercise creation attempts
-    let child = Command::new("echo").create_pidfd(true).spawn().unwrap();
+    let mut child = Command::new("false").create_pidfd(true).spawn().unwrap();
 
     // but only check if we know that the kernel supports pidfds
     if pidfd_open_available {
@@ -88,4 +89,12 @@ fn test_command_pidfd() {
         let flags = super::cvt(unsafe { libc::fcntl(pidfd.as_raw_fd(), libc::F_GETFD) }).unwrap();
         assert!(flags & libc::FD_CLOEXEC != 0);
     }
+    let status = child.wait().expect("error waiting on pidfd");
+    assert_eq!(status.code(), Some(1));
+
+    let mut child = Command::new("sleep").arg("1000").create_pidfd(true).spawn().unwrap();
+    assert_matches!(child.try_wait(), Ok(None));
+    child.kill().expect("failed to kill child");
+    let status = child.wait().expect("error waiting on pidfd");
+    assert_eq!(status.signal(), Some(libc::SIGKILL));
 }
