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.

Sign up for GitHub

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

Add ability to spawn Windows process with Proc Thread Attributes #88193

Closed
wants to merge 2 commits into from
Closed

Add ability to spawn Windows process with Proc Thread Attributes #88193

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0475e2f
add windows ext method for proc thread attributes
TyPR124 Aug 20, 2021
d20fc4e
Take generic value type
TyPR124 Sep 9, 2021
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
22 changes: 22 additions & 0 deletions library/std/src/os/windows/process.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,20 @@ pub trait CommandExt: Sealed {
/// `CommandLineToArgvW` escaping rules.
#[unstable(feature = "windows_process_extensions_raw_arg", issue = "29494")]
fn raw_arg<S: AsRef<OsStr>>(&mut self, text_to_append_as_is: S) -> &mut process::Command;

/// Attach the specified proc thread attribute. Multiple attributes may be attached.
///
/// # Safety
///
/// - 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<T: Copy + Send + Sync + 'static>(
&mut self,
attribute: usize,
value: T,
) -> &mut process::Command;
}

#[stable(feature = "windows_process_extensions", since = "1.16.0")]
Expand All @@ -179,4 +193,12 @@ impl CommandExt for process::Command {
self.as_inner_mut().raw_arg(raw_text.as_ref());
self
}
unsafe fn process_thread_attribute<T: Copy + Send + Sync + 'static>(
&mut self,
attribute: usize,
value: T,
) -> &mut process::Command {
self.as_inner_mut().process_thread_attribute(attribute, value);
self
}
}
34 changes: 34 additions & 0 deletions library/std/src/process/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,6 +405,40 @@ 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;
const PROC_THREAD_ATTRIBUTE_PARENT_PROCESS: DWORD_PTR = 0x00020000;

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,
parent.as_raw_handle() as isize,
);
}
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() {
Expand Down
25 changes: 25 additions & 0 deletions library/std/src/sys/windows/c.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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")]
Expand Down
134 changes: 107 additions & 27 deletions library/std/src/sys/windows/process.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -202,6 +204,7 @@ impl Command {
stdout: None,
stderr: None,
force_quotes_enabled: false,
proc_thread_attributes: Default::default(),
}
}

Expand Down Expand Up @@ -251,6 +254,16 @@ impl Command {
pub fn get_current_dir(&self) -> Option<&Path> {
self.cwd.as_ref().map(|cwd| Path::new(cwd))
}
pub unsafe fn process_thread_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,
Expand All @@ -259,9 +272,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::<c::STARTUPINFO>() as c::DWORD;
si.dwFlags = c::STARTF_USESTDHANDLES;
let mut si = zeroed_startupinfo_ex();
si.StartupInfo.cb = mem::size_of::<c::STARTUPINFO>() 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())
Expand Down Expand Up @@ -305,9 +318,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::<c::STARTUPINFOEX>() as u32;

buf.0.as_mut_ptr().cast()
});

let program = to_u16s(&program)?;
unsafe {
Expand All @@ -320,7 +348,7 @@ impl Command {
flags,
envp,
dirp,
&mut si,
&mut si as *mut _ as *mut _,
&mut pi,
))
}?;
Expand Down Expand Up @@ -687,26 +715,29 @@ impl From<u8> 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(),
}
}

Expand Down Expand Up @@ -843,6 +874,55 @@ fn make_dirp(d: Option<&OsString>) -> io::Result<(*const u16, Vec<u16>)> {
}
}

/// 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<Vec<u8>>);
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.
struct ProcThreadAttributeValue {
data: Box<dyn Send + Sync>,
size: usize,
}
fn make_proc_thread_attributes(
attributes: &BTreeMap<usize, ProcThreadAttributeValue>,
) -> io::Result<ProcThreadAttributeList> {
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));
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this need to be Pin? Surely the Vec itself can be moved around freely?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Pin prevents us from calling (for example) vec.push() or doing anything that may re-allocate (and move) the underlying data. I'm not 100% sure it is necessary to not move the data, as I don't know the internal workings of these Initialize/Update/Delete functions, but it seemed the right choice to me since there is at least no good reason to move or grow the data after that point.

The Pin doesn't prevent us from moving the Vec itself (after all, we must move it out of this function to return it), just the data it points to.

Copy link
Member

@bjorn3 bjorn3 Aug 21, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[u8] is Unpin, so Pin<Vec<u8>> so both Pin::into_inner and DerefMut are implemented allowing you to move the inner bytes anyway.

https://doc.rust-lang.org/stable/std/pin/struct.Pin.html#method.into_inner

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think perhaps a boxed slice would be enough to prevent accidentally reallocating? Or it might be a good idea to fully contain the owned data. So have an API like:

struct ProcThreadAttributeList(Vec<u8>); // or Box<[T]>
impl ProcThreadAttributeList {
    fn new(attribute_count: u32) -> io::Result<Self> { ... }
    unsafe fn update(&mut self, attribute: usize, value_ptr: *mut c_void, value_size: u32) -> io::Result<()> { ... }
}
impl Drop for ProcThreadAttributeList { ... }

That way there is no reason to touch the inner data.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ChrisDenton I could do that, just seemed like overkill to me when I was writing it.

@bjorn3 You could copy the bytes, or overwrite them, but you can't move them. let mut v = Pin::new(vec![]); v.push(()); does not compile. Unless I'm just being really dumb right now. You are right about Pin::into_inner but calling that explicitly indicates it is now okay to move the underlying data (which it might be anyway, but, again, I'm not certain).

// 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) {
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.size,
ptr::null_mut(),
ptr::null_mut(),
)
})?;
}
Ok(attribute_list)
}

pub struct CommandArgs<'a> {
iter: crate::slice::Iter<'a, Arg>,
}
Expand Down