Skip to content

Commit

Permalink
Added option to set ProcThreadAttributes for Windows processes
Browse files Browse the repository at this point in the history
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 <t.ruckinger@gmail.com>
  • Loading branch information
michaelvanstraten and TyPR124 committed Aug 25, 2023
1 parent d7e7510 commit edefa8b
Show file tree
Hide file tree
Showing 5 changed files with 306 additions and 2 deletions.
69 changes: 69 additions & 0 deletions library/std/src/os/windows/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<T: Copy + Send + Sync + 'static>(
&mut self,
attribute: usize,
value: T,
) -> &mut process::Command;
}

#[stable(feature = "windows_process_extensions", since = "1.16.0")]
Expand Down Expand Up @@ -219,6 +279,15 @@ impl CommandExt for process::Command {
let _ = always_async;
self
}

unsafe fn raw_attribute<T: Copy + Send + Sync + 'static>(
&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")]
Expand Down
85 changes: 85 additions & 0 deletions library/std/src/process/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<PROCESSENTRY32W>() 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: Send + Sync>(_: T) {}
Expand Down
5 changes: 5 additions & 0 deletions library/std/src/sys/windows/c/windows_sys.lst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
37 changes: 37 additions & 0 deletions library/std/src/sys/windows/c/windows_sys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
Loading

0 comments on commit edefa8b

Please sign in to comment.