Auto merge of #124101 - the8472:pidfd-methods, r=cuviper

bors · bors · commit 10e1f5d21218 · 2024-06-22T03:35:52.000Z
Add PidFd::{kill, wait, try_wait} #117957 changed `Child` kill/wait/try_wait to use its pidfd instead of the pid, when one is available. This PR extracts those implementations and makes them available on `PidFd` directly. The `PidFd` implementations differ significantly from the corresponding `Child` methods: * the methods can be called after the child has been reaped, which will result in an error but will be safe. This state is not observable in `Child` unless something stole the zombie child * the `ExitStatus` is not kept, meaning that only the first time a wait succeeds it will be returned * `wait` does not close stdin * `wait` only requires `&self` instead of `&mut self` since there is no state to maintain and subsequent calls are safe Tracking issue: #82971
diff --git a/library/std/src/os/linux/process.rs b/library/std/src/os/linux/process.rs
@@ -6,20 +6,20 @@
 
 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.
 ///
 /// A `PidFd` can be obtained by setting the corresponding option on [`Command`]
 /// with [`create_pidfd`]. Subsequently, the created pidfd can be retrieved
-/// from the [`Child`] by calling [`pidfd`] or [`take_pidfd`].
+/// from the [`Child`] by calling [`pidfd`] or [`into_pidfd`].
 ///
 /// Example:
 /// ```no_run
@@ -33,7 +33,7 @@ struct FileDesc;
 ///     .expect("Failed to spawn child");
 ///
 /// let pidfd = child
-///     .take_pidfd()
+///     .into_pidfd()
 ///     .expect("Failed to retrieve pidfd");
 ///
 /// // The file descriptor will be closed when `pidfd` is dropped.
@@ -44,66 +44,101 @@ struct FileDesc;
 /// [`create_pidfd`]: CommandExt::create_pidfd
 /// [`Child`]: process::Child
 /// [`pidfd`]: fn@ChildExt::pidfd
-/// [`take_pidfd`]: ChildExt::take_pidfd
+/// [`into_pidfd`]: ChildExt::into_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()
     }
 }
 
@@ -124,18 +159,26 @@ pub trait ChildExt: Sealed {
     /// [`Child`]: process::Child
     fn pidfd(&self) -> Result<&PidFd>;
 
-    /// Takes ownership of the [`PidFd`] created for this [`Child`], if available.
+    /// Returns the [`PidFd`] created for this [`Child`], if available.
+    /// Otherwise self is returned.
     ///
     /// A pidfd will only be available if its creation was requested with
     /// [`create_pidfd`] when the corresponding [`Command`] was created.
     ///
+    /// Taking ownership of the PidFd consumes the Child to avoid pid reuse
+    /// races. Use [`pidfd`] and [`BorrowedFd::try_clone_to_owned`] if
+    /// you don't want to disassemble the Child yet.
+    ///
     /// Even if requested, a pidfd may not be available due to an older
     /// version of Linux being in use, or if some other error occurred.
     ///
     /// [`Command`]: process::Command
     /// [`create_pidfd`]: CommandExt::create_pidfd
+    /// [`pidfd`]: ChildExt::pidfd
     /// [`Child`]: process::Child
-    fn take_pidfd(&mut self) -> Result<PidFd>;
+    fn into_pidfd(self) -> crate::result::Result<PidFd, Self>
+    where
+        Self: Sized;
 }
 
 /// Os-specific extensions for [`Command`]
@@ -146,7 +189,7 @@ pub trait CommandExt: Sealed {
     /// spawned by this [`Command`].
     /// By default, no pidfd will be created.
     ///
-    /// The pidfd can be retrieved from the child with [`pidfd`] or [`take_pidfd`].
+    /// The pidfd can be retrieved from the child with [`pidfd`] or [`into_pidfd`].
     ///
     /// A pidfd will only be created if it is possible to do so
     /// in a guaranteed race-free manner. Otherwise, [`pidfd`] will return an error.
@@ -160,7 +203,7 @@ pub trait CommandExt: Sealed {
     /// [`Command`]: process::Command
     /// [`Child`]: process::Child
     /// [`pidfd`]: fn@ChildExt::pidfd
-    /// [`take_pidfd`]: ChildExt::take_pidfd
+    /// [`into_pidfd`]: ChildExt::into_pidfd
     fn create_pidfd(&mut self, val: bool) -> &mut process::Command;
 }
 
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,87 @@
+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 child = Command::new("sleep")
+        .arg("1000")
+        .create_pidfd(true)
+        .spawn()
+        .expect("executing 'sleep' failed");
+
+    let fd = child.into_pidfd().unwrap();
+
+    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
