Add job work (ends up a PR merged upstream that enables this rust-lang/rust#114848

Signed-off-by: James Sturtevant <[email protected]>

Author: jsturtevant

Cargo.toml

@@ -18,7 +18,9 @@ windows-sys = { version = "0.48.0", features = [
     "Win32_Security_Cryptography",
     "Win32_System_Console",
     "Win32_System_WindowsProgramming",
-    
 ]}
 cvt = "0.1.2"
 libc = "0.2.147"
+
+[dev-dependencies]
+windows = { version="0.51", features = ["Win32_Foundation", "Win32_Security", "Win32_System_JobObjects"] }

src/args.rs

@@ -5,7 +5,7 @@
 
 use crate::c;
 use crate::path_ext::get_long_path;
-use crate::windows::{ensure_no_nuls, to_u16s};
+use crate::util::{ensure_no_nuls, to_u16s};
 
 use std::ffi::OsString;
 use std::fmt;
@@ -177,7 +177,7 @@ pub(crate) fn to_user_path(path: &Path) -> io::Result<Vec<u16>> {
     from_wide_to_user_path(to_u16s(path)?)
 }
 pub(crate) fn from_wide_to_user_path(mut path: Vec<u16>) -> io::Result<Vec<u16>> {
-    use crate::windows::fill_utf16_buf;
+    use crate::util::fill_utf16_buf;
     use std::ptr;
 
     // UTF-16 encoded code points, used in parsing and building UTF-16 paths.

src/c.rs

@@ -34,7 +34,7 @@ pub use windows_sys::Win32::System::Threading::TerminateProcess;
 pub use windows_sys::Win32::System::Threading::{
     CreateEventW, GetCurrentProcessId, GetExitCodeProcess, GetProcessId, SleepEx,
     WaitForSingleObject, CREATE_NEW_PROCESS_GROUP, CREATE_UNICODE_ENVIRONMENT, DETACHED_PROCESS,
-    INFINITE, STARTF_USESTDHANDLES,
+    INFINITE, STARTF_USESTDHANDLES, EXTENDED_STARTUPINFO_PRESENT,
 };
 use windows_sys::Win32::System::Threading::{PROCESS_CREATION_FLAGS, STARTUPINFOW_FLAGS};
 pub use windows_sys::Win32::System::IO::CancelIo;
@@ -382,6 +382,19 @@ impl ::core::clone::Clone for STARTUPINFOW {
         *self
     }
 }
+pub type LPPROC_THREAD_ATTRIBUTE_LIST = *mut ::core::ffi::c_void;
+
+#[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 PROCESS_INFORMATION {
@@ -396,3 +409,35 @@ impl ::core::clone::Clone for PROCESS_INFORMATION {
         *self
     }
 }
+
+#[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 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 DeleteProcThreadAttributeList(lpattributelist: LPPROC_THREAD_ATTRIBUTE_LIST) -> ();
+}
+
+#[link(name = "kernel32")]
+extern "system" {
+    pub fn CloseHandle(hobject: HANDLE) -> BOOL;
+}

src/child.rs

@@ -26,6 +26,12 @@ pub struct Child {
     pub(crate) main_thread_handle: Handle,
 }
 
+impl AsRawHandle for Child {
+    fn as_raw_handle(&self) -> c::HANDLE {
+        self.handle.as_raw_handle()
+    }
+}
+
 impl Child {
     pub fn kill(&mut self) -> io::Result<()> {
         let result = unsafe { c::TerminateProcess(self.handle.as_raw_handle() as isize, 1) };
@@ -118,3 +124,14 @@ pub fn wait_with_output(
     let status = process.wait()?;
     Ok((status, stdout, stderr))
 }
+
+pub trait ChildExt {
+    fn main_thread_handle(&self) -> BorrowedHandle<'_>;
+}
+
+
+impl ChildExt for Child {
+    fn main_thread_handle(&self) -> BorrowedHandle<'_> {
+        self.main_thread_handle.as_handle()
+    }
+}

src/command.rs

@@ -6,12 +6,14 @@ use crate::file::{open, OpenOptions};
 use crate::handle::Handle;
 use crate::path_ext;
 use crate::pipe::{self, AnonPipe};
-use crate::{c, windows};
+use crate::{c, util};
 use cvt::cvt;
 use std::collections::BTreeMap;
 use std::env::consts::{EXE_EXTENSION, EXE_SUFFIX};
 use std::ffi::{c_void, OsStr, OsString};
 use std::fs::File;
+use std::io::ErrorKind;
+use std::mem::MaybeUninit;
 use std::os::windows::prelude::{
     AsRawHandle, FromRawHandle, IntoRawHandle, OsStrExt, OsStringExt, RawHandle,
 };
@@ -45,6 +47,7 @@ pub struct Command {
     stdout: Option<Stdio>,
     stderr: Option<Stdio>,
     force_quotes_enabled: bool,
+    proc_thread_attributes: BTreeMap<usize, ProcThreadAttributeValue>,
 }
 
 impl Command {
@@ -60,6 +63,7 @@ impl Command {
             stdout: None,
             stderr: None,
             force_quotes_enabled: false,
+            proc_thread_attributes: Default::default(),
         }
     }
 
@@ -115,6 +119,17 @@ impl Command {
     //     self.env.iter()
     // }
 
+    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 get_current_dir(&self) -> Option<&Path> {
         self.cwd.as_ref().map(Path::new)
     }
@@ -191,7 +206,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`.
@@ -205,6 +219,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(),
@@ -215,7 +250,7 @@ impl Command {
                 flags,
                 envp,
                 dirp,
-                &si,
+                si_ptr,
                 &mut pi,
             ))
         }?;
@@ -339,7 +374,7 @@ pub struct CommandArgs<'a> {
 
 // Get `cmd.exe` for use with bat scripts, encoded as a UTF-16 string.
 fn command_prompt() -> io::Result<Vec<u16>> {
-    let mut system: Vec<u16> = windows::fill_utf16_buf(
+    let mut system: Vec<u16> = util::fill_utf16_buf(
         |buf, size| unsafe { c::GetSystemDirectoryW(buf, size) },
         |buf| buf.into(),
     )?;
@@ -557,15 +592,15 @@ where
     // 3 & 4. System paths
     // SAFETY: This uses `fill_utf16_buf` to safely call the OS functions.
     unsafe {
-        if let Ok(Some(path)) = windows::fill_utf16_buf(
+        if let Ok(Some(path)) = util::fill_utf16_buf(
             |buf, size| c::GetSystemDirectoryW(buf, size),
             |buf| exists(PathBuf::from(OsString::from_wide(buf))),
         ) {
             return Some(path);
         }
         #[cfg(not(target_vendor = "uwp"))]
         {
-            if let Ok(Some(path)) = windows::fill_utf16_buf(
+            if let Ok(Some(path)) = util::fill_utf16_buf(
                 |buf, size| c::GetWindowsDirectoryW(buf, size),
                 |buf| exists(PathBuf::from(OsString::from_wide(buf))),
             ) {
@@ -600,3 +635,79 @@ fn program_exists(path: &Path) -> Option<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::Error::new(
+            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)
+}
+