From edefa8b105cd4a35b689e85187689f4fb4085ede Mon Sep 17 00:00:00 2001 From: Michael van Straten Date: Fri, 25 Aug 2023 16:19:16 +0200 Subject: [PATCH] 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 --- library/std/src/os/windows/process.rs | 69 +++++++++++ library/std/src/process/tests.rs | 85 +++++++++++++ library/std/src/sys/windows/c/windows_sys.lst | 5 + library/std/src/sys/windows/c/windows_sys.rs | 37 ++++++ library/std/src/sys/windows/process.rs | 112 +++++++++++++++++- 5 files changed, 306 insertions(+), 2 deletions(-) diff --git a/library/std/src/os/windows/process.rs b/library/std/src/os/windows/process.rs index 073168cf2d209..5d87216ff1b4c 100644 --- a/library/std/src/os/windows/process.rs +++ b/library/std/src/os/windows/process.rs @@ -192,6 +192,66 @@ pub trait CommandExt: Sealed { /// ``` #[unstable(feature = "windows_process_extensions_async_pipes", issue = "98289")] fn async_pipes(&mut self, always_async: bool) -> &mut process::Command; + + /// Sets a raw attribute on the command, providing extended configuration options for Windows processes. + /// + /// This method allows you to specify custom attributes for a child process on Windows systems using raw attribute values. + /// Raw attributes provide extended configurability for process creation, but their usage can be complex and potentially unsafe. + /// + /// The `attribute` parameter specifies the raw attribute to be set, while the `value` parameter holds the value associated with that attribute. + /// Please refer to the [`windows-rs`](https://microsoft.github.io/windows-docs-rs/doc/windows/) documentation or the [`Win32 API documentation`](https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-updateprocthreadattribute) for detailed information about available attributes and their meanings. + /// + /// # Note + /// + /// The maximum number of raw attributes is the value of [`u32::MAX`]. + /// If this limit is exceeded, the call to [`process::Command::spawn`] will return an `Error` indicating that the maximum number of attributes has been exceeded. + /// # Safety + /// + /// The usage of raw attributes is potentially unsafe and should be done with caution. Incorrect attribute values or improper configuration can lead to unexpected behavior or errors. + /// + /// # Example + /// + /// The following example demonstrates how to create a child process with a specific parent process ID using a raw attribute. + /// + /// ```rust + /// #![feature(windows_process_extensions_raw_attribute)] + /// use std::os::windows::{process::CommandExt, io::AsRawHandle}; + /// use std::process::Command; + /// + /// # struct ProcessDropGuard(std::process::Child); + /// # impl Drop for ProcessDropGuard { + /// # fn drop(&mut self) { + /// # let _ = self.0.kill(); + /// # } + /// # } + /// + /// let parent = Command::new("cmd").spawn()?; + /// + /// let mut child_cmd = Command::new("cmd"); + /// + /// const PROC_THREAD_ATTRIBUTE_PARENT_PROCESS: usize = 0x00020000; + /// + /// unsafe { + /// child_cmd.raw_attribute(PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, parent.as_raw_handle() as isize); + /// } + /// # + /// # let parent = ProcessDropGuard(parent); + /// + /// let mut child = child_cmd.spawn()?; + /// + /// # child.kill()?; + /// # Ok::<(), std::io::Error>(()) + /// ``` + /// + /// # Safety Note + /// + /// Remember that improper use of raw attributes can lead to undefined behavior or security vulnerabilities. Always consult the documentation and ensure proper attribute values are used. + #[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")] + unsafe fn raw_attribute( + &mut self, + attribute: usize, + value: T, + ) -> &mut process::Command; } #[stable(feature = "windows_process_extensions", since = "1.16.0")] @@ -219,6 +279,15 @@ impl CommandExt for process::Command { let _ = always_async; self } + + unsafe fn raw_attribute( + &mut self, + attribute: usize, + value: T, + ) -> &mut process::Command { + self.as_inner_mut().raw_attribute(attribute, value); + self + } } #[unstable(feature = "windows_process_extensions_main_thread_handle", issue = "96723")] diff --git a/library/std/src/process/tests.rs b/library/std/src/process/tests.rs index 366b591466c07..0063623728821 100644 --- a/library/std/src/process/tests.rs +++ b/library/std/src/process/tests.rs @@ -434,6 +434,91 @@ 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::mem; + use crate::os::windows::io::AsRawHandle; + use crate::os::windows::process::CommandExt; + use crate::sys::c::{CloseHandle, BOOL, HANDLE}; + use crate::sys::cvt; + + #[repr(C)] + #[allow(non_snake_case)] + struct PROCESSENTRY32W { + dwSize: u32, + cntUsage: u32, + th32ProcessID: u32, + th32DefaultHeapID: usize, + th32ModuleID: u32, + cntThreads: u32, + th32ParentProcessID: u32, + pcPriClassBase: i32, + dwFlags: u32, + szExeFile: [u16; 260], + } + + extern "system" { + fn CreateToolhelp32Snapshot(dwflags: u32, th32processid: u32) -> HANDLE; + fn Process32First(hsnapshot: HANDLE, lppe: *mut PROCESSENTRY32W) -> BOOL; + fn Process32Next(hsnapshot: HANDLE, lppe: *mut PROCESSENTRY32W) -> BOOL; + } + + const PROC_THREAD_ATTRIBUTE_PARENT_PROCESS: usize = 0x00020000; + const TH32CS_SNAPPROCESS: u32 = 0x00000002; + + struct ProcessDropGuard(crate::process::Child); + + impl Drop for ProcessDropGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + } + } + + let parent = ProcessDropGuard(Command::new("cmd").spawn().unwrap()); + + let mut child_cmd = Command::new("cmd"); + + unsafe { + child_cmd + .raw_attribute(PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, parent.0.as_raw_handle() as isize); + } + + let child = ProcessDropGuard(child_cmd.spawn().unwrap()); + + let h_snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) }; + + let mut process_entry = PROCESSENTRY32W { + dwSize: mem::size_of::() as u32, + cntUsage: 0, + th32ProcessID: 0, + th32DefaultHeapID: 0, + th32ModuleID: 0, + cntThreads: 0, + th32ParentProcessID: 0, + pcPriClassBase: 0, + dwFlags: 0, + szExeFile: [0; 260], + }; + + unsafe { cvt(Process32First(h_snapshot, &mut process_entry as *mut _)) }.unwrap(); + + loop { + if child.0.id() == process_entry.th32ProcessID { + break; + } + unsafe { cvt(Process32Next(h_snapshot, &mut process_entry as *mut _)) }.unwrap(); + } + + unsafe { cvt(CloseHandle(h_snapshot)) }.unwrap(); + + assert_eq!(parent.0.id(), process_entry.th32ParentProcessID); + + drop(child) +} + #[test] fn test_command_implements_send_sync() { fn take_send_sync_type(_: T) {} diff --git a/library/std/src/sys/windows/c/windows_sys.lst b/library/std/src/sys/windows/c/windows_sys.lst index 631aedd26b4ea..5469ba42eb11a 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 02377087173a7..b24f6dc900fc5 100644 --- a/library/std/src/sys/windows/c/windows_sys.rs +++ b/library/std/src/sys/windows/c/windows_sys.rs @@ -155,6 +155,10 @@ extern "system" { pub fn DeleteFileW(lpfilename: PCWSTR) -> BOOL; } #[link(name = "kernel32")] +extern "system" { + pub fn DeleteProcThreadAttributeList(lpattributelist: LPPROC_THREAD_ATTRIBUTE_LIST) -> (); +} +#[link(name = "kernel32")] extern "system" { pub fn DeviceIoControl( hdevice: HANDLE, @@ -371,6 +375,15 @@ extern "system" { ) -> BOOL; } #[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, @@ -543,6 +556,18 @@ extern "system" { pub fn TryAcquireSRWLockShared(srwlock: *mut RTL_SRWLOCK) -> BOOLEAN; } #[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, @@ -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 2dd0c67acdb4e..92edd2f6112f3 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, stderr: Option, force_quotes_enabled: bool, + proc_thread_attributes: BTreeMap, } 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( + &mut self, + attribute: usize, + value: T, + ) { + self.proc_thread_attributes.insert( + attribute, + ProcThreadAttributeValue { size: mem::size_of::(), 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::() 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::() 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:: 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)> { } } +struct ProcThreadAttributeList(Box<[MaybeUninit]>); + +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, + size: usize, +} + +fn make_proc_thread_attribute_list( + attributes: &BTreeMap, +) -> io::Result { + // 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>, }