Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Abstract ProcThreadAttributeList into its own struct #123604

Merged
merged 1 commit into from
Dec 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
328 changes: 273 additions & 55 deletions library/std/src/os/windows/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,14 @@

#![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::{
AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle, OwnedHandle, RawHandle,
};
use crate::sealed::Sealed;
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")]
impl FromRawHandle for process::Stdio {
Expand Down Expand Up @@ -295,41 +296,25 @@ pub trait CommandExt: Sealed {
#[unstable(feature = "windows_process_extensions_async_pipes", issue = "98289")]
fn async_pipes(&mut self, always_async: bool) -> &mut process::Command;

/// Set a raw attribute on the command, providing extended configuration options for Windows
/// processes.
/// Executes the command as a child process with the given
/// [`ProcThreadAttributeList`], returning a handle to it.
///
/// 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` 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
/// This method enables the customization of attributes for the spawned
/// child process on Windows systems.
/// Attributes offer extended configurability for process creation,
/// but their usage can be intricate and potentially unsafe.
///
/// # 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.
/// By default, stdin, stdout, and stderr are inherited from the parent
/// process.
///
/// # 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::os::windows::io::AsRawHandle;
/// use std::os::windows::process::{CommandExt, ProcThreadAttributeList};
/// use std::process::Command;
///
/// # struct ProcessDropGuard(std::process::Child);
Expand All @@ -338,36 +323,27 @@ pub trait CommandExt: Sealed {
/// # let _ = self.0.kill();
/// # }
/// # }
///
/// #
/// let parent = Command::new("cmd").spawn()?;
///
/// let mut child_cmd = Command::new("cmd");
/// let parent_process_handle = parent.as_raw_handle();
/// # let parent = ProcessDropGuard(parent);
///
/// 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();
///
/// unsafe {
/// child_cmd.raw_attribute(PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, parent.as_raw_handle() as isize);
/// }
/// let mut child = Command::new("cmd").spawn_with_attributes(&attribute_list)?;
/// #
/// # 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<T: Copy + Send + Sync + 'static>(
fn spawn_with_attributes(
&mut self,
attribute: usize,
value: T,
) -> &mut process::Command;
attribute_list: &ProcThreadAttributeList<'_>,
) -> io::Result<process::Child>;
}

#[stable(feature = "windows_process_extensions", since = "1.16.0")]
Expand Down Expand Up @@ -401,13 +377,13 @@ impl CommandExt for process::Command {
self
}

unsafe fn raw_attribute<T: Copy + Send + Sync + 'static>(
fn spawn_with_attributes(
&mut self,
attribute: usize,
value: T,
) -> &mut process::Command {
unsafe { self.as_inner_mut().raw_attribute(attribute, value) };
self
attribute_list: &ProcThreadAttributeList<'_>,
) -> io::Result<process::Child> {
self.as_inner_mut()
.spawn_with_attributes(sys::process::Stdio::Inherit, true, Some(attribute_list))
.map(process::Child::from_inner)
}
}

Expand Down Expand Up @@ -447,3 +423,245 @@ impl ExitCodeExt for process::ExitCode {
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 ()>,
michaelvanstraten marked this conversation as resolved.
Show resolved Hide resolved
}

#[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,
}
Loading
Loading