1
Fork 0

Abstract ProcThreadAttributeList into its own struct

This commit is contained in:
Michael van Straten 2024-04-07 21:08:37 +02:00
parent d6f88291f3
commit f371952cde
3 changed files with 297 additions and 154 deletions

View file

@ -4,13 +4,14 @@
#![stable(feature = "process_extensions", since = "1.2.0")] #![stable(feature = "process_extensions", since = "1.2.0")]
use crate::ffi::OsStr; use crate::ffi::{OsStr, c_void};
use crate::mem::MaybeUninit;
use crate::os::windows::io::{ use crate::os::windows::io::{
AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle, OwnedHandle, RawHandle, AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle, OwnedHandle, RawHandle,
}; };
use crate::sealed::Sealed; use crate::sealed::Sealed;
use crate::sys_common::{AsInner, AsInnerMut, FromInner, IntoInner}; use crate::sys_common::{AsInner, AsInnerMut, FromInner, IntoInner};
use crate::{process, sys}; use crate::{io, marker, process, ptr, sys};
#[stable(feature = "process_extensions", since = "1.2.0")] #[stable(feature = "process_extensions", since = "1.2.0")]
impl FromRawHandle for process::Stdio { impl FromRawHandle for process::Stdio {
@ -295,41 +296,25 @@ pub trait CommandExt: Sealed {
#[unstable(feature = "windows_process_extensions_async_pipes", issue = "98289")] #[unstable(feature = "windows_process_extensions_async_pipes", issue = "98289")]
fn async_pipes(&mut self, always_async: bool) -> &mut process::Command; fn async_pipes(&mut self, always_async: bool) -> &mut process::Command;
/// Set a raw attribute on the command, providing extended configuration options for Windows /// Executes the command as a child process with the given
/// processes. /// [`ProcThreadAttributeList`], returning a handle to it.
/// ///
/// This method allows you to specify custom attributes for a child process on Windows systems /// This method enables the customization of attributes for the spawned
/// using raw attribute values. Raw attributes provide extended configurability for process /// child process on Windows systems.
/// creation, but their usage can be complex and potentially unsafe. /// Attributes offer extended configurability for process creation,
/// /// but their usage can be intricate 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` documentation] or the [Win32 API documentation] for detailed information
/// about available attributes and their meanings.
///
/// [`windows-rs` documentation]: https://microsoft.github.io/windows-docs-rs/doc/windows/
/// [Win32 API documentation]: https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-updateprocthreadattribute
/// ///
/// # Note /// # Note
/// ///
/// The maximum number of raw attributes is the value of [`u32::MAX`]. /// By default, stdin, stdout, and stderr are inherited from the parent
/// If this limit is exceeded, the call to [`process::Command::spawn`] will return an `Error` /// process.
/// 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 /// # 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)] /// #![feature(windows_process_extensions_raw_attribute)]
/// use std::os::windows::{process::CommandExt, io::AsRawHandle}; /// use std::os::windows::io::AsRawHandle;
/// use std::os::windows::process::{CommandExt, ProcThreadAttributeList};
/// use std::process::Command; /// use std::process::Command;
/// ///
/// # struct ProcessDropGuard(std::process::Child); /// # struct ProcessDropGuard(std::process::Child);
@ -338,36 +323,27 @@ pub trait CommandExt: Sealed {
/// # let _ = self.0.kill(); /// # 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 = Command::new("cmd").spawn()?;
/// let parent_process_handle = parent.as_raw_handle();
/// # let parent = ProcessDropGuard(parent); /// # let parent = ProcessDropGuard(parent);
/// ///
/// let mut child = child_cmd.spawn()?; /// const PROC_THREAD_ATTRIBUTE_PARENT_PROCESS: usize = 0x00020000;
/// let mut attribute_list = ProcThreadAttributeList::build()
/// .attribute(PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, &parent_process_handle)
/// .finish()
/// .unwrap();
/// ///
/// let mut child = Command::new("cmd").spawn_with_attributes(&attribute_list)?;
/// #
/// # child.kill()?; /// # child.kill()?;
/// # Ok::<(), std::io::Error>(()) /// # 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")] #[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")]
unsafe fn raw_attribute<T: Copy + Send + Sync + 'static>( fn spawn_with_attributes(
&mut self, &mut self,
attribute: usize, attribute_list: &ProcThreadAttributeList<'_>,
value: T, ) -> io::Result<process::Child>;
) -> &mut process::Command;
} }
#[stable(feature = "windows_process_extensions", since = "1.16.0")] #[stable(feature = "windows_process_extensions", since = "1.16.0")]
@ -401,13 +377,13 @@ impl CommandExt for process::Command {
self self
} }
unsafe fn raw_attribute<T: Copy + Send + Sync + 'static>( fn spawn_with_attributes(
&mut self, &mut self,
attribute: usize, attribute_list: &ProcThreadAttributeList<'_>,
value: T, ) -> io::Result<process::Child> {
) -> &mut process::Command { self.as_inner_mut()
unsafe { self.as_inner_mut().raw_attribute(attribute, value) }; .spawn_with_attributes(sys::process::Stdio::Inherit, true, Some(attribute_list))
self .map(process::Child::from_inner)
} }
} }
@ -447,3 +423,245 @@ impl ExitCodeExt for process::ExitCode {
process::ExitCode::from_inner(From::from(raw)) process::ExitCode::from_inner(From::from(raw))
} }
} }
/// A wrapper around windows [`ProcThreadAttributeList`][1].
///
/// [1]: <https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-initializeprocthreadattributelist>
#[derive(Debug)]
#[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")]
pub struct ProcThreadAttributeList<'a> {
attribute_list: Box<[MaybeUninit<u8>]>,
_lifetime_marker: marker::PhantomData<&'a ()>,
}
#[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")]
impl<'a> ProcThreadAttributeList<'a> {
/// Creates a new builder for constructing a [`ProcThreadAttributeList`].
pub fn build() -> ProcThreadAttributeListBuilder<'a> {
ProcThreadAttributeListBuilder::new()
}
/// Returns a pointer to the underling attribute list.
#[doc(hidden)]
pub fn as_ptr(&self) -> *const MaybeUninit<u8> {
self.attribute_list.as_ptr()
}
}
#[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")]
impl<'a> Drop for ProcThreadAttributeList<'a> {
/// Deletes the attribute list.
///
/// This method calls [`DeleteProcThreadAttributeList`][1] to delete the
/// underlying attribute list.
///
/// [1]: <https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-deleteprocthreadattributelist>
fn drop(&mut self) {
let lp_attribute_list = self.attribute_list.as_mut_ptr().cast::<c_void>();
unsafe { sys::c::DeleteProcThreadAttributeList(lp_attribute_list) }
}
}
/// Builder for constructing a [`ProcThreadAttributeList`].
#[derive(Clone, Debug)]
#[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")]
pub struct ProcThreadAttributeListBuilder<'a> {
attributes: alloc::collections::BTreeMap<usize, ProcThreadAttributeValue>,
_lifetime_marker: marker::PhantomData<&'a ()>,
}
#[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")]
impl<'a> ProcThreadAttributeListBuilder<'a> {
fn new() -> Self {
ProcThreadAttributeListBuilder {
attributes: alloc::collections::BTreeMap::new(),
_lifetime_marker: marker::PhantomData,
}
}
/// Sets an attribute on the attribute list.
///
/// 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 documentation][1] for a list of valid attributes.
///
/// # Note
///
/// The maximum number of attributes is the value of [`u32::MAX`]. If this
/// limit is exceeded, the call to [`Self::finish`] will return an `Error`
/// indicating that the maximum number of attributes has been exceeded.
///
/// # Safety Note
///
/// Remember that improper use of attributes can lead to undefined behavior
/// or security vulnerabilities. Always consult the documentation and ensure
/// proper attribute values are used.
///
/// [1]: <https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-updateprocthreadattribute#parameters>
pub fn attribute<T>(self, attribute: usize, value: &'a T) -> Self {
unsafe {
self.raw_attribute(
attribute,
ptr::addr_of!(*value).cast::<c_void>(),
crate::mem::size_of::<T>(),
)
}
}
/// Sets a raw attribute on the attribute list.
///
/// This function is useful for setting attributes with pointers or sizes
/// that cannot be derived directly from their values.
///
/// # Safety
///
/// This function is marked as `unsafe` because it deals with raw pointers
/// and sizes. It is the responsibility of the caller to ensure the value
/// lives longer than the resulting [`ProcThreadAttributeList`] as well as
/// the validity of the size parameter.
///
/// # Example
///
/// ```
/// #![feature(windows_process_extensions_raw_attribute)]
/// use std::ffi::c_void;
/// use std::os::windows::process::{CommandExt, ProcThreadAttributeList};
/// use std::os::windows::raw::HANDLE;
/// use std::process::Command;
///
/// #[repr(C)]
/// pub struct COORD {
/// pub X: i16,
/// pub Y: i16,
/// }
///
/// extern "system" {
/// fn CreatePipe(
/// hreadpipe: *mut HANDLE,
/// hwritepipe: *mut HANDLE,
/// lppipeattributes: *const c_void,
/// nsize: u32,
/// ) -> i32;
/// fn CreatePseudoConsole(
/// size: COORD,
/// hinput: HANDLE,
/// houtput: HANDLE,
/// dwflags: u32,
/// phpc: *mut isize,
/// ) -> i32;
/// fn CloseHandle(hobject: HANDLE) -> i32;
/// }
///
/// let [mut input_read_side, mut output_write_side, mut output_read_side, mut input_write_side] =
/// [unsafe { std::mem::zeroed::<HANDLE>() }; 4];
///
/// unsafe {
/// CreatePipe(&mut input_read_side, &mut input_write_side, std::ptr::null(), 0);
/// CreatePipe(&mut output_read_side, &mut output_write_side, std::ptr::null(), 0);
/// }
///
/// let size = COORD { X: 60, Y: 40 };
/// let mut h_pc = unsafe { std::mem::zeroed() };
/// unsafe { CreatePseudoConsole(size, input_read_side, output_write_side, 0, &mut h_pc) };
///
/// unsafe { CloseHandle(input_read_side) };
/// unsafe { CloseHandle(output_write_side) };
///
/// const PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE: usize = 131094;
///
/// let attribute_list = unsafe {
/// ProcThreadAttributeList::build()
/// .raw_attribute(
/// PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE,
/// h_pc as *const c_void,
/// std::mem::size_of::<isize>(),
/// )
/// .finish()?
/// };
///
/// let mut child = Command::new("cmd").spawn_with_attributes(&attribute_list)?;
/// #
/// # child.kill()?;
/// # Ok::<(), std::io::Error>(())
/// ```
pub unsafe fn raw_attribute<T>(
mut self,
attribute: usize,
value_ptr: *const T,
value_size: usize,
) -> Self {
self.attributes.insert(attribute, ProcThreadAttributeValue {
ptr: value_ptr.cast::<c_void>(),
size: value_size,
});
self
}
/// Finalizes the construction of the `ProcThreadAttributeList`.
///
/// # Errors
///
/// Returns an error if the maximum number of attributes is exceeded
/// or if there is an I/O error during initialization.
pub fn finish(&self) -> io::Result<ProcThreadAttributeList<'a>> {
// 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) = self.attributes.len().try_into() else {
return Err(io::const_error!(
io::ErrorKind::InvalidInput,
"maximum number of ProcThreadAttributes exceeded",
));
};
unsafe {
sys::c::InitializeProcThreadAttributeList(
ptr::null_mut(),
attribute_count,
0,
&mut required_size,
)
};
let mut attribute_list = vec![MaybeUninit::uninit(); required_size].into_boxed_slice();
// Once we've allocated the necessary memory, it's safe to invoke
// `InitializeProcThreadAttributeList` to properly initialize the list.
sys::cvt(unsafe {
sys::c::InitializeProcThreadAttributeList(
attribute_list.as_mut_ptr().cast::<c_void>(),
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 self.attributes.iter().take(attribute_count as usize) {
sys::cvt(unsafe {
sys::c::UpdateProcThreadAttribute(
attribute_list.as_mut_ptr().cast::<c_void>(),
0,
attribute,
value.ptr,
value.size,
ptr::null_mut(),
ptr::null_mut(),
)
})?;
}
Ok(ProcThreadAttributeList { attribute_list, _lifetime_marker: marker::PhantomData })
}
}
/// Wrapper around the value data to be used as a Process Thread Attribute.
#[derive(Clone, Debug)]
struct ProcThreadAttributeValue {
ptr: *const c_void,
size: usize,
}

View file

@ -450,7 +450,7 @@ fn test_creation_flags() {
fn test_proc_thread_attributes() { fn test_proc_thread_attributes() {
use crate::mem; use crate::mem;
use crate::os::windows::io::AsRawHandle; use crate::os::windows::io::AsRawHandle;
use crate::os::windows::process::CommandExt; use crate::os::windows::process::{CommandExt, ProcThreadAttributeList};
use crate::sys::c::{BOOL, CloseHandle, HANDLE}; use crate::sys::c::{BOOL, CloseHandle, HANDLE};
use crate::sys::cvt; use crate::sys::cvt;
@ -490,12 +490,14 @@ fn test_proc_thread_attributes() {
let mut child_cmd = Command::new("cmd"); let mut child_cmd = Command::new("cmd");
unsafe { let parent_process_handle = parent.0.as_raw_handle();
child_cmd
.raw_attribute(PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, parent.0.as_raw_handle() as isize);
}
let child = ProcessDropGuard(child_cmd.spawn().unwrap()); let mut attribute_list = ProcThreadAttributeList::build()
.attribute(PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, &parent_process_handle)
.finish()
.unwrap();
let child = ProcessDropGuard(child_cmd.spawn_with_attributes(&mut attribute_list).unwrap());
let h_snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) }; let h_snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };

View file

@ -10,10 +10,10 @@ use crate::collections::BTreeMap;
use crate::env::consts::{EXE_EXTENSION, EXE_SUFFIX}; use crate::env::consts::{EXE_EXTENSION, EXE_SUFFIX};
use crate::ffi::{OsStr, OsString}; use crate::ffi::{OsStr, OsString};
use crate::io::{self, Error, ErrorKind}; use crate::io::{self, Error, ErrorKind};
use crate::mem::MaybeUninit;
use crate::num::NonZero; use crate::num::NonZero;
use crate::os::windows::ffi::{OsStrExt, OsStringExt}; use crate::os::windows::ffi::{OsStrExt, OsStringExt};
use crate::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle}; use crate::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle};
use crate::os::windows::process::ProcThreadAttributeList;
use crate::path::{Path, PathBuf}; use crate::path::{Path, PathBuf};
use crate::sync::Mutex; use crate::sync::Mutex;
use crate::sys::args::{self, Arg}; use crate::sys::args::{self, Arg};
@ -162,7 +162,6 @@ pub struct Command {
stdout: Option<Stdio>, stdout: Option<Stdio>,
stderr: Option<Stdio>, stderr: Option<Stdio>,
force_quotes_enabled: bool, force_quotes_enabled: bool,
proc_thread_attributes: BTreeMap<usize, ProcThreadAttributeValue>,
} }
pub enum Stdio { pub enum Stdio {
@ -194,7 +193,6 @@ impl Command {
stdout: None, stdout: None,
stderr: None, stderr: None,
force_quotes_enabled: false, force_quotes_enabled: false,
proc_thread_attributes: Default::default(),
} }
} }
@ -248,21 +246,19 @@ impl Command {
self.cwd.as_ref().map(Path::new) self.cwd.as_ref().map(Path::new)
} }
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( pub fn spawn(
&mut self, &mut self,
default: Stdio, default: Stdio,
needs_stdin: bool, needs_stdin: bool,
) -> io::Result<(Process, StdioPipes)> {
self.spawn_with_attributes(default, needs_stdin, None)
}
pub fn spawn_with_attributes(
&mut self,
default: Stdio,
needs_stdin: bool,
proc_thread_attribute_list: Option<&ProcThreadAttributeList<'_>>,
) -> io::Result<(Process, StdioPipes)> { ) -> io::Result<(Process, StdioPipes)> {
let maybe_env = self.env.capture_if_changed(); let maybe_env = self.env.capture_if_changed();
@ -355,18 +351,18 @@ impl Command {
let si_ptr: *mut c::STARTUPINFOW; let si_ptr: *mut c::STARTUPINFOW;
let mut proc_thread_attribute_list;
let mut si_ex; let mut si_ex;
if !self.proc_thread_attributes.is_empty() { if let Some(proc_thread_attribute_list) = proc_thread_attribute_list {
si.cb = mem::size_of::<c::STARTUPINFOEXW>() as u32; si.cb = mem::size_of::<c::STARTUPINFOEXW>() as u32;
flags |= c::EXTENDED_STARTUPINFO_PRESENT; flags |= c::EXTENDED_STARTUPINFO_PRESENT;
proc_thread_attribute_list =
make_proc_thread_attribute_list(&self.proc_thread_attributes)?;
si_ex = c::STARTUPINFOEXW { si_ex = c::STARTUPINFOEXW {
StartupInfo: si, StartupInfo: si,
lpAttributeList: proc_thread_attribute_list.0.as_mut_ptr() as _, // SAFETY: Casting this `*const` pointer to a `*mut` pointer is "safe"
// here because windows does not internally mutate the attribute list.
// Ideally this should be reflected in the interface of the `windows-sys` crate.
lpAttributeList: proc_thread_attribute_list.as_ptr().cast::<c_void>().cast_mut(),
}; };
si_ptr = (&raw mut si_ex) as _; si_ptr = (&raw mut si_ex) as _;
} else { } else {
@ -896,79 +892,6 @@ 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_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].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 = (&raw const *value.data) 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> { pub struct CommandArgs<'a> {
iter: crate::slice::Iter<'a, Arg>, iter: crate::slice::Iter<'a, Arg>,
} }