Skip to content

Commit db52667

Browse files
MikailBagMikailBag
committed
refactor attr list handling
1 parent 3230496 commit db52667

File tree

2 files changed

+136
-78
lines changed

2 files changed

+136
-78
lines changed

‎src/windows/spawn.rs

Lines changed: 24 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
mod attr_list;
2+
13
use crate::{
24
windows::{
35
pipe::{self, ReadPipe, WritePipe},
@@ -6,6 +8,7 @@ use crate::{
68
},
79
InputSpecificationData, OutputSpecificationData,
810
};
11+
use attr_list::AttrList;
912
use std::{
1013
ffi::{OsStr, OsString},
1114
mem::size_of,
@@ -15,15 +18,10 @@ use std::{
1518
},
1619
};
1720
use winapi::{
18-
shared::{minwindef::TRUE, winerror::ERROR_INSUFFICIENT_BUFFER},
21+
shared::minwindef::TRUE,
1922
um::{
20-
errhandlingapi::GetLastError,
2123
handleapi::INVALID_HANDLE_VALUE,
22-
minwinbase::SECURITY_ATTRIBUTES,
23-
processthreadsapi::{
24-
CreateProcessW, DeleteProcThreadAttributeList, InitializeProcThreadAttributeList,
25-
UpdateProcThreadAttribute, PROCESS_INFORMATION, PROC_THREAD_ATTRIBUTE_LIST,
26-
},
24+
processthreadsapi::{CreateProcessW, PROCESS_INFORMATION},
2725
winbase::{
2826
CreateFileMappingA, CREATE_UNICODE_ENVIRONMENT, EXTENDED_STARTUPINFO_PRESENT,
2927
STARTF_USESTDHANDLES, STARTUPINFOEXW,
@@ -147,76 +145,22 @@ pub(in crate::windows) struct ChildParams {
147145
// TODO: upstream to winapi: https://github.com/retep998/winapi-rs/pull/933/
148146
const MAGIC_PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES: usize = 131081;
149147

150-
struct AlignedMemBlock(*mut u8, usize);
151-
152-
impl AlignedMemBlock {
153-
fn layout(cnt: usize) -> std::alloc::Layout {
154-
assert!(cnt > 0);
155-
std::alloc::Layout::from_size_align(cnt, 8).unwrap()
156-
}
157-
158-
fn new(cnt: usize) -> AlignedMemBlock {
159-
let ptr = unsafe { std::alloc::alloc_zeroed(Self::layout(cnt)) };
160-
AlignedMemBlock(ptr, cnt)
161-
}
162-
163-
fn ptr(&self) -> *mut u8 {
164-
self.0
165-
}
166-
}
167-
168-
impl Drop for AlignedMemBlock {
169-
fn drop(&mut self) {
170-
unsafe {
171-
std::alloc::dealloc(self.0, Self::layout(self.1));
172-
}
173-
}
174-
}
175-
176148
pub(in crate::windows) fn spawn(
177149
sandbox: &WindowsSandbox,
178150
stdio: Stdio,
179151
params: ChildParams,
180152
) -> Result<PROCESS_INFORMATION, Error> {
181-
let proc_thread_attr_list_storage;
182153
let mut security_capabilities;
154+
let mut proc_thread_attr_list = AttrList::new(1)?;
183155
let mut startup_info = unsafe {
184156
let mut startup_info: STARTUPINFOEXW = std::mem::zeroed();
185-
let mut proc_thread_attr_list_len = 0;
186-
{
187-
InitializeProcThreadAttributeList(
188-
std::ptr::null_mut(),
189-
// we need only one attribute: security capabilities.
190-
1,
191-
0,
192-
&mut proc_thread_attr_list_len,
193-
);
194-
if GetLastError() != ERROR_INSUFFICIENT_BUFFER {
195-
return Err(Error::last());
196-
}
197-
}
198-
proc_thread_attr_list_storage = AlignedMemBlock::new(proc_thread_attr_list_len);
199-
let proc_thread_attr_list = proc_thread_attr_list_storage.ptr();
200-
startup_info.lpAttributeList = proc_thread_attr_list.cast();
201-
Cvt::nonzero(InitializeProcThreadAttributeList(
202-
startup_info.lpAttributeList,
203-
1,
204-
0,
205-
&mut proc_thread_attr_list_len,
206-
))?;
157+
207158
security_capabilities = sandbox.profile.get_security_capabilities();
208-
Cvt::nonzero(UpdateProcThreadAttribute(
209-
startup_info.lpAttributeList,
210-
// reserved
211-
0,
159+
160+
proc_thread_attr_list.add_attr::<SECURITY_CAPABILITIES>(
212161
MAGIC_PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES,
213-
(&mut security_capabilities as *mut SECURITY_CAPABILITIES).cast(),
214-
std::mem::size_of::<SECURITY_ATTRIBUTES>(),
215-
// reserved
216-
std::ptr::null_mut(),
217-
// reserved
218-
std::ptr::null_mut(),
219-
))?;
162+
&mut security_capabilities,
163+
)?;
220164

221165
startup_info.StartupInfo.cb = size_of::<STARTUPINFOEXW>() as u32;
222166
startup_info.StartupInfo.dwFlags = STARTF_USESTDHANDLES;
@@ -229,21 +173,24 @@ pub(in crate::windows) fn spawn(
229173
startup_info.StartupInfo.hStdInput = cvt_handle(stdio.stdin);
230174
startup_info.StartupInfo.hStdOutput = cvt_handle(stdio.stdout);
231175
startup_info.StartupInfo.hStdError = cvt_handle(stdio.stderr);
176+
177+
startup_info.lpAttributeList = proc_thread_attr_list.borrow_ptr();
232178
startup_info
233179
};
234180
let creation_flags = CREATE_UNICODE_ENVIRONMENT | EXTENDED_STARTUPINFO_PRESENT;
235181
let mut info: PROCESS_INFORMATION = unsafe { std::mem::zeroed() };
182+
let application_name: Vec<u16> = params.exe.encode_wide().collect();
183+
let mut cmd_line = application_name.clone();
184+
for arg in params.argv {
185+
quote_arg(&mut cmd_line, &arg);
186+
}
187+
let (mut env, env_status) = encode_env(&params.env);
188+
if let EncodeEnvResult::Partial = env_status {
189+
tracing::warn!("skipped zero chars in provided environment");
190+
}
191+
let cwd: Vec<u16> = params.cwd.encode_wide().collect();
192+
236193
unsafe {
237-
let application_name: Vec<u16> = params.exe.encode_wide().collect();
238-
let mut cmd_line = application_name.clone();
239-
for arg in params.argv {
240-
quote_arg(&mut cmd_line, &arg);
241-
}
242-
let (mut env, env_status) = encode_env(&params.env);
243-
if let EncodeEnvResult::Partial = env_status {
244-
tracing::warn!("skipped zero chars in provided environment");
245-
}
246-
let cwd: Vec<u16> = params.cwd.encode_wide().collect();
247194
Cvt::nonzero(CreateProcessW(
248195
application_name.as_ptr(),
249196
cmd_line.as_mut_ptr(),
@@ -259,7 +206,6 @@ pub(in crate::windows) fn spawn(
259206
(&mut startup_info as *mut STARTUPINFOEXW).cast(),
260207
&mut info,
261208
))?;
262-
DeleteProcThreadAttributeList(startup_info.lpAttributeList);
263209
}
264210
Ok(info)
265211
}

‎src/windows/spawn/attr_list.rs

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
//! RAII wrapper for ProcThreadAttrList
2+
use std::marker::PhantomData;
3+
4+
use crate::windows::{Cvt, Error};
5+
use winapi::{
6+
shared::winerror::ERROR_INSUFFICIENT_BUFFER,
7+
um::{
8+
errhandlingapi::GetLastError,
9+
processthreadsapi::{
10+
DeleteProcThreadAttributeList, InitializeProcThreadAttributeList,
11+
UpdateProcThreadAttribute, PROC_THREAD_ATTRIBUTE_LIST,
12+
},
13+
},
14+
};
15+
16+
pub struct AttrList<'a> {
17+
storage: AlignedMemBlock,
18+
cap: usize,
19+
len: usize,
20+
phantom: PhantomData<&'a mut &'a mut ()>,
21+
}
22+
23+
impl<'a> AttrList<'a> {
24+
pub fn new(capacity: usize) -> Result<Self, Error> {
25+
let mut proc_thread_attr_list_len = 0;
26+
unsafe {
27+
InitializeProcThreadAttributeList(
28+
std::ptr::null_mut(),
29+
// we need only one attribute: security capabilities.
30+
1,
31+
0,
32+
&mut proc_thread_attr_list_len,
33+
);
34+
if GetLastError() != ERROR_INSUFFICIENT_BUFFER {
35+
return Err(Error::last());
36+
}
37+
}
38+
let storage = AlignedMemBlock::new(proc_thread_attr_list_len);
39+
unsafe {
40+
Cvt::nonzero(InitializeProcThreadAttributeList(
41+
storage.ptr().cast(),
42+
1,
43+
0,
44+
&mut proc_thread_attr_list_len,
45+
))?;
46+
}
47+
48+
Ok(AttrList {
49+
storage,
50+
cap: capacity,
51+
len: 0,
52+
phantom: PhantomData,
53+
})
54+
}
55+
56+
pub fn add_attr<T>(&mut self, attr_name: usize, attr_val: &'a mut T) -> Result<(), Error> {
57+
assert!(self.len < self.cap);
58+
unsafe {
59+
Cvt::nonzero(UpdateProcThreadAttribute(
60+
self.storage.ptr().cast(),
61+
// reserved
62+
0,
63+
attr_name,
64+
(attr_val as *mut T).cast(),
65+
std::mem::size_of::<T>(),
66+
// reserved
67+
std::ptr::null_mut(),
68+
// reserved
69+
std::ptr::null_mut(),
70+
))?;
71+
}
72+
self.len += 1;
73+
Ok(())
74+
}
75+
76+
pub fn borrow_ptr(&self) -> *mut PROC_THREAD_ATTRIBUTE_LIST {
77+
assert_eq!(self.len, self.cap);
78+
self.storage.ptr().cast()
79+
}
80+
}
81+
82+
impl<'a> Drop for AttrList<'a> {
83+
fn drop(&mut self) {
84+
unsafe { DeleteProcThreadAttributeList(self.storage.ptr().cast()) };
85+
}
86+
}
87+
88+
struct AlignedMemBlock(*mut u8, usize);
89+
90+
impl AlignedMemBlock {
91+
fn layout(cnt: usize) -> std::alloc::Layout {
92+
assert!(cnt > 0);
93+
std::alloc::Layout::from_size_align(cnt, 8).unwrap()
94+
}
95+
96+
fn new(cnt: usize) -> AlignedMemBlock {
97+
let ptr = unsafe { std::alloc::alloc_zeroed(Self::layout(cnt)) };
98+
AlignedMemBlock(ptr, cnt)
99+
}
100+
101+
fn ptr(&self) -> *mut u8 {
102+
self.0
103+
}
104+
}
105+
106+
impl Drop for AlignedMemBlock {
107+
fn drop(&mut self) {
108+
unsafe {
109+
std::alloc::dealloc(self.0, Self::layout(self.1));
110+
}
111+
}
112+
}

0 commit comments

Comments
 (0)