use ntapi::ntioapi::{IO_STATUS_BLOCK_u, IO_STATUS_BLOCK}; use ntapi::ntioapi::{NtCancelIoFileEx, NtDeviceIoControlFile}; use ntapi::ntrtl::RtlNtStatusToDosError; use std::fmt; use std::fs::File; use std::io; use std::mem::size_of; use std::os::windows::io::AsRawHandle; use std::ptr::null_mut; use winapi::shared::ntdef::{HANDLE, LARGE_INTEGER, NTSTATUS, PVOID, ULONG}; use winapi::shared::ntstatus::{STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS}; const IOCTL_AFD_POLL: ULONG = 0x00012024; /// Winsock2 AFD driver instance. /// /// All operations are unsafe due to IO_STATUS_BLOCK parameter are being used by Afd driver during STATUS_PENDING before I/O Completion Port returns its result. #[derive(Debug)] pub struct Afd { fd: File, } #[repr(C)] #[derive(Debug)] pub struct AfdPollHandleInfo { pub handle: HANDLE, pub events: ULONG, pub status: NTSTATUS, } unsafe impl Send for AfdPollHandleInfo {} #[repr(C)] pub struct AfdPollInfo { pub timeout: LARGE_INTEGER, // Can have only value 1. pub number_of_handles: ULONG, pub exclusive: ULONG, pub handles: [AfdPollHandleInfo; 1], } impl fmt::Debug for AfdPollInfo { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("AfdPollInfo").finish() } } impl Afd { /// Poll `Afd` instance with `AfdPollInfo`. /// /// # Unsafety /// /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`). /// `iosb` needs to be untouched after the call while operation is in effective at ALL TIME except for `cancel` method. /// So be careful not to `poll` twice while polling. /// User should deallocate there overlapped value when error to prevent memory leak. pub unsafe fn poll( &self, info: &mut AfdPollInfo, iosb: *mut IO_STATUS_BLOCK, overlapped: PVOID, ) -> io::Result { let info_ptr: PVOID = info as *mut _ as PVOID; (*iosb).u.Status = STATUS_PENDING; let status = NtDeviceIoControlFile( self.fd.as_raw_handle(), null_mut(), None, overlapped, iosb, IOCTL_AFD_POLL, info_ptr, size_of::() as u32, info_ptr, size_of::() as u32, ); match status { STATUS_SUCCESS => Ok(true), STATUS_PENDING => Ok(false), _ => Err(io::Error::from_raw_os_error( RtlNtStatusToDosError(status) as i32 )), } } /// Cancel previous polled request of `Afd`. /// /// iosb needs to be used by `poll` first for valid `cancel`. /// /// # Unsafety /// /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`). /// Use it only with request is still being polled so that you have valid `IO_STATUS_BLOCK` to use. /// User should NOT deallocate there overlapped value after the `cancel` to prevent double free. pub unsafe fn cancel(&self, iosb: *mut IO_STATUS_BLOCK) -> io::Result<()> { if (*iosb).u.Status != STATUS_PENDING { return Ok(()); } let mut cancel_iosb = IO_STATUS_BLOCK { u: IO_STATUS_BLOCK_u { Status: 0 }, Information: 0, }; let status = NtCancelIoFileEx(self.fd.as_raw_handle(), iosb, &mut cancel_iosb); if status == STATUS_SUCCESS || status == STATUS_NOT_FOUND { return Ok(()); } Err(io::Error::from_raw_os_error( RtlNtStatusToDosError(status) as i32 )) } } cfg_io_source! { use std::mem::zeroed; use std::os::windows::io::{FromRawHandle, RawHandle}; use std::sync::atomic::{AtomicUsize, Ordering}; use miow::iocp::CompletionPort; use ntapi::ntioapi::{NtCreateFile, FILE_OPEN}; use winapi::shared::ntdef::{OBJECT_ATTRIBUTES, UNICODE_STRING, USHORT, WCHAR}; use winapi::um::handleapi::INVALID_HANDLE_VALUE; use winapi::um::winbase::{SetFileCompletionNotificationModes, FILE_SKIP_SET_EVENT_ON_HANDLE}; use winapi::um::winnt::{SYNCHRONIZE, FILE_SHARE_READ, FILE_SHARE_WRITE}; const AFD_HELPER_ATTRIBUTES: OBJECT_ATTRIBUTES = OBJECT_ATTRIBUTES { Length: size_of::() as ULONG, RootDirectory: null_mut(), ObjectName: &AFD_OBJ_NAME as *const _ as *mut _, Attributes: 0, SecurityDescriptor: null_mut(), SecurityQualityOfService: null_mut(), }; const AFD_OBJ_NAME: UNICODE_STRING = UNICODE_STRING { Length: (AFD_HELPER_NAME.len() * size_of::()) as USHORT, MaximumLength: (AFD_HELPER_NAME.len() * size_of::()) as USHORT, Buffer: AFD_HELPER_NAME.as_ptr() as *mut _, }; const AFD_HELPER_NAME: &[WCHAR] = &[ '\\' as _, 'D' as _, 'e' as _, 'v' as _, 'i' as _, 'c' as _, 'e' as _, '\\' as _, 'A' as _, 'f' as _, 'd' as _, '\\' as _, 'M' as _, 'i' as _, 'o' as _ ]; static NEXT_TOKEN: AtomicUsize = AtomicUsize::new(0); impl AfdPollInfo { pub fn zeroed() -> AfdPollInfo { unsafe { zeroed() } } } impl Afd { /// Create new Afd instance. pub fn new(cp: &CompletionPort) -> io::Result { let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE; let mut iosb = IO_STATUS_BLOCK { u: IO_STATUS_BLOCK_u { Status: 0 }, Information: 0, }; unsafe { let status = NtCreateFile( &mut afd_helper_handle as *mut _, SYNCHRONIZE, &AFD_HELPER_ATTRIBUTES as *const _ as *mut _, &mut iosb, null_mut(), 0, FILE_SHARE_READ | FILE_SHARE_WRITE, FILE_OPEN, 0, null_mut(), 0, ); if status != STATUS_SUCCESS { let raw_err = io::Error::from_raw_os_error( RtlNtStatusToDosError(status) as i32 ); let msg = format!("Failed to open \\Device\\Afd\\Mio: {}", raw_err); return Err(io::Error::new(raw_err.kind(), msg)); } let fd = File::from_raw_handle(afd_helper_handle as RawHandle); // Increment by 2 to reserve space for other types of handles. // Non-AFD types (currently only NamedPipe), use odd numbered // tokens. This allows the selector to differentate between them // and dispatch events accordingly. let token = NEXT_TOKEN.fetch_add(2, Ordering::Relaxed) + 2; let afd = Afd { fd }; cp.add_handle(token, &afd.fd)?; match SetFileCompletionNotificationModes( afd_helper_handle, FILE_SKIP_SET_EVENT_ON_HANDLE, ) { 0 => Err(io::Error::last_os_error()), _ => Ok(afd), } } } } } pub const POLL_RECEIVE: u32 = 0b0_0000_0001; pub const POLL_RECEIVE_EXPEDITED: u32 = 0b0_0000_0010; pub const POLL_SEND: u32 = 0b0_0000_0100; pub const POLL_DISCONNECT: u32 = 0b0_0000_1000; pub const POLL_ABORT: u32 = 0b0_0001_0000; pub const POLL_LOCAL_CLOSE: u32 = 0b0_0010_0000; // Not used as it indicated in each event where a connection is connected, not // just the first time a connection is established. // Also see https://github.com/piscisaureus/wepoll/commit/8b7b340610f88af3d83f40fb728e7b850b090ece. pub const POLL_CONNECT: u32 = 0b0_0100_0000; pub const POLL_ACCEPT: u32 = 0b0_1000_0000; pub const POLL_CONNECT_FAIL: u32 = 0b1_0000_0000; pub const KNOWN_EVENTS: u32 = POLL_RECEIVE | POLL_RECEIVE_EXPEDITED | POLL_SEND | POLL_DISCONNECT | POLL_ABORT | POLL_LOCAL_CLOSE | POLL_ACCEPT | POLL_CONNECT_FAIL;