From 0475e2fea0520d1a40145faef0ce660e71d6a8d0 Mon Sep 17 00:00:00 2001 From: Tyler Ruckinger Date: Fri, 20 Aug 2021 14:04:26 -0400 Subject: [PATCH 1/2] add windows ext method for proc thread attributes Added doc comments, don't take too many attrs kill the cmds in test --- library/std/src/os/windows/process.rs | 28 +++++- library/std/src/process/tests.rs | 36 +++++++ library/std/src/sys/windows/c.rs | 25 +++++ library/std/src/sys/windows/process.rs | 129 +++++++++++++++++++------ 4 files changed, 190 insertions(+), 28 deletions(-) diff --git a/library/std/src/os/windows/process.rs b/library/std/src/os/windows/process.rs index 9510d104806db..c796e228b0eb1 100644 --- 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>(&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 index e5cdc4737068a..cc2e196038462 100644 --- 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::(), + ); + } + 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 index 9b61b2476d5bb..dea91a4ca6fa6 100644 --- 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 index fafd1412d4cb3..c53a8c6383a17 100644 --- 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, stderr: Option, force_quotes_enabled: bool, + proc_thread_attributes: BTreeMap, } 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::() as c::DWORD; - si.dwFlags = c::STARTF_USESTDHANDLES; + let mut si = zeroed_startupinfo_ex(); + si.StartupInfo.cb = mem::size_of::() 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::() 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 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)> { } } +/// 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>); +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, +) -> io::Result { + 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>, } From d20fc4ead76ffee0a6fccec91e44ca38633fd2c8 Mon Sep 17 00:00:00 2001 From: Tyler Ruckinger Date: Thu, 9 Sep 2021 08:51:05 -0400 Subject: [PATCH 2/2] Take generic value type --- library/std/src/os/windows/process.rs | 22 +++++++++------------ library/std/src/process/tests.rs | 8 +++----- library/std/src/sys/windows/process.rs | 27 +++++++++++++++----------- 3 files changed, 28 insertions(+), 29 deletions(-) diff --git a/library/std/src/os/windows/process.rs b/library/std/src/os/windows/process.rs index c796e228b0eb1..32911e57abfda 100644 --- 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::{c_void, OsStr}; +use crate::ffi::OsStr; use crate::os::windows::io::{ AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle, OwnedHandle, RawHandle, }; @@ -166,17 +166,14 @@ pub trait CommandExt: Sealed { /// /// # 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. + /// - The attribute and value pair must be supplied in accordance with [Win32 API usage][1]. + /// + /// [1]: https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-updateprocthreadattribute #[unstable(feature = "windows_process_extensions_proc_thread_attributes", issue = "none")] - unsafe fn process_thread_attribute( + unsafe fn process_thread_attribute( &mut self, attribute: usize, - value: *mut c_void, - size: usize, + value: T, ) -> &mut process::Command; } @@ -196,13 +193,12 @@ impl CommandExt for process::Command { self.as_inner_mut().raw_arg(raw_text.as_ref()); self } - unsafe fn process_thread_attribute( + unsafe fn process_thread_attribute( &mut self, attribute: usize, - value: *mut c_void, - size: usize, + value: T, ) -> &mut process::Command { - self.as_inner_mut().process_thread_attribute(attribute, value, size); + self.as_inner_mut().process_thread_attribute(attribute, value); self } } diff --git a/library/std/src/process/tests.rs b/library/std/src/process/tests.rs index cc2e196038462..e35d9d3bcf388 100644 --- a/library/std/src/process/tests.rs +++ b/library/std/src/process/tests.rs @@ -412,17 +412,15 @@ fn test_creation_flags() { fn test_proc_thread_attributes() { use crate::os::windows::io::AsRawHandle; use crate::os::windows::process::CommandExt; - use crate::sys::c::{DWORD_PTR, HANDLE}; + use crate::sys::c::DWORD_PTR; 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 parent = Command::new("cmd.exe").spawn().unwrap(); 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::(), + parent.as_raw_handle() as isize, ); } let mut child = child_cmd.spawn().unwrap(); diff --git a/library/std/src/sys/windows/process.rs b/library/std/src/sys/windows/process.rs index c53a8c6383a17..fa2659673798e 100644 --- a/library/std/src/sys/windows/process.rs +++ b/library/std/src/sys/windows/process.rs @@ -254,8 +254,15 @@ 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 unsafe fn process_thread_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( @@ -876,16 +883,10 @@ impl Drop for ProcThreadAttributeList { } } /// 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, + data: Box, size: usize, } -unsafe impl Send for ProcThreadAttributeValue {} -unsafe impl Sync for ProcThreadAttributeValue {} - fn make_proc_thread_attributes( attributes: &BTreeMap, ) -> io::Result { @@ -902,13 +903,17 @@ fn make_proc_thread_attributes( // 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) { + for (&attribute, value) in attributes.iter().take(count as usize) { + let value_ptr: *const (dyn Send + Sync) = &*value.data as *const _; + // let value_ptr = value_ptr as *const _; + let value_ptr = value_ptr as *const c_void; + let value_ptr = value_ptr as *mut c_void; cvt(unsafe { c::UpdateProcThreadAttribute( attribute_list.0.as_mut_ptr().cast(), 0, attribute, - value.ptr, + value_ptr, value.size, ptr::null_mut(), ptr::null_mut(),