diff --git a/.cspell.dict/cpython.txt b/.cspell.dict/cpython.txt index 5cc189ff72e..b0115cecc60 100644 --- a/.cspell.dict/cpython.txt +++ b/.cspell.dict/cpython.txt @@ -61,6 +61,7 @@ posonlyargs prec preinitialized pydecimal +pymain pyrepl pythonw PYTHREAD_NAME @@ -73,9 +74,9 @@ stackdepth stginfo stringlib structseq -subscr subkwargs subparams +subscr swappedbytes ticketer tok_oldval diff --git a/crates/stdlib/src/overlapped.rs b/crates/stdlib/src/overlapped.rs index b591a2f0f0f..573d1fc0c85 100644 --- a/crates/stdlib/src/overlapped.rs +++ b/crates/stdlib/src/overlapped.rs @@ -8,7 +8,7 @@ mod _overlapped { // straight-forward port of Modules/overlapped.c use crate::vm::{ - AsObject, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, + Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, builtins::{PyBaseExceptionRef, PyBytesRef, PyType}, common::lock::PyMutex, convert::{ToPyException, ToPyObject}, @@ -18,7 +18,7 @@ mod _overlapped { }; use windows_sys::Win32::{ Foundation::{self, GetLastError, HANDLE}, - Networking::WinSock::{AF_INET, AF_INET6, SOCKADDR_IN, SOCKADDR_IN6}, + Networking::WinSock::{AF_INET, AF_INET6, SOCKADDR, SOCKADDR_IN, SOCKADDR_IN6}, System::IO::OVERLAPPED, }; @@ -41,6 +41,95 @@ mod _overlapped { #[pyattr] const NULL: isize = 0; + // Function pointers for Winsock extension functions + static ACCEPT_EX: std::sync::OnceLock = std::sync::OnceLock::new(); + static CONNECT_EX: std::sync::OnceLock = std::sync::OnceLock::new(); + static DISCONNECT_EX: std::sync::OnceLock = std::sync::OnceLock::new(); + static TRANSMIT_FILE: std::sync::OnceLock = std::sync::OnceLock::new(); + + fn initialize_winsock_extensions(vm: &VirtualMachine) -> PyResult<()> { + use windows_sys::Win32::Networking::WinSock::{ + IPPROTO_TCP, SIO_GET_EXTENSION_FUNCTION_POINTER, SOCK_STREAM, SOCKET_ERROR, WSAIoctl, + closesocket, socket, + }; + + // GUIDs for extension functions + const WSAID_ACCEPTEX: windows_sys::core::GUID = windows_sys::core::GUID { + data1: 0xb5367df1, + data2: 0xcbac, + data3: 0x11cf, + data4: [0x95, 0xca, 0x00, 0x80, 0x5f, 0x48, 0xa1, 0x92], + }; + const WSAID_CONNECTEX: windows_sys::core::GUID = windows_sys::core::GUID { + data1: 0x25a207b9, + data2: 0xddf3, + data3: 0x4660, + data4: [0x8e, 0xe9, 0x76, 0xe5, 0x8c, 0x74, 0x06, 0x3e], + }; + const WSAID_DISCONNECTEX: windows_sys::core::GUID = windows_sys::core::GUID { + data1: 0x7fda2e11, + data2: 0x8630, + data3: 0x436f, + data4: [0xa0, 0x31, 0xf5, 0x36, 0xa6, 0xee, 0xc1, 0x57], + }; + const WSAID_TRANSMITFILE: windows_sys::core::GUID = windows_sys::core::GUID { + data1: 0xb5367df0, + data2: 0xcbac, + data3: 0x11cf, + data4: [0x95, 0xca, 0x00, 0x80, 0x5f, 0x48, 0xa1, 0x92], + }; + + // Check all four locks to prevent partial initialization + if ACCEPT_EX.get().is_some() + && CONNECT_EX.get().is_some() + && DISCONNECT_EX.get().is_some() + && TRANSMIT_FILE.get().is_some() + { + return Ok(()); + } + + let s = unsafe { socket(AF_INET as i32, SOCK_STREAM, IPPROTO_TCP) }; + if s == windows_sys::Win32::Networking::WinSock::INVALID_SOCKET { + return Err( + vm.new_os_error("Failed to create socket for WSA extension init".to_owned()) + ); + } + + let mut dw_bytes: u32 = 0; + + macro_rules! get_extension { + ($guid:expr, $lock:expr) => {{ + let mut func_ptr: usize = 0; + let ret = unsafe { + WSAIoctl( + s, + SIO_GET_EXTENSION_FUNCTION_POINTER, + &$guid as *const _ as *const _, + std::mem::size_of_val(&$guid) as u32, + &mut func_ptr as *mut _ as *mut _, + std::mem::size_of::() as u32, + &mut dw_bytes, + std::ptr::null_mut(), + None, + ) + }; + if ret == SOCKET_ERROR { + unsafe { closesocket(s) }; + return Err(vm.new_os_error("Failed to get WSA extension function".to_owned())); + } + let _ = $lock.set(func_ptr); + }}; + } + + get_extension!(WSAID_ACCEPTEX, ACCEPT_EX); + get_extension!(WSAID_CONNECTEX, CONNECT_EX); + get_extension!(WSAID_DISCONNECTEX, DISCONNECT_EX); + get_extension!(WSAID_TRANSMITFILE, TRANSMIT_FILE); + + unsafe { closesocket(s) }; + Ok(()) + } + #[pyattr] #[pyclass(name)] #[derive(PyPayload)] @@ -62,7 +151,6 @@ mod _overlapped { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { let zelf = self.inner.lock(); f.debug_struct("Overlapped") - // .field("overlapped", &(self.overlapped as *const _ as usize)) .field("handle", &zelf.handle) .field("error", &zelf.error) .field("data", &zelf.data) @@ -70,67 +158,75 @@ mod _overlapped { } } - #[allow(dead_code)] // TODO: remove when done #[derive(Debug)] enum OverlappedData { None, NotStarted, Read(PyBytesRef), + // Fields below store buffers that must be kept alive during async operations + #[allow(dead_code)] ReadInto(PyBuffer), + #[allow(dead_code)] Write(PyBuffer), - Accept(PyObjectRef), - Connect, + #[allow(dead_code)] + Accept(PyBytesRef), + Connect(Vec), // Store address bytes to keep them alive during async operation Disconnect, ConnectNamedPipe, + #[allow(dead_code)] // Reserved for named pipe support WaitNamedPipeAndConnect, TransmitFile, ReadFrom(OverlappedReadFrom), - WriteTo(PyBuffer), + WriteTo(OverlappedWriteTo), // Store address bytes for WSASendTo ReadFromInto(OverlappedReadFromInto), } struct OverlappedReadFrom { - // A (buffer, (host, port)) tuple - result: PyObjectRef, // The actual read buffer - allocated_buffer: PyObjectRef, - #[allow(dead_code)] - address: SOCKADDR_IN6, // TODO: remove when done - address_length: libc::c_int, + allocated_buffer: PyBytesRef, + address: SOCKADDR_IN6, + address_length: i32, } impl core::fmt::Debug for OverlappedReadFrom { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("OverlappedReadFrom") - .field("result", &self.result) .field("allocated_buffer", &self.allocated_buffer) - // .field("address", &self.address) .field("address_length", &self.address_length) .finish() } } struct OverlappedReadFromInto { - // A (number of bytes read, (host, port)) tuple - result: PyObjectRef, /* Buffer passed by the user */ user_buffer: PyBuffer, - #[allow(dead_code)] - address: SOCKADDR_IN6, // TODO: remove when done - address_length: libc::c_int, + address: SOCKADDR_IN6, + address_length: i32, } impl core::fmt::Debug for OverlappedReadFromInto { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("OverlappedReadFromInto") - .field("result", &self.result) .field("user_buffer", &self.user_buffer) - // .field("address", &self.address) .field("address_length", &self.address_length) .finish() } } + struct OverlappedWriteTo { + buf: PyBuffer, + address: Vec, // Keep address alive during async operation + } + + impl core::fmt::Debug for OverlappedWriteTo { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("OverlappedWriteTo") + .field("buf", &self.buf) + .field("address", &self.address.len()) + .finish() + } + } + fn mark_as_completed(ov: &mut OVERLAPPED) { ov.Internal = 0; if !ov.hEvent.is_null() { @@ -146,7 +242,6 @@ mod _overlapped { ERROR_CONNECTION_ABORTED => vm.ctx.exceptions.connection_aborted_error, err => return std::io::Error::from_raw_os_error(err as i32).to_pyexception(vm), }; - // TODO: set errno and winerror vm.new_exception_empty(exc.to_owned()) } @@ -154,14 +249,103 @@ mod _overlapped { overlapped.Internal != (Foundation::STATUS_PENDING as usize) } - /// Parse a SOCKADDR_IN6 (which can also hold IPv4 addresses) to a Python address tuple - fn unparse_address( - addr: &SOCKADDR_IN6, - _addr_len: libc::c_int, - vm: &VirtualMachine, - ) -> PyObjectRef { - use crate::vm::convert::ToPyObject; + /// Parse a Python address tuple to SOCKADDR + fn parse_address(addr_obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult<(Vec, i32)> { + use crate::vm::builtins::PyTuple; + use windows_sys::Win32::Networking::WinSock::WSAStringToAddressW; + + let tuple = addr_obj + .downcast_ref::() + .ok_or_else(|| vm.new_type_error("address must be a tuple".to_owned()))?; + + let tuple_len = tuple.len(); + + if tuple_len == 2 { + // IPv4: (host, port) + let host: String = tuple[0].try_to_value(vm)?; + let port: u16 = tuple[1].try_to_value(vm)?; + + let mut addr: SOCKADDR_IN = unsafe { std::mem::zeroed() }; + addr.sin_family = AF_INET; + addr.sin_port = port.to_be(); + + // Convert host string to address + let host_wide: Vec = host.encode_utf16().chain(std::iter::once(0)).collect(); + let mut addr_len = std::mem::size_of::() as i32; + + let ret = unsafe { + WSAStringToAddressW( + host_wide.as_ptr(), + AF_INET as i32, + std::ptr::null(), + &mut addr as *mut _ as *mut SOCKADDR, + &mut addr_len, + ) + }; + + if ret != 0 { + return Err(vm.new_os_error(format!("Invalid IPv4 address: {}", host))); + } + + // Restore port (WSAStringToAddressW overwrites it) + addr.sin_port = port.to_be(); + + let bytes = unsafe { + std::slice::from_raw_parts( + &addr as *const _ as *const u8, + std::mem::size_of::(), + ) + }; + Ok((bytes.to_vec(), std::mem::size_of::() as i32)) + } else if tuple_len == 4 { + // IPv6: (host, port, flowinfo, scope_id) + let host: String = tuple[0].try_to_value(vm)?; + let port: u16 = tuple[1].try_to_value(vm)?; + let flowinfo: u32 = tuple[2].try_to_value(vm)?; + let scope_id: u32 = tuple[3].try_to_value(vm)?; + + let mut addr: SOCKADDR_IN6 = unsafe { std::mem::zeroed() }; + addr.sin6_family = AF_INET6; + addr.sin6_port = port.to_be(); + addr.sin6_flowinfo = flowinfo; + addr.Anonymous.sin6_scope_id = scope_id; + + let host_wide: Vec = host.encode_utf16().chain(std::iter::once(0)).collect(); + let mut addr_len = std::mem::size_of::() as i32; + + let ret = unsafe { + WSAStringToAddressW( + host_wide.as_ptr(), + AF_INET6 as i32, + std::ptr::null(), + &mut addr as *mut _ as *mut SOCKADDR, + &mut addr_len, + ) + }; + + if ret != 0 { + return Err(vm.new_os_error(format!("Invalid IPv6 address: {}", host))); + } + + // Restore fields that WSAStringToAddressW might overwrite + addr.sin6_port = port.to_be(); + addr.sin6_flowinfo = flowinfo; + addr.Anonymous.sin6_scope_id = scope_id; + + let bytes = unsafe { + std::slice::from_raw_parts( + &addr as *const _ as *const u8, + std::mem::size_of::(), + ) + }; + Ok((bytes.to_vec(), std::mem::size_of::() as i32)) + } else { + Err(vm.new_value_error("address tuple must have 2 or 4 elements".to_owned())) + } + } + /// Parse a SOCKADDR_IN6 (which can also hold IPv4 addresses) to a Python address tuple + fn unparse_address(addr: &SOCKADDR_IN6, _addr_len: i32, vm: &VirtualMachine) -> PyObjectRef { unsafe { let family = addr.sin6_family; if family == AF_INET { @@ -214,73 +398,16 @@ mod _overlapped { && !matches!(inner.data, OverlappedData::NotStarted) } - fn WSARecv_inner( - inner: &mut OverlappedInner, - handle: isize, - buf: &[u8], - mut flags: u32, - vm: &VirtualMachine, - ) -> PyResult { - use windows_sys::Win32::Foundation::{ - ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS, - }; - - let wsabuf = windows_sys::Win32::Networking::WinSock::WSABUF { - buf: buf.as_ptr() as *mut _, - len: buf.len() as _, - }; - let mut n_read: u32 = 0; - // TODO: optimization with MaybeUninit - let ret = unsafe { - windows_sys::Win32::Networking::WinSock::WSARecv( - handle as _, - &wsabuf, - 1, - &mut n_read, - &mut flags, - &mut inner.overlapped, - None, - ) - }; - let err = if ret < 0 { - unsafe { windows_sys::Win32::Networking::WinSock::WSAGetLastError() as u32 } - } else { - Foundation::ERROR_SUCCESS - }; - inner.error = err; - match err { - ERROR_BROKEN_PIPE => { - mark_as_completed(&mut inner.overlapped); - Err(from_windows_err(err, vm)) - } - ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()), - _ => Err(from_windows_err(err, vm)), - } + #[pygetset] + fn error(&self, _vm: &VirtualMachine) -> u32 { + let inner = self.inner.lock(); + inner.error } - #[pymethod] - fn WSARecv( - zelf: &Py, - handle: isize, - size: u32, - flags: u32, - vm: &VirtualMachine, - ) -> PyResult { - let mut inner = zelf.inner.lock(); - if !matches!(inner.data, OverlappedData::None) { - return Err(vm.new_value_error("operation already attempted")); - } - - #[cfg(target_pointer_width = "32")] - let size = core::cmp::min(size, std::isize::MAX as _); - - let buf = vec![0u8; std::cmp::max(size, 1) as usize]; - let buf = vm.ctx.new_bytes(buf); - inner.handle = handle as _; - - let r = Self::WSARecv_inner(&mut inner, handle as _, buf.as_bytes(), flags, vm); - inner.data = OverlappedData::Read(buf); - r + #[pygetset] + fn event(&self, _vm: &VirtualMachine) -> isize { + let inner = self.inner.lock(); + inner.overlapped.hEvent as isize } #[pymethod] @@ -349,16 +476,7 @@ mod _overlapped { // For read operations, broken pipe is acceptable match &inner.data { OverlappedData::Read(_) | OverlappedData::ReadInto(_) => {} - OverlappedData::ReadFrom(rf) - if rf.result.is(&vm.ctx.none()) - || rf.allocated_buffer.is(&vm.ctx.none()) => - { - return Err(from_windows_err(err, vm)); - } OverlappedData::ReadFrom(_) => {} - OverlappedData::ReadFromInto(rfi) if rfi.result.is(&vm.ctx.none()) => { - return Err(from_windows_err(err, vm)); - } OverlappedData::ReadFromInto(_) => {} _ => return Err(from_windows_err(err, vm)), } @@ -372,7 +490,6 @@ mod _overlapped { // Return result based on operation type match &inner.data { OverlappedData::Read(buf) => { - // Resize buffer to actual bytes read let bytes = buf.as_bytes(); let result = if transferred as usize != bytes.len() { vm.ctx.new_bytes(bytes[..transferred as usize].to_vec()) @@ -381,45 +498,22 @@ mod _overlapped { }; Ok(result.into()) } - OverlappedData::ReadInto(_) => { - // Return number of bytes read - Ok(vm.ctx.new_int(transferred).into()) - } - OverlappedData::Write(_) => { - // Return number of bytes written + OverlappedData::ReadInto(_) => Ok(vm.ctx.new_int(transferred).into()), + OverlappedData::Write(_) | OverlappedData::WriteTo(_) => { Ok(vm.ctx.new_int(transferred).into()) } - OverlappedData::Accept(_) => { - // Return None for accept - Ok(vm.ctx.none()) - } - OverlappedData::Connect => { - // Return None for connect - Ok(vm.ctx.none()) - } - OverlappedData::Disconnect => { - // Return None for disconnect - Ok(vm.ctx.none()) - } - OverlappedData::ConnectNamedPipe => { - // Return None for connect named pipe - Ok(vm.ctx.none()) - } - OverlappedData::WaitNamedPipeAndConnect => { - // Return None - Ok(vm.ctx.none()) - } + OverlappedData::Accept(_) => Ok(vm.ctx.none()), + OverlappedData::Connect(_) => Ok(vm.ctx.none()), + OverlappedData::Disconnect => Ok(vm.ctx.none()), + OverlappedData::ConnectNamedPipe => Ok(vm.ctx.none()), + OverlappedData::WaitNamedPipeAndConnect => Ok(vm.ctx.none()), + OverlappedData::TransmitFile => Ok(vm.ctx.none()), OverlappedData::ReadFrom(rf) => { - // Return (resized_buffer, (host, port)) tuple - let buf = rf - .allocated_buffer - .downcast_ref::() - .unwrap(); - let bytes = buf.as_bytes(); + let bytes = rf.allocated_buffer.as_bytes(); let resized_buf = if transferred as usize != bytes.len() { vm.ctx.new_bytes(bytes[..transferred as usize].to_vec()) } else { - buf.to_owned() + rf.allocated_buffer.clone() }; let addr_tuple = unparse_address(&rf.address, rf.address_length, vm); Ok(vm @@ -428,7 +522,6 @@ mod _overlapped { .into()) } OverlappedData::ReadFromInto(rfi) => { - // Return (transferred, (host, port)) tuple let addr_tuple = unparse_address(&rfi.address, rfi.address_length, vm); Ok(vm .ctx @@ -438,108 +531,1299 @@ mod _overlapped { _ => Ok(vm.ctx.none()), } } - } - impl Constructor for Overlapped { - type Args = (isize,); + // ReadFile + #[pymethod] + fn ReadFile(zelf: &Py, handle: isize, size: u32, vm: &VirtualMachine) -> PyResult { + use windows_sys::Win32::Foundation::{ + ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS, + }; + use windows_sys::Win32::Storage::FileSystem::ReadFile; - fn py_new( - _cls: &Py, - (mut event,): Self::Args, - vm: &VirtualMachine, - ) -> PyResult { - if event == INVALID_HANDLE_VALUE { - event = unsafe { - windows_sys::Win32::System::Threading::CreateEventA( - core::ptr::null(), - Foundation::TRUE, - Foundation::FALSE, - core::ptr::null(), - ) as isize - }; - if event == NULL { - return Err(vm.new_last_os_error()); + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + #[cfg(target_pointer_width = "32")] + let size = core::cmp::min(size, isize::MAX as u32); + + let buf = vec![0u8; std::cmp::max(size, 1) as usize]; + let buf = vm.ctx.new_bytes(buf); + inner.handle = handle as HANDLE; + inner.data = OverlappedData::Read(buf.clone()); + + let mut nread: u32 = 0; + let ret = unsafe { + ReadFile( + handle as HANDLE, + buf.as_bytes().as_ptr() as *mut _, + size, + &mut nread, + &mut inner.overlapped, + ) + }; + + let err = if ret != 0 { + ERROR_SUCCESS + } else { + unsafe { GetLastError() } + }; + inner.error = err; + + match err { + ERROR_BROKEN_PIPE => { + mark_as_completed(&mut inner.overlapped); + Err(from_windows_err(err, vm)) + } + ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) } } + } - let mut overlapped: OVERLAPPED = unsafe { std::mem::zeroed() }; - if event != NULL { - overlapped.hEvent = event as _; + // ReadFileInto + #[pymethod] + fn ReadFileInto( + zelf: &Py, + handle: isize, + buf: PyBuffer, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ + ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS, + }; + use windows_sys::Win32::Storage::FileSystem::ReadFile; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); } - let inner = OverlappedInner { - overlapped, - handle: NULL as _, - error: 0, - data: OverlappedData::None, + + inner.handle = handle as HANDLE; + let buf_len = buf.desc.len; + + // For async read, buffer must be contiguous - we can't use a temporary copy + // because Windows writes data directly to the buffer after this call returns + let Some(contiguous) = buf.as_contiguous_mut() else { + return Err(vm.new_buffer_error("buffer is not contiguous".to_owned())); }; - Ok(Overlapped { - inner: PyMutex::new(inner), - }) - } - } - unsafe fn u64_to_handle(raw_ptr_value: u64) -> HANDLE { - raw_ptr_value as HANDLE - } + inner.data = OverlappedData::ReadInto(buf.clone()); - #[pyfunction] - fn CreateIoCompletionPort( - handle: isize, - port: isize, - key: usize, - concurrency: u32, - vm: &VirtualMachine, - ) -> PyResult { - let r = unsafe { - windows_sys::Win32::System::IO::CreateIoCompletionPort( - handle as _, - port as _, - key, - concurrency, - ) as isize - }; - if r as usize == 0 { - return Err(vm.new_last_os_error()); - } - Ok(r) - } + let mut nread: u32 = 0; + let ret = unsafe { + ReadFile( + handle as HANDLE, + contiguous.as_ptr() as *mut _, + buf_len as u32, + &mut nread, + &mut inner.overlapped, + ) + }; - #[pyfunction] - fn GetQueuedCompletionStatus(port: isize, msecs: u32, vm: &VirtualMachine) -> PyResult { - let mut bytes_transferred = 0; - let mut completion_key = 0; - let mut overlapped: *mut OVERLAPPED = std::ptr::null_mut(); - let ret = unsafe { - windows_sys::Win32::System::IO::GetQueuedCompletionStatus( - port as _, - &mut bytes_transferred, - &mut completion_key, - &mut overlapped, - msecs, - ) - }; - let err = if ret != 0 { - Foundation::ERROR_SUCCESS - } else { - unsafe { Foundation::GetLastError() } - }; - if overlapped.is_null() { - if err == Foundation::WAIT_TIMEOUT { - return Ok(vm.ctx.none()); + let err = if ret != 0 { + ERROR_SUCCESS } else { - return Err(vm.new_last_os_error()); + unsafe { GetLastError() } + }; + inner.error = err; + + match err { + ERROR_BROKEN_PIPE => { + mark_as_completed(&mut inner.overlapped); + Err(from_windows_err(err, vm)) + } + ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } } } - let value = vm.ctx.new_tuple(vec![ - err.to_pyobject(vm), - completion_key.to_pyobject(vm), - bytes_transferred.to_pyobject(vm), + // WSARecv + #[pymethod] + fn WSARecv( + zelf: &Py, + handle: isize, + size: u32, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ + ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS, + }; + use windows_sys::Win32::Networking::WinSock::{WSABUF, WSAGetLastError, WSARecv}; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + let mut flags = flags.unwrap_or(0); + + #[cfg(target_pointer_width = "32")] + let size = core::cmp::min(size, isize::MAX as u32); + + let buf = vec![0u8; std::cmp::max(size, 1) as usize]; + let buf = vm.ctx.new_bytes(buf); + inner.handle = handle as HANDLE; + inner.data = OverlappedData::Read(buf.clone()); + + let wsabuf = WSABUF { + buf: buf.as_bytes().as_ptr() as *mut _, + len: size, + }; + let mut nread: u32 = 0; + + let ret = unsafe { + WSARecv( + handle as _, + &wsabuf, + 1, + &mut nread, + &mut flags, + &mut inner.overlapped, + None, + ) + }; + + let err = if ret < 0 { + unsafe { WSAGetLastError() as u32 } + } else { + ERROR_SUCCESS + }; + inner.error = err; + + match err { + ERROR_BROKEN_PIPE => { + mark_as_completed(&mut inner.overlapped); + Err(from_windows_err(err, vm)) + } + ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // WSARecvInto + #[pymethod] + fn WSARecvInto( + zelf: &Py, + handle: isize, + buf: PyBuffer, + flags: u32, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ + ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS, + }; + use windows_sys::Win32::Networking::WinSock::{WSABUF, WSAGetLastError, WSARecv}; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + let mut flags = flags; + inner.handle = handle as HANDLE; + let buf_len = buf.desc.len; + + let Some(contiguous) = buf.as_contiguous_mut() else { + return Err(vm.new_buffer_error("buffer is not contiguous".to_owned())); + }; + + inner.data = OverlappedData::ReadInto(buf.clone()); + + let wsabuf = WSABUF { + buf: contiguous.as_ptr() as *mut _, + len: buf_len as u32, + }; + let mut nread: u32 = 0; + + let ret = unsafe { + WSARecv( + handle as _, + &wsabuf, + 1, + &mut nread, + &mut flags, + &mut inner.overlapped, + None, + ) + }; + + let err = if ret < 0 { + unsafe { WSAGetLastError() as u32 } + } else { + ERROR_SUCCESS + }; + inner.error = err; + + match err { + ERROR_BROKEN_PIPE => { + mark_as_completed(&mut inner.overlapped); + Err(from_windows_err(err, vm)) + } + ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // WriteFile + #[pymethod] + fn WriteFile( + zelf: &Py, + handle: isize, + buf: PyBuffer, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ERROR_IO_PENDING, ERROR_SUCCESS}; + use windows_sys::Win32::Storage::FileSystem::WriteFile; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + inner.handle = handle as HANDLE; + let buf_len = buf.desc.len; + + // For async write, buffer must be contiguous - we can't use a temporary copy + // because Windows reads from the buffer after this call returns + let Some(contiguous) = buf.as_contiguous() else { + return Err(vm.new_buffer_error("buffer is not contiguous".to_owned())); + }; + + inner.data = OverlappedData::Write(buf.clone()); + + let mut written: u32 = 0; + let ret = unsafe { + WriteFile( + handle as HANDLE, + contiguous.as_ptr() as *const _, + buf_len as u32, + &mut written, + &mut inner.overlapped, + ) + }; + + let err = if ret != 0 { + ERROR_SUCCESS + } else { + unsafe { GetLastError() } + }; + inner.error = err; + + match err { + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // WSASend + #[pymethod] + fn WSASend( + zelf: &Py, + handle: isize, + buf: PyBuffer, + flags: u32, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ERROR_IO_PENDING, ERROR_SUCCESS}; + use windows_sys::Win32::Networking::WinSock::{WSABUF, WSAGetLastError, WSASend}; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + inner.handle = handle as HANDLE; + let buf_len = buf.desc.len; + + let Some(contiguous) = buf.as_contiguous() else { + return Err(vm.new_buffer_error("buffer is not contiguous".to_owned())); + }; + + inner.data = OverlappedData::Write(buf.clone()); + + let wsabuf = WSABUF { + buf: contiguous.as_ptr() as *mut _, + len: buf_len as u32, + }; + let mut written: u32 = 0; + + let ret = unsafe { + WSASend( + handle as _, + &wsabuf, + 1, + &mut written, + flags, + &mut inner.overlapped, + None, + ) + }; + + let err = if ret < 0 { + unsafe { WSAGetLastError() as u32 } + } else { + ERROR_SUCCESS + }; + inner.error = err; + + match err { + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // AcceptEx + #[pymethod] + fn AcceptEx( + zelf: &Py, + listen_socket: isize, + accept_socket: isize, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ERROR_IO_PENDING, ERROR_SUCCESS}; + use windows_sys::Win32::Networking::WinSock::WSAGetLastError; + + initialize_winsock_extensions(vm)?; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + // Buffer size: local address + remote address + let size = std::mem::size_of::() + 16; + let buf = vec![0u8; size * 2]; + let buf = vm.ctx.new_bytes(buf); + + inner.handle = listen_socket as HANDLE; + inner.data = OverlappedData::Accept(buf.clone()); + + let mut bytes_received: u32 = 0; + + type AcceptExFn = unsafe extern "system" fn( + sListenSocket: usize, + sAcceptSocket: usize, + lpOutputBuffer: *mut core::ffi::c_void, + dwReceiveDataLength: u32, + dwLocalAddressLength: u32, + dwRemoteAddressLength: u32, + lpdwBytesReceived: *mut u32, + lpOverlapped: *mut OVERLAPPED, + ) -> i32; + + let accept_ex: AcceptExFn = unsafe { std::mem::transmute(*ACCEPT_EX.get().unwrap()) }; + + let ret = unsafe { + accept_ex( + listen_socket as _, + accept_socket as _, + buf.as_bytes().as_ptr() as *mut _, + 0, + size as u32, + size as u32, + &mut bytes_received, + &mut inner.overlapped, + ) + }; + + let err = if ret != 0 { + ERROR_SUCCESS + } else { + unsafe { WSAGetLastError() as u32 } + }; + inner.error = err; + + match err { + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // ConnectEx + #[pymethod] + fn ConnectEx( + zelf: &Py, + socket: isize, + address: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ERROR_IO_PENDING, ERROR_SUCCESS}; + use windows_sys::Win32::Networking::WinSock::WSAGetLastError; + + initialize_winsock_extensions(vm)?; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + let (addr_bytes, addr_len) = parse_address(&address, vm)?; + + inner.handle = socket as HANDLE; + // Store addr_bytes in OverlappedData to keep it alive during async operation + inner.data = OverlappedData::Connect(addr_bytes); + + type ConnectExFn = unsafe extern "system" fn( + s: usize, + name: *const SOCKADDR, + namelen: i32, + lpSendBuffer: *const core::ffi::c_void, + dwSendDataLength: u32, + lpdwBytesSent: *mut u32, + lpOverlapped: *mut OVERLAPPED, + ) -> i32; + + let connect_ex: ConnectExFn = + unsafe { std::mem::transmute(*CONNECT_EX.get().unwrap()) }; + + // Get pointer to the stored address data + let addr_ptr = match &inner.data { + OverlappedData::Connect(bytes) => bytes.as_ptr(), + _ => unreachable!(), + }; + + let ret = unsafe { + connect_ex( + socket as _, + addr_ptr as *const SOCKADDR, + addr_len, + std::ptr::null(), + 0, + std::ptr::null_mut(), + &mut inner.overlapped, + ) + }; + + let err = if ret != 0 { + ERROR_SUCCESS + } else { + unsafe { WSAGetLastError() as u32 } + }; + inner.error = err; + + match err { + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // DisconnectEx + #[pymethod] + fn DisconnectEx( + zelf: &Py, + socket: isize, + flags: u32, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ERROR_IO_PENDING, ERROR_SUCCESS}; + use windows_sys::Win32::Networking::WinSock::WSAGetLastError; + + initialize_winsock_extensions(vm)?; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + inner.handle = socket as HANDLE; + inner.data = OverlappedData::Disconnect; + + type DisconnectExFn = unsafe extern "system" fn( + s: usize, + lpOverlapped: *mut OVERLAPPED, + dwFlags: u32, + dwReserved: u32, + ) -> i32; + + let disconnect_ex: DisconnectExFn = + unsafe { std::mem::transmute(*DISCONNECT_EX.get().unwrap()) }; + + let ret = unsafe { disconnect_ex(socket as _, &mut inner.overlapped, flags, 0) }; + + let err = if ret != 0 { + ERROR_SUCCESS + } else { + unsafe { WSAGetLastError() as u32 } + }; + inner.error = err; + + match err { + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // TransmitFile + #[allow(clippy::too_many_arguments)] + #[pymethod] + fn TransmitFile( + zelf: &Py, + socket: isize, + file: isize, + offset: u32, + offset_high: u32, + count_to_write: u32, + count_per_send: u32, + flags: u32, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ERROR_IO_PENDING, ERROR_SUCCESS}; + use windows_sys::Win32::Networking::WinSock::WSAGetLastError; + + initialize_winsock_extensions(vm)?; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + inner.handle = socket as HANDLE; + inner.data = OverlappedData::TransmitFile; + inner.overlapped.Anonymous.Anonymous.Offset = offset; + inner.overlapped.Anonymous.Anonymous.OffsetHigh = offset_high; + + type TransmitFileFn = unsafe extern "system" fn( + hSocket: usize, + hFile: HANDLE, + nNumberOfBytesToWrite: u32, + nNumberOfBytesPerSend: u32, + lpOverlapped: *mut OVERLAPPED, + lpTransmitBuffers: *const core::ffi::c_void, + dwReserved: u32, + ) -> i32; + + let transmit_file: TransmitFileFn = + unsafe { std::mem::transmute(*TRANSMIT_FILE.get().unwrap()) }; + + let ret = unsafe { + transmit_file( + socket as _, + file as HANDLE, + count_to_write, + count_per_send, + &mut inner.overlapped, + std::ptr::null(), + flags, + ) + }; + + let err = if ret != 0 { + ERROR_SUCCESS + } else { + unsafe { WSAGetLastError() as u32 } + }; + inner.error = err; + + match err { + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // ConnectNamedPipe + #[pymethod] + fn ConnectNamedPipe(zelf: &Py, pipe: isize, vm: &VirtualMachine) -> PyResult { + use windows_sys::Win32::Foundation::{ + ERROR_IO_PENDING, ERROR_PIPE_CONNECTED, ERROR_SUCCESS, + }; + use windows_sys::Win32::System::Pipes::ConnectNamedPipe; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + inner.handle = pipe as HANDLE; + inner.data = OverlappedData::ConnectNamedPipe; + + let ret = unsafe { ConnectNamedPipe(pipe as HANDLE, &mut inner.overlapped) }; + + let err = if ret != 0 { + ERROR_SUCCESS + } else { + unsafe { GetLastError() } + }; + inner.error = err; + + match err { + ERROR_PIPE_CONNECTED => { + mark_as_completed(&mut inner.overlapped); + Ok(true) + } + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(false), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // ConnectPipe - this is a static method that returns a handle + #[pymethod] + fn ConnectPipe(address: String, vm: &VirtualMachine) -> PyResult { + use windows_sys::Win32::Storage::FileSystem::{ + CreateFileW, FILE_FLAG_OVERLAPPED, FILE_GENERIC_READ, FILE_GENERIC_WRITE, + OPEN_EXISTING, + }; + + let address_wide: Vec = address.encode_utf16().chain(std::iter::once(0)).collect(); + + let handle = unsafe { + CreateFileW( + address_wide.as_ptr(), + FILE_GENERIC_READ | FILE_GENERIC_WRITE, + 0, + std::ptr::null(), + OPEN_EXISTING, + FILE_FLAG_OVERLAPPED, + std::ptr::null_mut(), + ) + }; + + if handle == windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE { + return Err(vm.new_last_os_error()); + } + + Ok(handle as isize) + } + + // WSASendTo + #[pymethod] + fn WSASendTo( + zelf: &Py, + handle: isize, + buf: PyBuffer, + flags: u32, + address: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ERROR_IO_PENDING, ERROR_SUCCESS}; + use windows_sys::Win32::Networking::WinSock::{WSABUF, WSAGetLastError, WSASendTo}; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + let (addr_bytes, addr_len) = parse_address(&address, vm)?; + + inner.handle = handle as HANDLE; + let buf_len = buf.desc.len; + + let Some(contiguous) = buf.as_contiguous() else { + return Err(vm.new_buffer_error("buffer is not contiguous".to_owned())); + }; + + // Store both buffer and address in OverlappedData to keep them alive + inner.data = OverlappedData::WriteTo(OverlappedWriteTo { + buf: buf.clone(), + address: addr_bytes, + }); + + let wsabuf = WSABUF { + buf: contiguous.as_ptr() as *mut _, + len: buf_len as u32, + }; + let mut written: u32 = 0; + + // Get pointer to the stored address data + let addr_ptr = match &inner.data { + OverlappedData::WriteTo(wt) => wt.address.as_ptr(), + _ => unreachable!(), + }; + + let ret = unsafe { + WSASendTo( + handle as _, + &wsabuf, + 1, + &mut written, + flags, + addr_ptr as *const SOCKADDR, + addr_len, + &mut inner.overlapped, + None, + ) + }; + + let err = if ret < 0 { + unsafe { WSAGetLastError() as u32 } + } else { + ERROR_SUCCESS + }; + inner.error = err; + + match err { + ERROR_SUCCESS | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // WSARecvFrom + #[pymethod] + fn WSARecvFrom( + zelf: &Py, + handle: isize, + size: u32, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ + ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS, + }; + use windows_sys::Win32::Networking::WinSock::{WSABUF, WSAGetLastError, WSARecvFrom}; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + let mut flags = flags.unwrap_or(0); + + #[cfg(target_pointer_width = "32")] + let size = core::cmp::min(size, isize::MAX as u32); + + let buf = vec![0u8; std::cmp::max(size, 1) as usize]; + let buf = vm.ctx.new_bytes(buf); + inner.handle = handle as HANDLE; + + let address: SOCKADDR_IN6 = unsafe { std::mem::zeroed() }; + let address_length = std::mem::size_of::() as i32; + + inner.data = OverlappedData::ReadFrom(OverlappedReadFrom { + allocated_buffer: buf.clone(), + address, + address_length, + }); + + let wsabuf = WSABUF { + buf: buf.as_bytes().as_ptr() as *mut _, + len: size, + }; + let mut nread: u32 = 0; + + // Get mutable reference to address in inner.data + let (addr_ptr, addr_len_ptr) = match &mut inner.data { + OverlappedData::ReadFrom(rf) => ( + &mut rf.address as *mut SOCKADDR_IN6, + &mut rf.address_length as *mut i32, + ), + _ => unreachable!(), + }; + + let ret = unsafe { + WSARecvFrom( + handle as _, + &wsabuf, + 1, + &mut nread, + &mut flags, + addr_ptr as *mut SOCKADDR, + addr_len_ptr, + &mut inner.overlapped, + None, + ) + }; + + let err = if ret < 0 { + unsafe { WSAGetLastError() as u32 } + } else { + ERROR_SUCCESS + }; + inner.error = err; + + match err { + ERROR_BROKEN_PIPE => { + mark_as_completed(&mut inner.overlapped); + Err(from_windows_err(err, vm)) + } + ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + + // WSARecvFromInto + #[pymethod] + fn WSARecvFromInto( + zelf: &Py, + handle: isize, + buf: PyBuffer, + size: u32, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::Foundation::{ + ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS, + }; + use windows_sys::Win32::Networking::WinSock::{WSABUF, WSAGetLastError, WSARecvFrom}; + + let mut inner = zelf.inner.lock(); + if !matches!(inner.data, OverlappedData::None) { + return Err(vm.new_value_error("operation already attempted".to_owned())); + } + + let mut flags = flags.unwrap_or(0); + inner.handle = handle as HANDLE; + + let Some(contiguous) = buf.as_contiguous_mut() else { + return Err(vm.new_buffer_error("buffer is not contiguous".to_owned())); + }; + + // Validate size against buffer length to prevent buffer overflow + let buf_len = buf.desc.len as u32; + if size > buf_len { + return Err(vm.new_value_error("size exceeds buffer length".to_owned())); + } + + let address: SOCKADDR_IN6 = unsafe { std::mem::zeroed() }; + let address_length = std::mem::size_of::() as i32; + + inner.data = OverlappedData::ReadFromInto(OverlappedReadFromInto { + user_buffer: buf.clone(), + address, + address_length, + }); + + let wsabuf = WSABUF { + buf: contiguous.as_ptr() as *mut _, + len: size, + }; + let mut nread: u32 = 0; + + // Get mutable reference to address in inner.data + let (addr_ptr, addr_len_ptr) = match &mut inner.data { + OverlappedData::ReadFromInto(rfi) => ( + &mut rfi.address as *mut SOCKADDR_IN6, + &mut rfi.address_length as *mut i32, + ), + _ => unreachable!(), + }; + + let ret = unsafe { + WSARecvFrom( + handle as _, + &wsabuf, + 1, + &mut nread, + &mut flags, + addr_ptr as *mut SOCKADDR, + addr_len_ptr, + &mut inner.overlapped, + None, + ) + }; + + let err = if ret < 0 { + unsafe { WSAGetLastError() as u32 } + } else { + ERROR_SUCCESS + }; + inner.error = err; + + match err { + ERROR_BROKEN_PIPE => { + mark_as_completed(&mut inner.overlapped); + Err(from_windows_err(err, vm)) + } + ERROR_SUCCESS | ERROR_MORE_DATA | ERROR_IO_PENDING => Ok(vm.ctx.none()), + _ => { + inner.data = OverlappedData::NotStarted; + Err(from_windows_err(err, vm)) + } + } + } + } + + impl Constructor for Overlapped { + type Args = (OptionalArg,); + + fn py_new(_cls: &Py, (event,): Self::Args, vm: &VirtualMachine) -> PyResult { + let mut event = event.unwrap_or(INVALID_HANDLE_VALUE); + + if event == INVALID_HANDLE_VALUE { + event = unsafe { + windows_sys::Win32::System::Threading::CreateEventA( + core::ptr::null(), + Foundation::TRUE, + Foundation::FALSE, + core::ptr::null(), + ) as isize + }; + if event == NULL { + return Err(vm.new_last_os_error()); + } + } + + let mut overlapped: OVERLAPPED = unsafe { std::mem::zeroed() }; + if event != NULL { + overlapped.hEvent = event as HANDLE; + } + let inner = OverlappedInner { + overlapped, + handle: NULL as HANDLE, + error: 0, + data: OverlappedData::None, + }; + Ok(Overlapped { + inner: PyMutex::new(inner), + }) + } + } + + #[pyfunction] + fn CreateIoCompletionPort( + handle: isize, + port: isize, + key: usize, + concurrency: u32, + vm: &VirtualMachine, + ) -> PyResult { + let r = unsafe { + windows_sys::Win32::System::IO::CreateIoCompletionPort( + handle as HANDLE, + port as HANDLE, + key, + concurrency, + ) as isize + }; + if r == 0 { + return Err(vm.new_last_os_error()); + } + Ok(r) + } + + #[pyfunction] + fn GetQueuedCompletionStatus(port: isize, msecs: u32, vm: &VirtualMachine) -> PyResult { + let mut bytes_transferred = 0; + let mut completion_key = 0; + let mut overlapped: *mut OVERLAPPED = std::ptr::null_mut(); + let ret = unsafe { + windows_sys::Win32::System::IO::GetQueuedCompletionStatus( + port as HANDLE, + &mut bytes_transferred, + &mut completion_key, + &mut overlapped, + msecs, + ) + }; + let err = if ret != 0 { + Foundation::ERROR_SUCCESS + } else { + unsafe { GetLastError() } + }; + if overlapped.is_null() { + if err == Foundation::WAIT_TIMEOUT { + return Ok(vm.ctx.none()); + } else { + return Err(vm.new_last_os_error()); + } + } + + let value = vm.ctx.new_tuple(vec![ + err.to_pyobject(vm), + bytes_transferred.to_pyobject(vm), + completion_key.to_pyobject(vm), (overlapped as usize).to_pyobject(vm), ]); Ok(value.into()) } + #[pyfunction] + fn PostQueuedCompletionStatus( + port: isize, + bytes: u32, + key: usize, + address: usize, + vm: &VirtualMachine, + ) -> PyResult<()> { + let ret = unsafe { + windows_sys::Win32::System::IO::PostQueuedCompletionStatus( + port as HANDLE, + bytes, + key, + address as *mut OVERLAPPED, + ) + }; + if ret == 0 { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + + // Registry to track callback data for proper cleanup + // Uses Arc for reference counting to prevent use-after-free when callback + // and UnregisterWait race - the data stays alive until both are done + static WAIT_CALLBACK_REGISTRY: std::sync::OnceLock< + std::sync::Mutex>>, + > = std::sync::OnceLock::new(); + + fn wait_callback_registry() + -> &'static std::sync::Mutex>> + { + WAIT_CALLBACK_REGISTRY + .get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new())) + } + + // Callback data for RegisterWaitWithQueue + // Uses Arc to ensure the data stays alive while callback is executing + struct PostCallbackData { + completion_port: HANDLE, + overlapped: *mut OVERLAPPED, + } + + // SAFETY: The pointers are handles/addresses passed from Python and are + // only used to call Windows APIs. They are not dereferenced as Rust pointers. + unsafe impl Send for PostCallbackData {} + unsafe impl Sync for PostCallbackData {} + + unsafe extern "system" fn post_to_queue_callback( + parameter: *mut core::ffi::c_void, + timer_or_wait_fired: bool, + ) { + // Reconstruct Arc from raw pointer - this gives us ownership of one reference + // The Arc prevents use-after-free since we own a reference count + let data = unsafe { std::sync::Arc::from_raw(parameter as *const PostCallbackData) }; + + unsafe { + let _ = windows_sys::Win32::System::IO::PostQueuedCompletionStatus( + data.completion_port, + if timer_or_wait_fired { 1 } else { 0 }, + 0, + data.overlapped, + ); + } + // Arc is dropped here, decrementing refcount + // Memory is freed only when all references (callback + registry) are gone + } + + #[pyfunction] + fn RegisterWaitWithQueue( + object: isize, + completion_port: isize, + overlapped: usize, + timeout: u32, + vm: &VirtualMachine, + ) -> PyResult { + use windows_sys::Win32::System::Threading::{ + RegisterWaitForSingleObject, WT_EXECUTEINWAITTHREAD, WT_EXECUTEONLYONCE, + }; + + let data = std::sync::Arc::new(PostCallbackData { + completion_port: completion_port as HANDLE, + overlapped: overlapped as *mut OVERLAPPED, + }); + + // Create raw pointer for the callback - this increments refcount + let data_ptr = std::sync::Arc::into_raw(data.clone()); + + let mut new_wait_object: HANDLE = std::ptr::null_mut(); + let ret = unsafe { + RegisterWaitForSingleObject( + &mut new_wait_object, + object as HANDLE, + Some(post_to_queue_callback), + data_ptr as *mut _, + timeout, + WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE, + ) + }; + + if ret == 0 { + // Registration failed - reconstruct Arc to drop the extra reference + unsafe { + let _ = std::sync::Arc::from_raw(data_ptr); + } + return Err(vm.new_last_os_error()); + } + + // Store in registry for cleanup tracking + let wait_handle = new_wait_object as isize; + if let Ok(mut registry) = wait_callback_registry().lock() { + registry.insert(wait_handle, data); + } + + Ok(wait_handle) + } + + // Helper to cleanup callback data when unregistering + // Just removes from registry - Arc ensures memory stays alive if callback is running + fn cleanup_wait_callback_data(wait_handle: isize) { + if let Ok(mut registry) = wait_callback_registry().lock() { + // Removing from registry drops one Arc reference + // If callback already ran, this frees the memory + // If callback is still pending/running, it holds the other reference + registry.remove(&wait_handle); + } + } + + #[pyfunction] + fn UnregisterWait(wait_handle: isize, vm: &VirtualMachine) -> PyResult<()> { + use windows_sys::Win32::System::Threading::UnregisterWait; + + let ret = unsafe { UnregisterWait(wait_handle as HANDLE) }; + // Cleanup callback data regardless of UnregisterWait result + // (callback may have already fired, or may never fire) + cleanup_wait_callback_data(wait_handle); + if ret == 0 { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + + #[pyfunction] + fn UnregisterWaitEx(wait_handle: isize, event: isize, vm: &VirtualMachine) -> PyResult<()> { + use windows_sys::Win32::System::Threading::UnregisterWaitEx; + + let ret = unsafe { UnregisterWaitEx(wait_handle as HANDLE, event as HANDLE) }; + // Cleanup callback data regardless of UnregisterWaitEx result + cleanup_wait_callback_data(wait_handle); + if ret == 0 { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + + #[pyfunction] + fn BindLocal(socket: isize, family: i32, vm: &VirtualMachine) -> PyResult<()> { + use windows_sys::Win32::Networking::WinSock::{INADDR_ANY, SOCKET_ERROR, bind}; + + let ret = if family == AF_INET as i32 { + let mut addr: SOCKADDR_IN = unsafe { std::mem::zeroed() }; + addr.sin_family = AF_INET; + addr.sin_port = 0; + addr.sin_addr.S_un.S_addr = INADDR_ANY; + unsafe { + bind( + socket as _, + &addr as *const _ as *const SOCKADDR, + std::mem::size_of::() as i32, + ) + } + } else if family == AF_INET6 as i32 { + // in6addr_any is all zeros, which we have from zeroed() + let mut addr: SOCKADDR_IN6 = unsafe { std::mem::zeroed() }; + addr.sin6_family = AF_INET6; + addr.sin6_port = 0; + unsafe { + bind( + socket as _, + &addr as *const _ as *const SOCKADDR, + std::mem::size_of::() as i32, + ) + } + } else { + return Err(vm.new_value_error("family must be AF_INET or AF_INET6".to_owned())); + }; + + if ret == SOCKET_ERROR { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + + #[pyfunction] + fn FormatMessage(error_code: u32, _vm: &VirtualMachine) -> PyResult { + use windows_sys::Win32::Foundation::LocalFree; + use windows_sys::Win32::System::Diagnostics::Debug::{ + FORMAT_MESSAGE_ALLOCATE_BUFFER, FORMAT_MESSAGE_FROM_SYSTEM, + FORMAT_MESSAGE_IGNORE_INSERTS, FormatMessageW, + }; + + // LANG_NEUTRAL = 0, SUBLANG_DEFAULT = 1 + const LANG_NEUTRAL: u32 = 0; + const SUBLANG_DEFAULT: u32 = 1; + + let mut buffer: *mut u16 = std::ptr::null_mut(); + + let len = unsafe { + FormatMessageW( + FORMAT_MESSAGE_ALLOCATE_BUFFER + | FORMAT_MESSAGE_FROM_SYSTEM + | FORMAT_MESSAGE_IGNORE_INSERTS, + std::ptr::null(), + error_code, + (SUBLANG_DEFAULT << 10) | LANG_NEUTRAL, + &mut buffer as *mut _ as *mut u16, + 0, + std::ptr::null(), + ) + }; + + if len == 0 || buffer.is_null() { + return Ok(format!("unknown error code {}", error_code)); + } + + // Convert to Rust string, trimming trailing whitespace + let slice = unsafe { std::slice::from_raw_parts(buffer, len as usize) }; + let msg = String::from_utf16_lossy(slice).trim_end().to_string(); + + unsafe { LocalFree(buffer as *mut _) }; + + Ok(msg) + } + + #[pyfunction] + fn WSAConnect(socket: isize, address: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + use windows_sys::Win32::Networking::WinSock::{SOCKET_ERROR, WSAConnect}; + + let (addr_bytes, addr_len) = parse_address(&address, vm)?; + + let ret = unsafe { + WSAConnect( + socket as _, + addr_bytes.as_ptr() as *const SOCKADDR, + addr_len, + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null(), + std::ptr::null(), + ) + }; + + if ret == SOCKET_ERROR { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + #[pyfunction] fn CreateEvent( event_attributes: PyObjectRef, @@ -549,22 +1833,22 @@ mod _overlapped { vm: &VirtualMachine, ) -> PyResult { if !vm.is_none(&event_attributes) { - return Err(vm.new_value_error("EventAttributes must be None")); + return Err(vm.new_value_error("EventAttributes must be None".to_owned())); } - let name = match name { - Some(name) => { - let name = widestring::WideCString::from_str(&name).unwrap(); - name.as_ptr() - } - None => core::ptr::null(), - }; + let name_wide: Option> = + name.map(|n| n.encode_utf16().chain(std::iter::once(0)).collect()); + let name_ptr = name_wide + .as_ref() + .map(|n| n.as_ptr()) + .unwrap_or(std::ptr::null()); + let event = unsafe { windows_sys::Win32::System::Threading::CreateEventW( - core::ptr::null(), - manual_reset as _, - initial_state as _, - name, + std::ptr::null(), + if manual_reset { 1 } else { 0 }, + if initial_state { 1 } else { 0 }, + name_ptr, ) as isize }; if event == NULL { @@ -574,8 +1858,8 @@ mod _overlapped { } #[pyfunction] - fn SetEvent(handle: u64, vm: &VirtualMachine) -> PyResult<()> { - let ret = unsafe { windows_sys::Win32::System::Threading::SetEvent(u64_to_handle(handle)) }; + fn SetEvent(handle: isize, vm: &VirtualMachine) -> PyResult<()> { + let ret = unsafe { windows_sys::Win32::System::Threading::SetEvent(handle as HANDLE) }; if ret == 0 { return Err(vm.new_last_os_error()); } @@ -583,9 +1867,8 @@ mod _overlapped { } #[pyfunction] - fn ResetEvent(handle: u64, vm: &VirtualMachine) -> PyResult<()> { - let ret = - unsafe { windows_sys::Win32::System::Threading::ResetEvent(u64_to_handle(handle)) }; + fn ResetEvent(handle: isize, vm: &VirtualMachine) -> PyResult<()> { + let ret = unsafe { windows_sys::Win32::System::Threading::ResetEvent(handle as HANDLE) }; if ret == 0 { return Err(vm.new_last_os_error()); } diff --git a/scripts/generate_opcode_metadata.py b/scripts/generate_opcode_metadata.py index dc913e1a7ee..42fb55a7c01 100644 --- a/scripts/generate_opcode_metadata.py +++ b/scripts/generate_opcode_metadata.py @@ -49,7 +49,7 @@ def extract_enum_body(contents: str, enum_name: str) -> str: ) -contents = BYTECODE_FILE.read_text() +contents = BYTECODE_FILE.read_text(encoding="utf-8") enum_body = "\n".join( extract_enum_body(contents, enum_name) for enum_name in ("Instruction", "PseudoInstruction") @@ -78,4 +78,4 @@ def extract_enum_body(contents: str, enum_name: str) -> str: MIN_INSTRUMENTED_OPCODE = 236 """ -OPCODE_METADATA_FILE.write_text(output) +OPCODE_METADATA_FILE.write_text(output, encoding="utf-8")