8000 _overlapped · RustPython/RustPython@427ec50 · GitHub
[go: up one dir, main page]
Skip to content

Commit 427ec50

Browse files
youknowoneyouknowone
committed
_overlapped
1 parent 52ce150 commit 427ec50

File tree

3 files changed

+376
-4
lines changed

3 files changed

+376
-4
lines changed

‎stdlib/src/lib.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,8 @@ mod posixsubprocess;
5454
// libc is missing constants on redox
5555
#[cfg(all(unix, not(any(target_os = "android", target_os = "redox"))))]
5656
mod grp;
57+
#[cfg(windows)]
58+
mod overlapped;
5759
#[cfg(all(unix, not(target_os = "redox")))]
5860
mod resource;
5961
#[cfg(target_os = "macos")]
@@ -152,6 +154,10 @@ pub fn get_module_inits() -> impl Iterator<Item = (Cow<'static, str>, StdlibInit
152154
{
153155
"_bz2" => bz2::make_module,
154156
}
157+
#[cfg(windows)]
158+
{
159+
"_overlapped" => overlapped::make_module,
160+
}
155161
// Unix-only
156162
#[cfg(unix)]
157163
{

‎stdlib/src/overlapped.rs

Lines changed: 345 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,345 @@
1+
pub(crate) use _overlapped::make_module;
2+
3+
#[allow(non_snake_case)]
4+
#[pymodule]
5+
mod _overlapped {
6+
// straight-forward port of Modules/overlapped.c
7+
8+
use crate::vm::{
9+
builtins::{PyBaseExceptionRef, PyBytesRef, PyTypeRef},
10+
common::lock::PyMutex,
11+
convert::{ToPyException, ToPyObject},
12+
protocol::PyBuffer,
13+
stdlib::os::errno_err,
14+
types::Constructor,
15+
Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
16+
};
17+
use windows_sys::Win32::{
18+
Foundation::{self, GetLastError, HANDLE},
19+
Networking::WinSock::SOCKADDR_IN6,
20+
System::IO::OVERLAPPED,
21+
};
22+
23+
#[pyattr]
24+
use windows_sys::Win32::{
25+
Foundation::{
26+
ERROR_IO_PENDING, ERROR_NETNAME_DELETED, ERROR_OPERATION_ABORTED, ERROR_PIPE_BUSY,
27+
ERROR_PORT_UNREACHABLE, ERROR_SEM_TIMEOUT, INVALID_HANDLE_VALUE,
28+
},
29+
Networking::WinSock::{
30+
SO_UPDATE_ACCEPT_CONTEXT, SO_UPDATE_CONNECT_CONTEXT, TF_REUSE_SOCKET,
31+
},
32+
System::Threading::INFINITE,
33+
};
34+
#[pyattr]
35+
const NULL: isize = 0;
36+
37+
#[pyattr]
38+
#[pyclass(name)]
39+
#[derive(PyPayload)]
40+
struct Overlapped {
41+
inner: PyMutex<OverlappedInner>,
42+
}
43+
44+
struct OverlappedInner {
45+
overlapped: OVERLAPPED,
46+
handle: HANDLE,
47+
error: u32,
48+
data: OverlappedData,
49+
}
50+
51+
unsafe impl Sync for OverlappedInner {}
52+
unsafe impl Send for OverlappedInner {}
53+
54+
impl std::fmt::Debug for Overlapped {
55+
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
56+
let zelf = self.inner.lock();
57+
f.debug_struct("Overlapped")
58+
// .field("overlapped", &(self.overlapped as *const _ as usize))
59+
.field("handle", &zelf.handle)
60+
.field("error", &zelf.error)
61+
.field("data", &zelf.data)
62+
.finish()
63+
}
64+
}
65+
66+
#[allow(dead_code)] // TODO: remove when done
67+
#[derive(Debug)]
68+
enum OverlappedData {
69+
None,
70+
NotStarted,
71+
Read(PyBytesRef),
72+
ReadInto(PyBuffer),
73+
Write(PyBuffer),
74+
Accept(PyObjectRef),
75+
Connect,
76+
Disconnect,
77+
ConnectNamedPipe,
78+
WaitNamedPipeAndConnect,
79+
TransmitFile,
80+
ReadFrom(OverlappedReadFrom),
81+
WriteTo(PyBuffer),
82+
ReadFromInto(OverlappedReadFromInto),
83+
}
84+
85+
struct OverlappedReadFrom {
86+
// A (buffer, (host, port)) tuple
87+
result: PyObjectRef,
88+
// The actual read buffer
89+
allocated_buffer: PyObjectRef,
90+
#[allow(dead_code)]
91+
address: SOCKADDR_IN6, // TODO: remove when done
92+
address_length: libc::c_int,
93+
}
94+
95+
impl std::fmt::Debug for OverlappedReadFrom {
96+
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
97+
f.debug_struct("OverlappedReadFrom")
98+
.field("result", &self.result)
99+
.field("allocated_buffer", &self.allocated_buffer)
100+
// .field("address", &self.address)
101+
.field("address_length", &self.address_length)
102+
.finish()
103+
}
104+
}
105+
106+
struct OverlappedReadFromInto {
107+
// A (number of bytes read, (host, port)) tuple
108+
result: PyObjectRef,
109+
/* Buffer passed by the user */
110+
user_buffer: PyBuffer,
111+
#[allow(dead_code)]
112+
address: SOCKADDR_IN6, // TODO: remove when done
113+
address_length: libc::c_int,
114+
}
115+
116+
impl std::fmt::Debug for OverlappedReadFromInto {
117+
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
118+
f.debug_struct("OverlappedReadFromInto")
119+
.field("result", &self.result)
120+
.field("user_buffer", &self.user_buffer)
121+
// .field("address", &self.address)
122+
.field("address_length", &self.address_length)
123+
.finish()
124+
}
125+
}
126+
127+
fn mark_as_completed(ov: &mut OVERLAPPED) {
128+
ov.Internal = 0;
129< 1C6A /td>+
if ov.hEvent != 0 {
130+
unsafe { windows_sys::Win32::System::Threading::SetEvent(ov.hEvent) };
131+
}
132+
}
133+
134+
fn from_windows_err(err: u32, vm: &VirtualMachine) -> PyBaseExceptionRef {
135+
use Foundation::{ERROR_CONNECTION_ABORTED, ERROR_CONNECTION_REFUSED};
136+
debug_assert_ne!(err, 0, "call errno_err instead");
137+
let exc = match err {
138+
ERROR_CONNECTION_REFUSED => vm.ctx.exceptions.connection_refused_error,
139+
ERROR_CONNECTION_ABORTED => vm.ctx.exceptions.connection_aborted_error,
140+
err => return std::io::Error::from_raw_os_error(err as i32).to_pyexception(vm),
141+
};
142+
// TODO: set errno and winerror
143+
vm.new_exception_empty(exc.to_owned())
144+
}
145+
146+
fn HasOverlappedIoCompleted(overlapped: &OVERLAPPED) -> bool {
147+
overlapped.Internal != (Foundation::STATUS_PENDING as usize)
148+
}
149+
150+
#[pyclass(with(Constructor))]
151+
impl Overlapped {
152+
#[pygetset]
153+
fn address(&self, _vm: &VirtualMachine) -> usize {
154+
let inner = self.inner.lock();
155+
&inner.overlapped as *const _ as usize
156+
}
157+
158+
#[pygetset]
159+
fn pending(&self, _vm: &VirtualMachine) -> bool {
160+
let inner = self.inner.lock();
161+
!HasOverlappedIoCompleted(&inner.overlapped)
162+
&& !matches!(inner.data, OverlappedData::NotStarted)
163+
}
164+
165+
fn WSARecv_inner(
166+
inner: &mut OverlappedInner,
167+
handle: HANDLE,
168+
buf: &[u8],
169+
mut flags: u32,
170+
vm: &VirtualMachine,
171+
) -> PyResult {
172+
use windows_sys::Win32::Foundation::{
173+
ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS,
174+
};
175+
176+
let wsabuf = windows_sys::Win32::Networking::WinSock::WSABUF {
177+
buf: buf.as_ptr() as *mut _,
178+
len: buf.len() as _,
179+
};
180+
let mut nread: u32 = 0;
181+
// TODO: optimization with MaybeUninit
182+
let ret = unsafe {
183+
windows_sys::Win32::Networking::WinSock::WSARecv(
184+
handle as _,
185+
&wsabuf,
186+
1,
187+
&mut nread,
188+
&mut flags,
189+
&mut inner.overlapped,
190+
None,
191+
)
192+
};
193+
let err = if ret < 0 {
194+
unsafe { windows_sys::Win32::Networking::WinSock::WSAGetLastError() as u32 }
195+
} else {
196+
Foundation::ERROR_SUCCESS
197+
};
198+
inner.error = err;
199+
match err {
200+
ERROR_BROKEN_PIPE => {
201+
mark_as_completed(&mut inner.overlapped);
202+
Err(from_windows_err(err, vm))
203+
}
204+
ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()),
205+
_ => Err(from_windows_err(err, vm)),
206+
}
207+
}
208+
209+
#[pymethod]
210+
fn WSARecv(
211+
zelf: &Py<Self>,
212+
handle: HANDLE,
213+
size: u32,
214+
flags: u32,
215+
vm: &VirtualMachine,
216+
) -> PyResult {
217+
let mut inner = zelf.inner.lock();
218+
if !matches!(inner.data, OverlappedData::None) {
219+
return Err(vm.new_value_error("operation already attempted".to_owned()));
220+
}
221+
222+
#[cfg(target_pointer_width = "32")]
223+
let size = std::cmp::min(size, std::isize::MAX as _);
224+
225+
let buf = vec![0u8; std::cmp::max(size, 1) as usize];
226+
let buf = vm.ctx.new_bytes(buf);
227+
inner.handle = handle;
228+
229+
let r = Self::WSARecv_inner(&mut inner, handle, buf.as_bytes(), flags, vm);
230+
inner.data = OverlappedData::Read(buf);
231+
r
232+
}
233+
234+
#[pymethod]
235+
fn cancel(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
236+
let inner = zelf.inner.lock();
237+
if matches!(
238+
inner.data,
239+
OverlappedData::NotStarted | OverlappedData::WaitNamedPipeAndConnect
240+
) {
241+
return Ok(());
242+
}
243+
let ret = if !HasOverlappedIoCompleted(&inner.overlapped) {
244+
unsafe {
245+
windows_sys::Win32::System::IO::CancelIoEx(inner.handle, &inner.overlapped)
246+
}
247+
} else {
248+
1
249+
};
250+
// CancelIoEx returns ERROR_NOT_FOUND if the I/O completed in-between
251+
if ret == 0 && unsafe { GetLastError() } != Foundation::ERROR_NOT_FOUND {
252+
return Err(errno_err(vm));
253+
}
254+
Ok(())
255+
}
256+
}
257+
258+
impl Constructor for Overlapped {
259+
type Args = (HANDLE,);
260+
261+
fn py_new(cls: PyTypeRef, (mut event,): Self::Args, vm: &VirtualMachine) -> PyResult {
262+
if event == INVALID_HANDLE_VALUE {
263+
event = unsafe {
264+
windows_sys::Win32::System::Threading::CreateEventA(
265+
std::ptr::null(),
266+
Foundation::TRUE,
267+
Foundation::FALSE,
268+
std::ptr::null(),
269+
)
270+
};
271+
if event == NULL {
272+
return Err(errno_err(vm));
273+
}
274+
}
275+
276+
let mut overlapped: OVERLAPPED = unsafe { std::mem::zeroed() };
277+
if event != NULL {
278+
overlapped.hEvent = event;
279+
}
280+
let inner = OverlappedInner {
281+
overlapped,
282+
handle: NULL,
283+
error: 0,
284+
data: OverlappedData::None,
285+
};
286+
let overlapped = Overlapped {
287+
inner: PyMutex::new(inner),
288+
};
289+
overlapped.into_ref_with_type(vm, cls).map(Into::into)
290+
}
291+
}
292+
293+
#[pyfunction]
294+
fn CreateIoCompletionPort(
295+
handle: HANDLE,
296+
port: HANDLE,
297+
key: usize,
298+
concurrency: u32,
299+
vm: &VirtualMachine,
300+
) -> PyResult<HANDLE> {
301+
let r = unsafe {
302+
windows_sys::Win32::System::IO::CreateIoCompletionPort(handle, port, key, concurrency)
303+
};
304+
if r == 0 {
305+
return Err(errno_err(vm));
306+
}
307+
Ok(r)
308+
}
309+
310+
#[pyfunction]
311+
fn GetQueuedCompletionStatus(port: HANDLE, msecs: u32, vm: &VirtualMachine) -> PyResult {
312+
let mut bytes_transferred = 0;
313+
let mut completion_key = 0;
314+
let mut overlapped: *mut OVERLAPPED = std::ptr::null_mut();
315+
let ret = unsafe {
316+
windows_sys::Win32::System::IO::GetQueuedCompletionStatus(
317+
port,
318+
&mut bytes_transferred,
319+
&mut completion_key,
320+
&mut overlapped,
321+
msecs,
322+
)
323+
};
324+
let err = if ret != 0 {
325+
Foundation::ERROR_SUCCESS
326+
} else {
327+
unsafe { Foundation::GetLastError() }
328+
};
329+
if overlapped.is_null() {
330+
if err == Foundation::WAIT_TIMEOUT {
331+
return Ok(vm.ctx.none());
332+
} else {
333+
return Err(errno_err(vm));
334+
}
335+
}
336+
337+
let value = vm.ctx.new_tuple(vec![
338+
err.to_pyobject(vm),
339+
completion_key.to_pyobject(vm),
340+
bytes_transferred.to_pyobject(vm),
341+
(overlapped as usize).to_pyobject(vm),
342+
]);
343+
Ok(value.into())
344+
}
345+
}

0 commit comments

Comments
 (0)
0