Add PidFd::{kill, wait, try_wait}

the8472 · the8472 · commit 0787c7308cb0 · 2024-06-22T00:46:55.000+02:00
diff --git a/library/std/src/os/linux/process.rs b/library/std/src/os/linux/process.rs
@@ -6,14 +6,14 @@
 
 use crate::io::Result;
 use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
-use crate::process;
+use crate::process::{self, ExitStatus};
 use crate::sealed::Sealed;
 #[cfg(not(doc))]
-use crate::sys::fd::FileDesc;
+use crate::sys::{fd::FileDesc, linux::pidfd::PidFd as InnerPidFd};
 use crate::sys_common::{AsInner, AsInnerMut, FromInner, IntoInner};
 
 #[cfg(doc)]
-struct FileDesc;
+struct InnerPidFd;
 
 /// This type represents a file descriptor that refers to a process.
 ///
@@ -47,63 +47,98 @@ struct FileDesc;
 /// [`take_pidfd`]: ChildExt::take_pidfd
 /// [`pidfd_open(2)`]: https://man7.org/linux/man-pages/man2/pidfd_open.2.html
 #[derive(Debug)]
+#[repr(transparent)]
 pub struct PidFd {
-    inner: FileDesc,
+    inner: InnerPidFd,
 }
 
-impl AsInner<FileDesc> for PidFd {
+impl PidFd {
+    /// Forces the child process to exit.
+    ///
+    /// Unlike [`Child::kill`] it is possible to attempt to kill
+    /// reaped children since PidFd does not suffer from pid recycling
+    /// races. But doing so will return an Error.
+    ///
+    /// [`Child::kill`]: process::Child::kill
+    pub fn kill(&self) -> Result<()> {
+        self.inner.kill()
+    }
+
+    /// Waits for the child to exit completely, returning the status that it exited with.
+    ///
+    /// Unlike [`Child::wait`] it does not ensure that the stdin handle is closed.
+    /// Additionally it will not return an `ExitStatus` if the child
+    /// has already been reaped. Instead an error will be returned.
+    ///
+    /// [`Child::wait`]: process::Child::wait
+    pub fn wait(&self) -> Result<ExitStatus> {
+        self.inner.wait().map(FromInner::from_inner)
+    }
+
+    /// Attempts to collect the exit status of the child if it has already exited.
+    ///
+    /// Unlike [`Child::try_wait`] this method will return an Error
+    /// if the child has already been reaped.
+    ///
+    /// [`Child::try_wait`]: process::Child::try_wait
+    pub fn try_wait(&self) -> Result<Option<ExitStatus>> {
+        Ok(self.inner.try_wait()?.map(FromInner::from_inner))
+    }
+}
+
+impl AsInner<InnerPidFd> for PidFd {
     #[inline]
-    fn as_inner(&self) -> &FileDesc {
+    fn as_inner(&self) -> &InnerPidFd {
         &self.inner
     }
 }
 
-impl FromInner<FileDesc> for PidFd {
-    fn from_inner(inner: FileDesc) -> PidFd {
+impl FromInner<InnerPidFd> for PidFd {
+    fn from_inner(inner: InnerPidFd) -> PidFd {
         PidFd { inner }
     }
 }
 
-impl IntoInner<FileDesc> for PidFd {
-    fn into_inner(self) -> FileDesc {
+impl IntoInner<InnerPidFd> for PidFd {
+    fn into_inner(self) -> InnerPidFd {
         self.inner
     }
 }
 
 impl AsRawFd for PidFd {
     #[inline]
     fn as_raw_fd(&self) -> RawFd {
-        self.as_inner().as_raw_fd()
+        self.as_inner().as_inner().as_raw_fd()
     }
 }
 
 impl FromRawFd for PidFd {
     unsafe fn from_raw_fd(fd: RawFd) -> Self {
-        Self::from_inner(FileDesc::from_raw_fd(fd))
+        Self::from_inner(InnerPidFd::from_raw_fd(fd))
     }
 }
 
 impl IntoRawFd for PidFd {
     fn into_raw_fd(self) -> RawFd {
-        self.into_inner().into_raw_fd()
+        self.into_inner().into_inner().into_raw_fd()
     }
 }
 
 impl AsFd for PidFd {
     fn as_fd(&self) -> BorrowedFd<'_> {
-        self.as_inner().as_fd()
+        self.as_inner().as_inner().as_fd()
     }
 }
 
 impl From<OwnedFd> for PidFd {
     fn from(fd: OwnedFd) -> Self {
-        Self::from_inner(FileDesc::from_inner(fd))
+        Self::from_inner(InnerPidFd::from_inner(FileDesc::from_inner(fd)))
     }
 }
 
 impl From<PidFd> for OwnedFd {
     fn from(pid_fd: PidFd) -> Self {
-        pid_fd.into_inner().into_inner()
+        pid_fd.into_inner().into_inner().into_inner()
     }
 }
 
diff --git a/library/std/src/sys/pal/unix/linux/mod.rs b/library/std/src/sys/pal/unix/linux/mod.rs
@@ -0,0 +1 @@
+pub mod pidfd;
diff --git a/library/std/src/sys/pal/unix/linux/pidfd.rs b/library/std/src/sys/pal/unix/linux/pidfd.rs
@@ -0,0 +1,76 @@
+use crate::io;
+use crate::os::fd::{AsRawFd, FromRawFd, RawFd};
+use crate::sys::cvt;
+use crate::sys::pal::unix::fd::FileDesc;
+use crate::sys::process::ExitStatus;
+use crate::sys_common::{AsInner, FromInner, IntoInner};
+
+#[cfg(test)]
+mod tests;
+
+#[derive(Debug)]
+pub(crate) struct PidFd(FileDesc);
+
+impl PidFd {
+    pub fn kill(&self) -> io::Result<()> {
+        return cvt(unsafe {
+            libc::syscall(
+                libc::SYS_pidfd_send_signal,
+                self.0.as_raw_fd(),
+                libc::SIGKILL,
+                crate::ptr::null::<()>(),
+                0,
+            )
+        })
+        .map(drop);
+    }
+
+    pub fn wait(&self) -> io::Result<ExitStatus> {
+        let mut siginfo: libc::siginfo_t = unsafe { crate::mem::zeroed() };
+        cvt(unsafe {
+            libc::waitid(libc::P_PIDFD, self.0.as_raw_fd() as u32, &mut siginfo, libc::WEXITED)
+        })?;
+        return Ok(ExitStatus::from_waitid_siginfo(siginfo));
+    }
+
+    pub fn try_wait(&self) -> io::Result<Option<ExitStatus>> {
+        let mut siginfo: libc::siginfo_t = unsafe { crate::mem::zeroed() };
+
+        cvt(unsafe {
+            libc::waitid(
+                libc::P_PIDFD,
+                self.0.as_raw_fd() as u32,
+                &mut siginfo,
+                libc::WEXITED | libc::WNOHANG,
+            )
+        })?;
+        if unsafe { siginfo.si_pid() } == 0 {
+            return Ok(None);
+        }
+        return Ok(Some(ExitStatus::from_waitid_siginfo(siginfo)));
+    }
+}
+
+impl AsInner<FileDesc> for PidFd {
+    fn as_inner(&self) -> &FileDesc {
+        &self.0
+    }
+}
+
+impl IntoInner<FileDesc> for PidFd {
+    fn into_inner(self) -> FileDesc {
+        self.0
+    }
+}
+
+impl FromInner<FileDesc> for PidFd {
+    fn from_inner(inner: FileDesc) -> Self {
+        Self(inner)
+    }
+}
+
+impl FromRawFd for PidFd {
+    unsafe fn from_raw_fd(fd: RawFd) -> Self {
+        Self(FileDesc::from_raw_fd(fd))
+    }
+}
diff --git a/library/std/src/sys/pal/unix/linux/pidfd/tests.rs b/library/std/src/sys/pal/unix/linux/pidfd/tests.rs
@@ -0,0 +1,88 @@
+use crate::assert_matches::assert_matches;
+use crate::os::fd::{AsRawFd, RawFd};
+use crate::os::linux::process::{ChildExt, CommandExt};
+use crate::os::unix::process::ExitStatusExt;
+use crate::process::Command;
+
+#[test]
+fn test_command_pidfd() {
+    let pidfd_open_available = probe_pidfd_support();
+
+    // always exercise creation attempts
+    let mut child = Command::new("false").create_pidfd(true).spawn().unwrap();
+
+    // but only check if we know that the kernel supports pidfds.
+    // We don't assert the precise value, since the standard library
+    // might have opened other file descriptors before our code runs.
+    if pidfd_open_available {
+        assert!(child.pidfd().is_ok());
+    }
+    if let Ok(pidfd) = child.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));
+
+    let _ = Command::new("echo")
+        .create_pidfd(false)
+        .spawn()
+        .unwrap()
+        .pidfd()
+        .expect_err("pidfd should not have been created when create_pid(false) is set");
+
+    let _ = Command::new("echo")
+        .spawn()
+        .unwrap()
+        .pidfd()
+        .expect_err("pidfd should not have been created");
+}
+
+#[test]
+fn test_pidfd() {
+    if !probe_pidfd_support() {
+        return;
+    }
+
+    let mut child = Command::new("sleep")
+        .arg("1000")
+        .create_pidfd(true)
+        .spawn()
+        .expect("executing 'sleep' failed");
+
+    let fd = child.take_pidfd().unwrap();
+    drop(child);
+
+    assert_matches!(fd.try_wait(), Ok(None));
+    fd.kill().expect("kill failed");
+    fd.kill().expect("sending kill twice failed");
+    let status = fd.wait().expect("1st wait failed");
+    assert_eq!(status.signal(), Some(libc::SIGKILL));
+
+    // Trying to wait again for a reaped child is safe since there's no pid-recycling race.
+    // But doing so will return an error.
+    let res = fd.wait();
+    assert_matches!(res, Err(e) if e.raw_os_error() == Some(libc::ECHILD));
+
+    // Ditto for additional attempts to kill an already-dead child.
+    let res = fd.kill();
+    assert_matches!(res, Err(e) if e.raw_os_error() == Some(libc::ESRCH));
+}
+
+fn probe_pidfd_support() -> bool {
+    // pidfds require the pidfd_open syscall
+    let our_pid = crate::process::id();
+    let pidfd = unsafe { libc::syscall(libc::SYS_pidfd_open, our_pid, 0) };
+    if pidfd >= 0 {
+        unsafe { libc::close(pidfd as RawFd) };
+        true
+    } else {
+        false
+    }
+}
diff --git a/library/std/src/sys/pal/unix/mod.rs b/library/std/src/sys/pal/unix/mod.rs
@@ -20,6 +20,8 @@ pub mod io;
 pub mod kernel_copy;
 #[cfg(target_os = "l4re")]
 mod l4re;
+#[cfg(target_os = "linux")]
+pub mod linux;
 #[cfg(not(target_os = "l4re"))]
 pub mod net;
 #[cfg(target_os = "l4re")]
diff --git a/library/std/src/sys/pal/unix/process/process_unix.rs b/library/std/src/sys/pal/unix/process/process_unix.rs
diff --git a/library/std/src/sys/pal/unix/process/process_unix/tests.rs b/library/std/src/sys/pal/unix/process/process_unix/tests.rs
