run-make-support: add unsafe set_aux_fd helper

Co-authored-by: Noa <[email protected]>
Co-authored-by: Oneirical <[email protected]>

Author: 3 people

src/tools/run-make-support/src/command.rs

@@ -151,6 +151,53 @@ impl Command {
         self
     }
 
+    /// Set an auxiliary stream passed to the process, besides the stdio streams.
+    ///
+    /// # Safety
+    ///
+    /// Use with caution! Ideally, only set one aux fd; if there are multiple, their old `fd` may
+    /// overlap with another's `new_fd`, and may break. The caller must make sure this is not the
+    /// case.
+    #[cfg(unix)]
+    pub unsafe fn set_aux_fd<F: Into<std::os::fd::OwnedFd>>(
+        &mut self,
+        new_fd: std::os::fd::RawFd,
+        fd: F,
+    ) -> &mut Self {
+        // NOTE: If more than 1 auxiliary file descriptor is needed, this function should be
+        // rewritten.
+
+        use std::os::fd::AsRawFd;
+        use std::os::unix::process::CommandExt;
+
+        let cvt = |x| if x == -1 { Err(std::io::Error::last_os_error()) } else { Ok(()) };
+
+        let fd = fd.into();
+        if fd.as_raw_fd() == new_fd {
+            // If the new file descriptor is already the same as fd, just turn off `FD_CLOEXEC`.
+
+            // SAFETY(io-safety): `fd` is already owned.
+            cvt(unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFD, 0) })
+                .expect("disabling CLOEXEC failed");
+            // The `pre_exec` function should be unconditionally set, since it captures `fd`, and
+            // this ensures that it stays open until the fork.
+        }
+        let pre_exec = move || {
+            if fd.as_raw_fd() != new_fd {
+                // SAFETY(io-safety): `new_fd` is not necessarily an unused fd. However, we're
+                // ensuring that `new_fd` will now refer to the same file descriptor as `fd`, which
+                // is safe as long as we manage the lifecycle of both descriptors correctly. This
+                // operation will replace the file descriptor referred to by `new_fd` with the one
+                // from `fd`, allowing for shared access to the same underlying file or resource.
+                cvt(unsafe { libc::dup2(fd.as_raw_fd(), new_fd) })?;
+            }
+            Ok(())
+        };
+        // SAFETY(pre-exec-safe): `dup2` is pre-exec-safe.
+        unsafe { self.cmd.pre_exec(pre_exec) };
+        self
+    }
+
     /// Run the constructed command and assert that it is successfully run.
     ///
     /// By default, std{in,out,err} are [`Stdio::piped()`].

src/tools/run-make-support/src/macros.rs

@@ -104,6 +104,23 @@ macro_rules! impl_common_helpers {
                 self
             }
 
+            /// Set an auxiliary stream passed to the process, besides the stdio streams.
+            ///
+            /// # Safety
+            ///
+            /// Use with caution! Ideally, only set one aux fd; if there are multiple, their old
+            /// `fd` may overlap with another's `new_fd`, and may break. The caller must make sure
+            /// this is not the case.
+            #[cfg(unix)]
+            pub unsafe fn set_aux_fd<F: Into<std::os::fd::OwnedFd>>(
+                &mut self,
+                new_fd: std::os::fd::RawFd,
+                fd: F,
+            ) -> &mut Self {
+                self.cmd.set_aux_fd(new_fd, fd);
+                self
+            }
+
             /// Run the constructed command and assert that it is successfully run.
             #[track_caller]
             pub fn run(&mut self) -> crate::command::CompletedProcess {