Added option to set ProcThreadAttributes for Windows processes

This implements the ability to add arbitrary attributes to a command on Windows targets using a new `raw_attribute` method on the [`CommandExt`](https://doc.rust-lang.org/stable/std/os/windows/process/trait.CommandExt.html) trait. Setting these attributes provides extended configuration options for Windows processes.

Co-authored-by: Tyler Ruckinger <t.ruckinger@gmail.com>

3 files changed, 152 insertions, 2 deletions

diff --git a/library/std/src/sys/windows/c/windows_sys.lst b/library/std/src/sys/windows/c/windows_sys.lst
index 631aedd26b4..5469ba42eb1 100644
--- a/library/std/src/sys/windows/c/windows_sys.lst
+++ b/library/std/src/sys/windows/c/windows_sys.lst
@@ -2510,6 +2510,7 @@ Windows.Win32.System.Threading.CreateProcessW
 Windows.Win32.System.Threading.CreateThread
 Windows.Win32.System.Threading.DEBUG_ONLY_THIS_PROCESS
 Windows.Win32.System.Threading.DEBUG_PROCESS
+Windows.Win32.System.Threading.DeleteProcThreadAttributeList
 Windows.Win32.System.Threading.DETACHED_PROCESS
 Windows.Win32.System.Threading.ExitProcess
 Windows.Win32.System.Threading.EXTENDED_STARTUPINFO_PRESENT
@@ -2524,8 +2525,10 @@ Windows.Win32.System.Threading.INFINITE
 Windows.Win32.System.Threading.INHERIT_CALLER_PRIORITY
 Windows.Win32.System.Threading.INHERIT_PARENT_AFFINITY
 Windows.Win32.System.Threading.INIT_ONCE_INIT_FAILED
+Windows.Win32.System.Threading.InitializeProcThreadAttributeList
 Windows.Win32.System.Threading.InitOnceBeginInitialize
 Windows.Win32.System.Threading.InitOnceComplete
+Windows.Win32.System.Threading.LPPROC_THREAD_ATTRIBUTE_LIST
 Windows.Win32.System.Threading.LPTHREAD_START_ROUTINE
 Windows.Win32.System.Threading.NORMAL_PRIORITY_CLASS
 Windows.Win32.System.Threading.OpenProcessToken
@@ -2561,6 +2564,7 @@ Windows.Win32.System.Threading.STARTF_USEPOSITION
 Windows.Win32.System.Threading.STARTF_USESHOWWINDOW
 Windows.Win32.System.Threading.STARTF_USESIZE
 Windows.Win32.System.Threading.STARTF_USESTDHANDLES
+Windows.Win32.System.Threading.STARTUPINFOEXW
 Windows.Win32.System.Threading.STARTUPINFOW
 Windows.Win32.System.Threading.STARTUPINFOW_FLAGS
 Windows.Win32.System.Threading.SwitchToThread
@@ -2575,6 +2579,7 @@ Windows.Win32.System.Threading.TlsGetValue
 Windows.Win32.System.Threading.TlsSetValue
 Windows.Win32.System.Threading.TryAcquireSRWLockExclusive
 Windows.Win32.System.Threading.TryAcquireSRWLockShared
+Windows.Win32.System.Threading.UpdateProcThreadAttribute
 Windows.Win32.System.Threading.WaitForMultipleObjects
 Windows.Win32.System.Threading.WaitForSingleObject
 Windows.Win32.System.Threading.WakeAllConditionVariable
diff --git a/library/std/src/sys/windows/c/windows_sys.rs b/library/std/src/sys/windows/c/windows_sys.rs
index 02377087173..b24f6dc900f 100644
--- a/library/std/src/sys/windows/c/windows_sys.rs
+++ b/library/std/src/sys/windows/c/windows_sys.rs
@@ -156,6 +156,10 @@ extern "system" {
 }
 #[link(name = "kernel32")]
 extern "system" {
+    pub fn DeleteProcThreadAttributeList(lpattributelist: LPPROC_THREAD_ATTRIBUTE_LIST) -> ();
+}
+#[link(name = "kernel32")]
+extern "system" {
     pub fn DeviceIoControl(
         hdevice: HANDLE,
         dwiocontrolcode: u32,
@@ -372,6 +376,15 @@ extern "system" {
 }
 #[link(name = "kernel32")]
 extern "system" {
+    pub fn InitializeProcThreadAttributeList(
+        lpattributelist: LPPROC_THREAD_ATTRIBUTE_LIST,
+        dwattributecount: u32,
+        dwflags: u32,
+        lpsize: *mut usize,
+    ) -> BOOL;
+}
+#[link(name = "kernel32")]
+extern "system" {
     pub fn MoveFileExW(
         lpexistingfilename: PCWSTR,
         lpnewfilename: PCWSTR,
@@ -544,6 +557,18 @@ extern "system" {
 }
 #[link(name = "kernel32")]
 extern "system" {
+    pub fn UpdateProcThreadAttribute(
+        lpattributelist: LPPROC_THREAD_ATTRIBUTE_LIST,
+        dwflags: u32,
+        attribute: usize,
+        lpvalue: *const ::core::ffi::c_void,
+        cbsize: usize,
+        lppreviousvalue: *mut ::core::ffi::c_void,
+        lpreturnsize: *const usize,
+    ) -> BOOL;
+}
+#[link(name = "kernel32")]
+extern "system" {
     pub fn WaitForMultipleObjects(
         ncount: u32,
         lphandles: *const HANDLE,
@@ -3567,6 +3592,7 @@ pub type LPOVERLAPPED_COMPLETION_ROUTINE = ::core::option::Option<
         lpoverlapped: *mut OVERLAPPED,
     ) -> (),
 >;
+pub type LPPROC_THREAD_ATTRIBUTE_LIST = *mut ::core::ffi::c_void;
 pub type LPPROGRESS_ROUTINE = ::core::option::Option<
     unsafe extern "system" fn(
         totalfilesize: i64,
@@ -3833,6 +3859,17 @@ pub const STARTF_USESHOWWINDOW: STARTUPINFOW_FLAGS = 1u32;
 pub const STARTF_USESIZE: STARTUPINFOW_FLAGS = 2u32;
 pub const STARTF_USESTDHANDLES: STARTUPINFOW_FLAGS = 256u32;
 #[repr(C)]
+pub struct STARTUPINFOEXW {
+    pub StartupInfo: STARTUPINFOW,
+    pub lpAttributeList: LPPROC_THREAD_ATTRIBUTE_LIST,
+}
+impl ::core::marker::Copy for STARTUPINFOEXW {}
+impl ::core::clone::Clone for STARTUPINFOEXW {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+#[repr(C)]
 pub struct STARTUPINFOW {
     pub cb: u32,
     pub lpReserved: PWSTR,
diff --git a/library/std/src/sys/windows/process.rs b/library/std/src/sys/windows/process.rs
index 2dd0c67acdb..92edd2f6112 100644
--- a/library/std/src/sys/windows/process.rs
+++ b/library/std/src/sys/windows/process.rs
@@ -11,6 +11,7 @@ use crate::ffi::{OsStr, OsString};
 use crate::fmt;
 use crate::io::{self, Error, ErrorKind};
 use crate::mem;
+use crate::mem::MaybeUninit;
 use crate::num::NonZeroI32;
 use crate::os::windows::ffi::{OsStrExt, OsStringExt};
 use crate::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle};
@@ -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 {
@@ -195,6 +197,7 @@ impl Command {
             stdout: None,
             stderr: None,
             force_quotes_enabled: false,
+            proc_thread_attributes: Default::default(),
         }
     }
 
@@ -245,6 +248,17 @@ impl Command {
         self.cwd.as_ref().map(|cwd| Path::new(cwd))
     }
 
+    pub unsafe fn raw_attribute<T: Copy + Send + Sync + 'static>(
+        &mut self,
+        attribute: usize,
+        value: T,
+    ) {
+        self.proc_thread_attributes.insert(
+            attribute,
+            ProcThreadAttributeValue { size: mem::size_of::<T>(), data: Box::new(value) },
+        );
+    }
+
     pub fn spawn(
         &mut self,
         default: Stdio,
@@ -308,7 +322,6 @@ impl Command {
         let stderr = stderr.to_handle(c::STD_ERROR_HANDLE, &mut pipes.stderr)?;
 
         let mut si = zeroed_startupinfo();
-        si.cb = mem::size_of::<c::STARTUPINFOW>() as c::DWORD;
 
         // If at least one of stdin, stdout or stderr are set (i.e. are non null)
         // then set the `hStd` fields in `STARTUPINFO`.
@@ -322,6 +335,27 @@ impl Command {
             si.hStdError = stderr.as_raw_handle();
         }
 
+        let si_ptr: *mut c::STARTUPINFOW;
+
+        let mut proc_thread_attribute_list;
+        let mut si_ex;
+
+        if !self.proc_thread_attributes.is_empty() {
+            si.cb = mem::size_of::<c::STARTUPINFOEXW>() as u32;
+            flags |= c::EXTENDED_STARTUPINFO_PRESENT;
+
+            proc_thread_attribute_list =
+                make_proc_thread_attribute_list(&self.proc_thread_attributes)?;
+            si_ex = c::STARTUPINFOEXW {
+                StartupInfo: si,
+                lpAttributeList: proc_thread_attribute_list.0.as_mut_ptr() as _,
+            };
+            si_ptr = &mut si_ex as *mut _ as _;
+        } else {
+            si.cb = mem::size_of::<c::STARTUPINFOW> as c::DWORD;
+            si_ptr = &mut si as *mut _ as _;
+        }
+
         unsafe {
             cvt(c::CreateProcessW(
                 program.as_ptr(),
@@ -332,7 +366,7 @@ impl Command {
                 flags,
                 envp,
                 dirp,
-                &si,
+                si_ptr,
                 &mut pi,
             ))
         }?;
@@ -831,6 +865,80 @@ fn make_dirp(d: Option<&OsString>) -> io::Result<(*const u16, Vec<u16>)> {
     }
 }
 
+struct ProcThreadAttributeList(Box<[MaybeUninit<u8>]>);
+
+impl Drop for ProcThreadAttributeList {
+    fn drop(&mut self) {
+        let lp_attribute_list = self.0.as_mut_ptr() as _;
+        unsafe { c::DeleteProcThreadAttributeList(lp_attribute_list) }
+    }
+}
+
+/// Wrapper around the value data to be used as a Process Thread Attribute.
+struct ProcThreadAttributeValue {
+    data: Box<dyn Send + Sync>,
+    size: usize,
+}
+
+fn make_proc_thread_attribute_list(
+    attributes: &BTreeMap<usize, ProcThreadAttributeValue>,
+) -> io::Result<ProcThreadAttributeList> {
+    // To initialize our ProcThreadAttributeList, we need to determine
+    // how many bytes to allocate for it. The Windows API simplifies this process
+    // by allowing us to call `InitializeProcThreadAttributeList` with
+    // a null pointer to retrieve the required size.
+    let mut required_size = 0;
+    let Ok(attribute_count) = attributes.len().try_into() else {
+        return Err(io::const_io_error!(
+            ErrorKind::InvalidInput,
+            "maximum number of ProcThreadAttributes exceeded",
+        ));
+    };
+    unsafe {
+        c::InitializeProcThreadAttributeList(
+            ptr::null_mut(),
+            attribute_count,
+            0,
+            &mut required_size,
+        )
+    };
+
+    let mut proc_thread_attribute_list = ProcThreadAttributeList(
+        vec![MaybeUninit::uninit(); required_size as usize].into_boxed_slice(),
+    );
+
+    // Once we've allocated the necessary memory, it's safe to invoke
+    // `InitializeProcThreadAttributeList` to properly initialize the list.
+    cvt(unsafe {
+        c::InitializeProcThreadAttributeList(
+            proc_thread_attribute_list.0.as_mut_ptr() as *mut _,
+            attribute_count,
+            0,
+            &mut required_size,
+        )
+    })?;
+
+    // # Add our attributes to the buffer.
+    // It's theoretically possible for the attribute count to exceed a u32 value.
+    // Therefore, we ensure that we don't add more attributes than the buffer was initialized for.
+    for (&attribute, value) in attributes.iter().take(attribute_count as usize) {
+        let value_ptr = &*value.data as *const (dyn Send + Sync) as _;
+        cvt(unsafe {
+            c::UpdateProcThreadAttribute(
+                proc_thread_attribute_list.0.as_mut_ptr() as _,
+                0,
+                attribute,
+                value_ptr,
+                value.size,
+                ptr::null_mut(),
+                ptr::null_mut(),
+            )
+        })?;
+    }
+
+    Ok(proc_thread_attribute_list)
+}
+
 pub struct CommandArgs<'a> {
     iter: crate::slice::Iter<'a, Arg>,
 }