user0/rust
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Added option to set ProcThreadAttributes for Windows processes

Browse source 
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>
This commit is contained in:
Michael van Straten 2023-08-25 16:19:16 +02:00
parent d7e751006c
commit edefa8b105
5 changed files with 306 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
library/std/src/sys/windows/c/windows_sys.lst
View file

@ -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

37
library/std/src/sys/windows/c/windows_sys.rs
View file

@ -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,

112
library/std/src/sys/windows/process.rs
View file

@ -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<Stdio>,
stderr: Option<Stdio>,
force_quotes_enabled: bool,
proc_thread_attributes: BTreeMap<usize, ProcThreadAttributeValue>,
}
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<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(
&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::<c::STARTUPINFOW>() 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::<c::STARTUPINFOEXW>() 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::<c::STARTUPINFOW> 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<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_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>,
}