Add windows::process::Command::inherit_handles

quark-zju · quark-zju · commit 5ddfa2e29ae3 · 2020-08-15T15:27:42.000-07:00
The API limits handles to inherit. The implementation follows the "Programmatically controlling which handles are inherited by new processes in Win32" blog from Microsoft: https://devblogs.microsoft.com/oldnewthing/20111216-00/?p=8873 If the API is not called, the old behavior of inheriting all inheritable handles is not changed.
diff --git a/library/std/src/process.rs b/library/std/src/process.rs
@@ -2104,6 +2104,93 @@ mod tests {
         assert!(events > 0);
     }
 
+    #[test]
+    #[cfg(windows)]
+    fn test_inherit_handles() {
+        use crate::io;
+        use crate::os::windows::process::CommandExt;
+        use crate::sys::pipe::anon_pipe;
+
+        // Count handles of a child process using PowerShell.
+        fn count_child_handles(f: impl FnOnce(&mut Command) -> &mut Command) -> io::Result<usize> {
+            let mut command = Command::new("powershell");
+            let command = command.args(&[
+                "-Command",
+                "Get-Process -Id $PID | Format-Table -HideTableHeaders Handles",
+            ]);
+            let out = f(command).output()?.stdout;
+            String::from_utf8_lossy(&out)
+                .trim()
+                .parse::<usize>()
+                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
+        }
+
+        // The exact count is unstable because tests run in other threads might
+        // create inheritable handles, and PowerShell is a GC runtime.
+        // Increase epsilon if the test becomes flaky.
+        let epsilon = 25;
+        let base_count = count_child_handles(|c| c).unwrap();
+
+        // Create `n` inheritable pipes.
+        let n = epsilon * 6;
+        let pipes: Vec<_> = (1..n).map(|_| anon_pipe(true, true).unwrap()).collect();
+
+        // Without `inherit_handles`, all pipes are inherited.
+        let inherit_count = count_child_handles(|c| c).unwrap();
+        assert!(
+            inherit_count > base_count + n - epsilon,
+            "pipes do not seem to be inherited (base: {}, inherit: {}, n: {})",
+            base_count,
+            inherit_count,
+            n
+        );
+
+        // With `inherit_handles`, none of the pipes are inherited.
+        let inherit_count = count_child_handles(|c| c.inherit_handles(vec![])).unwrap();
+        assert!(
+            inherit_count < base_count + epsilon,
+            "pipe inheritance does not seem to be limited (base: {}, inherit: {}, n: {})",
+            base_count,
+            inherit_count,
+            n
+        );
+
+        // Use `inherit_handles` to inherit half of the pipes.
+        let half = n / 2;
+        let handles: Vec<_> = pipes.iter().take(half).map(|p| p.theirs.handle().raw()).collect();
+        let inherit_count = count_child_handles(move |c| c.inherit_handles(handles)).unwrap();
+        assert!(
+            inherit_count > base_count + half - epsilon,
+            "pipe inheritance seems over-limited (base: {}, inherit: {}, n: {}, half: {})",
+            base_count,
+            inherit_count,
+            n,
+            half
+        );
+        assert!(
+            inherit_count < base_count + half + epsilon,
+            "pipe inheritance seems under-limited (base: {}, inherit: {}, n: {}, half: {})",
+            base_count,
+            inherit_count,
+            n,
+            half
+        );
+
+        // Call `inherit_handles` multiple times to inherit all pipes.
+        let handles: Vec<_> = pipes.iter().map(|p| p.theirs.handle().raw()).collect();
+        let inherit_count = count_child_handles(move |c| {
+            handles.into_iter().fold(c, |c, h| c.inherit_handles(vec![h]))
+        })
+        .unwrap();
+        assert!(
+            inherit_count > base_count + n - epsilon,
+            "pipe inheritance seems over-limited (base: {}, inherit: {}, n: {})",
+            base_count,
+            inherit_count,
+            n,
+        );
+    }
+
     #[test]
     fn test_command_implements_send_sync() {
         fn take_send_sync_type<T: Send + Sync>(_: T) {}
diff --git a/library/std/src/sys/windows/c.rs b/library/std/src/sys/windows/c.rs
@@ -13,6 +13,7 @@ pub use self::EXCEPTION_DISPOSITION::*;
 pub use self::FILE_INFO_BY_HANDLE_CLASS::*;
 
 pub type DWORD = c_ulong;
+pub type DWORD_PTR = ULONG_PTR;
 pub type HANDLE = LPVOID;
 pub type HINSTANCE = HANDLE;
 pub type HMODULE = HINSTANCE;
@@ -39,6 +40,7 @@ pub type LPCWSTR = *const WCHAR;
 pub type LPDWORD = *mut DWORD;
 pub type LPHANDLE = *mut HANDLE;
 pub type LPOVERLAPPED = *mut OVERLAPPED;
+pub type LPPROC_THREAD_ATTRIBUTE_LIST = *mut PROC_THREAD_ATTRIBUTE_LIST;
 pub type LPPROCESS_INFORMATION = *mut PROCESS_INFORMATION;
 pub type LPSECURITY_ATTRIBUTES = *mut SECURITY_ATTRIBUTES;
 pub type LPSTARTUPINFO = *mut STARTUPINFO;
@@ -56,7 +58,9 @@ pub type LPWSAOVERLAPPED_COMPLETION_ROUTINE = *mut c_void;
 
 pub type PCONDITION_VARIABLE = *mut CONDITION_VARIABLE;
 pub type PLARGE_INTEGER = *mut c_longlong;
+pub type PSIZE_T = *mut ULONG_PTR;
 pub type PSRWLOCK = *mut SRWLOCK;
+pub type PVOID = *mut c_void;
 
 pub type SOCKET = crate::os::windows::raw::SOCKET;
 pub type socklen_t = c_int;
@@ -101,6 +105,8 @@ pub const SECURITY_SQOS_PRESENT: DWORD = 0x00100000;
 
 pub const FIONBIO: c_ulong = 0x8004667e;
 
+pub const PROC_THREAD_ATTRIBUTE_HANDLE_LIST: DWORD_PTR = 0x00020002;
+
 #[repr(C)]
 #[derive(Copy)]
 pub struct WIN32_FIND_DATAW {
@@ -217,6 +223,7 @@ pub const CONDITION_VARIABLE_INIT: CONDITION_VARIABLE = CONDITION_VARIABLE { ptr
 pub const SRWLOCK_INIT: SRWLOCK = SRWLOCK { ptr: ptr::null_mut() };
 
 pub const DETACHED_PROCESS: DWORD = 0x00000008;
+pub const EXTENDED_STARTUPINFO_PRESENT: DWORD = 0x00080000;
 pub const CREATE_NEW_PROCESS_GROUP: DWORD = 0x00000200;
 pub const CREATE_UNICODE_ENVIRONMENT: DWORD = 0x00000400;
 pub const STARTF_USESTDHANDLES: DWORD = 0x00000100;
@@ -483,6 +490,11 @@ pub struct SECURITY_ATTRIBUTES {
     pub bInheritHandle: BOOL,
 }
 
+#[repr(C)]
+pub struct PROC_THREAD_ATTRIBUTE_LIST {
+    pub dummy: *mut c_void,
+}
+
 #[repr(C)]
 pub struct PROCESS_INFORMATION {
     pub hProcess: HANDLE,
@@ -513,6 +525,12 @@ pub struct STARTUPINFO {
     pub hStdError: HANDLE,
 }
 
+#[repr(C)]
+pub struct STARTUPINFOEX {
+    pub StartupInfo: STARTUPINFO,
+    pub lpAttributeList: LPPROC_THREAD_ATTRIBUTE_LIST,
+}
+
 #[repr(C)]
 pub struct SOCKADDR {
     pub sa_family: ADDRESS_FAMILY,
@@ -887,6 +905,23 @@ extern "system" {
         lpUsedDefaultChar: LPBOOL,
     ) -> c_int;
 
+    pub fn InitializeProcThreadAttributeList(
+        lpAttributeList: LPPROC_THREAD_ATTRIBUTE_LIST,
+        dwAttributeCount: DWORD,
+        dwFlags: DWORD,
+        lpSize: PSIZE_T,
+    ) -> BOOL;
+    pub fn UpdateProcThreadAttribute(
+        lpAttributeList: LPPROC_THREAD_ATTRIBUTE_LIST,
+        dwFlags: DWORD,
+        Attribute: DWORD_PTR,
+        lpValue: PVOID,
+        cbSize: SIZE_T,
+        lpPreviousValue: PVOID,
+        lpReturnSize: PSIZE_T,
+    ) -> BOOL;
+    pub fn DeleteProcThreadAttributeList(lpAttributeList: LPPROC_THREAD_ATTRIBUTE_LIST);
+
     pub fn closesocket(socket: SOCKET) -> c_int;
     pub fn recv(socket: SOCKET, buf: *mut c_void, len: c_int, flags: c_int) -> c_int;
     pub fn send(socket: SOCKET, buf: *const c_void, len: c_int, flags: c_int) -> c_int;
diff --git a/library/std/src/sys/windows/ext/process.rs b/library/std/src/sys/windows/ext/process.rs
@@ -102,6 +102,28 @@ pub trait CommandExt {
     /// [1]: https://docs.microsoft.com/en-us/windows/win32/procthread/process-creation-flags
     #[stable(feature = "windows_process_extensions", since = "1.16.0")]
     fn creation_flags(&mut self, flags: u32) -> &mut process::Command;
+
+    /// Specifies additional [inheritable handles][1] for `CreateProcess`.
+    ///
+    /// Handles must be created as inheritable and must not be pseudo
+    /// such as those returned by the `GetCurrentProcess`.
+    ///
+    /// Handles for stdio redirection are always inherited. Handles are
+    /// passed via `STARTUPINFOEX`. Process creation flags will be ORed
+    /// by `EXTENDED_STARTUPINFO_PRESENT` automatically.
+    ///
+    /// Calling this function multiple times is equivalent to calling
+    /// one time with a chained iterator.
+    ///
+    /// If this function is not called, all inheritable handles will be
+    /// inherited. Note this may change in the future.
+    ///
+    /// [1]: https://docs.microsoft.com/en-us/windows/win32/procthread/inheritance#inheriting-handles
+    #[unstable(feature = "windows_process_inherit_handles", issue = "none")]
+    fn inherit_handles(
+        &mut self,
+        handles: impl IntoIterator<Item = RawHandle>,
+    ) -> &mut process::Command;
 }
 
 #[stable(feature = "windows_process_extensions", since = "1.16.0")]
@@ -110,4 +132,12 @@ impl CommandExt for process::Command {
         self.as_inner_mut().creation_flags(flags);
         self
     }
+
+    fn inherit_handles(
+        &mut self,
+        handles: impl IntoIterator<Item = RawHandle>,
+    ) -> &mut process::Command {
+        self.as_inner_mut().inherit_handles(handles.into_iter().collect());
+        self
+    }
 }
diff --git a/library/std/src/sys/windows/process.rs b/library/std/src/sys/windows/process.rs
@@ -10,6 +10,7 @@ use crate::fs;
 use crate::io::{self, Error, ErrorKind};
 use crate::mem;
 use crate::os::windows::ffi::OsStrExt;
+use crate::os::windows::io::RawHandle;
 use crate::path::Path;
 use crate::ptr;
 use crate::sys::c;
@@ -72,11 +73,17 @@ pub struct Command {
     cwd: Option<OsString>,
     flags: u32,
     detach: bool, // not currently exposed in std::process
+    inherit_handles: Option<SendHandles>,
     stdin: Option<Stdio>,
     stdout: Option<Stdio>,
     stderr: Option<Stdio>,
 }
 
+struct SendHandles(Vec<RawHandle>);
+
+unsafe impl Send for SendHandles {}
+unsafe impl Sync for SendHandles {}
+
 pub enum Stdio {
     Inherit,
     Null,
@@ -94,6 +101,13 @@ struct DropGuard<'a> {
     lock: &'a Mutex,
 }
 
+#[derive(Default)]
+struct ProcAttributeList {
+    buf: Vec<u8>,
+    inherit_handles: Vec<RawHandle>,
+    initialized: bool,
+}
+
 impl Command {
     pub fn new(program: &OsStr) -> Command {
         Command {
@@ -103,6 +117,7 @@ impl Command {
             cwd: None,
             flags: 0,
             detach: false,
+            inherit_handles: None,
             stdin: None,
             stdout: None,
             stderr: None,
@@ -130,6 +145,12 @@ impl Command {
     pub fn creation_flags(&mut self, flags: u32) {
         self.flags = flags;
     }
+    pub fn inherit_handles(&mut self, handles: Vec<RawHandle>) {
+        match self.inherit_handles {
+            Some(ref mut list) => list.0.extend(handles),
+            None => self.inherit_handles = Some(SendHandles(handles)),
+        }
+    }
 
     pub fn spawn(
         &mut self,
@@ -156,7 +177,11 @@ impl Command {
             None
         });
 
-        let mut si = zeroed_startupinfo();
+        let mut si_ex = c::STARTUPINFOEX {
+            StartupInfo: zeroed_startupinfo(),
+            lpAttributeList: ptr::null_mut(),
+        };
+        let si = &mut si_ex.StartupInfo;
         si.cb = mem::size_of::<c::STARTUPINFO>() as c::DWORD;
         si.dwFlags = c::STARTF_USESTDHANDLES;
 
@@ -199,6 +224,20 @@ impl Command {
         si.hStdOutput = stdout.raw();
         si.hStdError = stderr.raw();
 
+        // Limit handles to inherit.
+        // See https://devblogs.microsoft.com/oldnewthing/20111216-00/?p=8873
+        let mut proc_attribute_list = ProcAttributeList::default();
+        if let Some(SendHandles(ref handles)) = self.inherit_handles {
+            // Unconditionally inherit stdio redirection handles.
+            let mut inherit_handles = handles.clone();
+            inherit_handles.extend_from_slice(&[stdin.raw(), stdout.raw(), stderr.raw()]);
+            proc_attribute_list.init(1)?;
+            proc_attribute_list.set_handle_list(inherit_handles)?;
+            si.cb = mem::size_of::<c::STARTUPINFOEX>() as c::DWORD;
+            si_ex.lpAttributeList = proc_attribute_list.ptr();
+            flags |= c::EXTENDED_STARTUPINFO_PRESENT;
+        }
+
         unsafe {
             cvt(c::CreateProcessW(
                 ptr::null(),
@@ -209,7 +248,7 @@ impl Command {
                 flags,
                 envp,
                 dirp,
-                &mut si,
+                si,
                 &mut pi,
             ))
         }?;
@@ -250,6 +289,60 @@ impl<'a> Drop for DropGuard<'a> {
     }
 }
 
+impl ProcAttributeList {
+    fn init(&mut self, attribute_count: c::DWORD) -> io::Result<()> {
+        // Allocate
+        let mut size: c::SIZE_T = 0;
+        cvt(unsafe {
+            c::InitializeProcThreadAttributeList(ptr::null_mut(), attribute_count, 0, &mut size)
+        })
+        .or_else(|e| match e.raw_os_error().unwrap_or_default() as c::DWORD {
+            c::ERROR_INSUFFICIENT_BUFFER => Ok(0),
+            _ => Err(e),
+        })?;
+        self.buf.resize(size as usize, 0);
+
+        // Initialize
+        cvt(unsafe {
+            c::InitializeProcThreadAttributeList(self.ptr(), attribute_count, 0, &mut size)
+        })?;
+        self.initialized = true;
+
+        Ok(())
+    }
+
+    fn set_handle_list(&mut self, handles: Vec<RawHandle>) -> io::Result<()> {
+        // Take ownership. The PROC_THREAD_ATTRIBUTE_LIST struct might
+        // refer to the handles using pointers.
+        self.inherit_handles = handles;
+        cvt(unsafe {
+            c::UpdateProcThreadAttribute(
+                self.ptr(),
+                0,
+                c::PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
+                self.inherit_handles.as_mut_ptr() as c::PVOID,
+                self.inherit_handles.len() * mem::size_of::<c::HANDLE>(),
+                ptr::null_mut(),
+                ptr::null_mut(),
+            )
+        })?;
+
+        Ok(())
+    }
+
+    fn ptr(&mut self) -> c::LPPROC_THREAD_ATTRIBUTE_LIST {
+        self.buf.as_mut_ptr() as _
+    }
+}
+
+impl Drop for ProcAttributeList {
+    fn drop(&mut self) {
+        if self.initialized {
+            unsafe { c::DeleteProcThreadAttributeList(self.ptr()) };
+        }
+    }
+}
+
 impl Stdio {
     fn to_handle(&self, stdio_id: c::DWORD, pipe: &mut Option<AnonPipe>) -> io::Result<Handle> {
         match *self {
