add windows ext method for proc thread attributes

TyPR124 · TyPR124 · commit 0475e2fea052 · 2022-03-06T22:16:49.000-05:00
Added doc comments, don't take too many attrs

kill the cmds in test
diff --git a/library/std/src/os/windows/process.rs b/library/std/src/os/windows/process.rs
@@ -4,7 +4,7 @@
 
 #![stable(feature = "process_extensions", since = "1.2.0")]
 
-use crate::ffi::OsStr;
+use crate::ffi::{c_void, OsStr};
 use crate::os::windows::io::{
     AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle, OwnedHandle, RawHandle,
 };
@@ -161,6 +161,23 @@ pub trait CommandExt: Sealed {
     /// `CommandLineToArgvW` escaping rules.
     #[unstable(feature = "windows_process_extensions_raw_arg", issue = "29494")]
     fn raw_arg<S: AsRef<OsStr>>(&mut self, text_to_append_as_is: S) -> &mut process::Command;
+
+    /// Attach the specified proc thread attribute. Multiple attributes may be attached.
+    ///
+    /// # Safety
+    ///
+    /// - The data pointed to by `value` pointer must not be moved or aliased for the entire lifetime of
+    ///   the `Command`.
+    /// - The data pointed to by `value` pointer must be initialized.
+    /// - `size` must not exceed the size of the object pointed to by the `value` pointer.
+    /// - It must be safe to read the data pointed to by `value` from another thread.
+    #[unstable(feature = "windows_process_extensions_proc_thread_attributes", issue = "none")]
+    unsafe fn process_thread_attribute(
+        &mut self,
+        attribute: usize,
+        value: *mut c_void,
+        size: usize,
+    ) -> &mut process::Command;
 }
 
 #[stable(feature = "windows_process_extensions", since = "1.16.0")]
@@ -179,4 +196,13 @@ impl CommandExt for process::Command {
         self.as_inner_mut().raw_arg(raw_text.as_ref());
         self
     }
+    unsafe fn process_thread_attribute(
+        &mut self,
+        attribute: usize,
+        value: *mut c_void,
+        size: usize,
+    ) -> &mut process::Command {
+        self.as_inner_mut().process_thread_attribute(attribute, value, size);
+        self
+    }
 }
diff --git a/library/std/src/process/tests.rs b/library/std/src/process/tests.rs
@@ -405,6 +405,42 @@ fn test_creation_flags() {
     }
     assert!(events > 0);
 }
+/// Tests proc thread attributes by spawning a process with a custom parent process,
+/// then comparing the parent process ID with the expected parent process ID.
+#[test]
+#[cfg(windows)]
+fn test_proc_thread_attributes() {
+    use crate::os::windows::io::AsRawHandle;
+    use crate::os::windows::process::CommandExt;
+    use crate::sys::c::{DWORD_PTR, HANDLE};
+    const PROC_THREAD_ATTRIBUTE_PARENT_PROCESS: DWORD_PTR = 0x00020000;
+    let mut parent = Command::new("cmd.exe").spawn().unwrap();
+    let mut parent_handle: HANDLE = parent.as_raw_handle();
+
+    let mut child_cmd = Command::new("cmd.exe");
+    unsafe {
+        child_cmd.process_thread_attribute(
+            PROC_THREAD_ATTRIBUTE_PARENT_PROCESS,
+            &mut parent_handle as *mut _ as *mut _,
+            crate::mem::size_of::<HANDLE>(),
+        );
+    }
+    let mut child = child_cmd.spawn().unwrap();
+
+    let wmic = format!("wmic process where processid={} get parentprocessid", child.id());
+    assert_eq!(
+        parent.id(),
+        String::from_utf8(Command::new("cmd.exe").args(["/c", &wmic]).output().unwrap().stdout)
+            .unwrap()
+            .lines()
+            .skip(1)
+            .map(|line| line.trim().parse().unwrap())
+            .next()
+            .unwrap()
+    );
+    child.kill().unwrap();
+    parent.kill().unwrap();
+}
 
 #[test]
 fn test_command_implements_send_sync() {
diff --git a/library/std/src/sys/windows/c.rs b/library/std/src/sys/windows/c.rs
@@ -51,6 +51,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 c_void;
 pub type LPPROCESS_INFORMATION = *mut PROCESS_INFORMATION;
 pub type LPSECURITY_ATTRIBUTES = *mut SECURITY_ATTRIBUTES;
 pub type LPSTARTUPINFO = *mut STARTUPINFO;
@@ -68,6 +69,7 @@ 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 usize;
 pub type PSRWLOCK = *mut SRWLOCK;
 
 pub type SOCKET = crate::os::windows::raw::SOCKET;
@@ -196,6 +198,7 @@ pub const SRWLOCK_INIT: SRWLOCK = SRWLOCK { ptr: ptr::null_mut() };
 pub const DETACHED_PROCESS: DWORD = 0x00000008;
 pub const CREATE_NEW_PROCESS_GROUP: DWORD = 0x00000200;
 pub const CREATE_UNICODE_ENVIRONMENT: DWORD = 0x00000400;
+pub const EXTENDED_STARTUPINFO_PRESENT: DWORD = 0x00080000;
 pub const STARTF_USESTDHANDLES: DWORD = 0x00000100;
 
 pub const AF_INET: c_int = 2;
@@ -587,6 +590,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,
@@ -1078,6 +1087,22 @@ extern "system" {
         lpFilePart: *mut LPWSTR,
     ) -> DWORD;
     pub fn GetFileAttributesW(lpFileName: LPCWSTR) -> DWORD;
+    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: LPVOID,
+        cbSize: SIZE_T,
+        lpPreviousValue: LPVOID,
+        lpReturnSize: PSIZE_T,
+    ) -> BOOL;
+    pub fn DeleteProcThreadAttributeList(lpAttributeList: LPPROC_THREAD_ATTRIBUTE_LIST);
 }
 
 #[link(name = "ws2_32")]
diff --git a/library/std/src/sys/windows/process.rs b/library/std/src/sys/windows/process.rs
@@ -16,6 +16,7 @@ use crate::num::NonZeroI32;
 use crate::os::windows::ffi::{OsStrExt, OsStringExt};
 use crate::os::windows::io::{AsRawHandle, FromRawHandle, IntoRawHandle};
 use crate::path::{Path, PathBuf};
+use crate::pin::Pin;
 use crate::ptr;
 use crate::sys::c;
 use crate::sys::c::NonZeroDWORD;
@@ -166,6 +167,7 @@ pub struct Command {
     stdout: Option<Stdio>,
     stderr: Option<Stdio>,
     force_quotes_enabled: bool,
+    proc_thread_attributes: BTreeMap<usize, ProcThreadAttributeValue>,
 }
 
 pub enum Stdio {
@@ -202,6 +204,7 @@ impl Command {
             stdout: None,
             stderr: None,
             force_quotes_enabled: false,
+            proc_thread_attributes: Default::default(),
         }
     }
 
@@ -251,6 +254,9 @@ impl Command {
     pub fn get_current_dir(&self) -> Option<&Path> {
         self.cwd.as_ref().map(|cwd| Path::new(cwd))
     }
+    pub fn process_thread_attribute(&mut self, attribute: usize, ptr: *mut c_void, size: usize) {
+        self.proc_thread_attributes.insert(attribute, ProcThreadAttributeValue { ptr, size });
+    }
 
     pub fn spawn(
         &mut self,
@@ -259,9 +265,9 @@ impl Command {
     ) -> io::Result<(Process, StdioPipes)> {
         let maybe_env = self.env.capture_if_changed();
 
-        let mut si = zeroed_startupinfo();
-        si.cb = mem::size_of::<c::STARTUPINFO>() as c::DWORD;
-        si.dwFlags = c::STARTF_USESTDHANDLES;
+        let mut si = zeroed_startupinfo_ex();
+        si.StartupInfo.cb = mem::size_of::<c::STARTUPINFO>() as c::DWORD;
+        si.StartupInfo.dwFlags = c::STARTF_USESTDHANDLES;
 
         let child_paths = if let Some(env) = maybe_env.as_ref() {
             env.get(&EnvKey::new("PATH")).map(|s| s.as_os_str())
@@ -305,9 +311,24 @@ impl Command {
         let stdin = stdin.to_handle(c::STD_INPUT_HANDLE, &mut pipes.stdin)?;
         let stdout = stdout.to_handle(c::STD_OUTPUT_HANDLE, &mut pipes.stdout)?;
         let stderr = stderr.to_handle(c::STD_ERROR_HANDLE, &mut pipes.stderr)?;
-        si.hStdInput = stdin.as_raw_handle();
-        si.hStdOutput = stdout.as_raw_handle();
-        si.hStdError = stderr.as_raw_handle();
+        si.StartupInfo.hStdInput = stdin.as_raw_handle();
+        si.StartupInfo.hStdOutput = stdout.as_raw_handle();
+        si.StartupInfo.hStdError = stderr.as_raw_handle();
+
+        let mut attributes = if !self.proc_thread_attributes.is_empty() {
+            Some(make_proc_thread_attributes(&self.proc_thread_attributes)?)
+        } else {
+            None
+        };
+        si.lpAttributeList = attributes.as_mut().map_or(ptr::null_mut(), |buf| {
+            // Indicate that lpAttributeList is present by setting
+            // EXTENDED_STARTUPINFO_PRESENT and adjusting the size
+            // value of the struct.
+            flags |= c::EXTENDED_STARTUPINFO_PRESENT;
+            si.StartupInfo.cb = mem::size_of::<c::STARTUPINFOEX>() as u32;
+
+            buf.0.as_mut_ptr().cast()
+        });
 
         let program = to_u16s(&program)?;
         unsafe {
@@ -320,7 +341,7 @@ impl Command {
                 flags,
                 envp,
                 dirp,
-                &mut si,
+                &mut si as *mut _ as *mut _,
                 &mut pi,
             ))
         }?;
@@ -687,26 +708,29 @@ impl From<u8> for ExitCode {
     }
 }
 
-fn zeroed_startupinfo() -> c::STARTUPINFO {
-    c::STARTUPINFO {
-        cb: 0,
-        lpReserved: ptr::null_mut(),
-        lpDesktop: ptr::null_mut(),
-        lpTitle: ptr::null_mut(),
-        dwX: 0,
-        dwY: 0,
-        dwXSize: 0,
-        dwYSize: 0,
-        dwXCountChars: 0,
-        dwYCountCharts: 0,
-        dwFillAttribute: 0,
-        dwFlags: 0,
-        wShowWindow: 0,
-        cbReserved2: 0,
-        lpReserved2: ptr::null_mut(),
-        hStdInput: c::INVALID_HANDLE_VALUE,
-        hStdOutput: c::INVALID_HANDLE_VALUE,
-        hStdError: c::INVALID_HANDLE_VALUE,
+fn zeroed_startupinfo_ex() -> c::STARTUPINFOEX {
+    c::STARTUPINFOEX {
+        StartupInfo: c::STARTUPINFO {
+            cb: 0,
+            lpReserved: ptr::null_mut(),
+            lpDesktop: ptr::null_mut(),
+            lpTitle: ptr::null_mut(),
+            dwX: 0,
+            dwY: 0,
+            dwXSize: 0,
+            dwYSize: 0,
+            dwXCountChars: 0,
+            dwYCountCharts: 0,
+            dwFillAttribute: 0,
+            dwFlags: 0,
+            wShowWindow: 0,
+            cbReserved2: 0,
+            lpReserved2: ptr::null_mut(),
+            hStdInput: c::INVALID_HANDLE_VALUE,
+            hStdOutput: c::INVALID_HANDLE_VALUE,
+            hStdError: c::INVALID_HANDLE_VALUE,
+        },
+        lpAttributeList: ptr::null_mut(),
     }
 }
 
@@ -843,6 +867,57 @@ fn make_dirp(d: Option<&OsString>) -> io::Result<(*const u16, Vec<u16>)> {
     }
 }
 
+/// Wrapper around a u8 buffer which ensures that a ProcThreadAttributeList
+/// is not moved during use and is not freed before calling Delete.
+struct ProcThreadAttributeList(Pin<Vec<u8>>);
+impl Drop for ProcThreadAttributeList {
+    fn drop(&mut self) {
+        unsafe { c::DeleteProcThreadAttributeList(self.0.as_mut_ptr().cast()) }
+    }
+}
+/// Wrapper around the value data to be used as a Process Thread Attribute.
+/// This exists primarily to force the raw pointer type to be `Send` and `Sync`
+/// without needing to `unsafe impl` them for `Command`.
+#[derive(Copy, Clone)]
+struct ProcThreadAttributeValue {
+    ptr: *mut c_void,
+    size: usize,
+}
+unsafe impl Send for ProcThreadAttributeValue {}
+unsafe impl Sync for ProcThreadAttributeValue {}
+
+fn make_proc_thread_attributes(
+    attributes: &BTreeMap<usize, ProcThreadAttributeValue>,
+) -> io::Result<ProcThreadAttributeList> {
+    let mut buffer_size = 0;
+    let count = attributes.len() as u32;
+    // First call to get required size. Error is expected.
+    unsafe { c::InitializeProcThreadAttributeList(ptr::null_mut(), count, 0, &mut buffer_size) };
+    let mut buf = Pin::new(crate::vec::from_elem(0u8, buffer_size as usize));
+    // Second call to really initialize the buffer.
+    cvt(unsafe {
+        c::InitializeProcThreadAttributeList(buf.as_mut_ptr().cast(), count, 0, &mut buffer_size)
+    })?;
+    let mut attribute_list = ProcThreadAttributeList(buf);
+    // Add our attributes to the buffer.
+    // It's theoretically possible for the count to overflow a u32. Therefore, make
+    // sure we don't add more attributes than we actually initialized the buffer for.
+    for (&attribute, &value) in attributes.iter().take(count as usize) {
+        cvt(unsafe {
+            c::UpdateProcThreadAttribute(
+                attribute_list.0.as_mut_ptr().cast(),
+                0,
+                attribute,
+                value.ptr,
+                value.size,
+                ptr::null_mut(),
+                ptr::null_mut(),
+            )
+        })?;
+    }
+    Ok(attribute_list)
+}
+
 pub struct CommandArgs<'a> {
     iter: crate::slice::Iter<'a, Arg>,
 }
