1 use ntapi::ntioapi::{IO_STATUS_BLOCK_u, IO_STATUS_BLOCK}; 2 use ntapi::ntioapi::{NtCancelIoFileEx, NtDeviceIoControlFile}; 3 use ntapi::ntrtl::RtlNtStatusToDosError; 4 use std::fmt; 5 use std::fs::File; 6 use std::io; 7 use std::mem::size_of; 8 use std::os::windows::io::AsRawHandle; 9 use std::ptr::null_mut; 10 use winapi::shared::ntdef::{HANDLE, LARGE_INTEGER, NTSTATUS, PVOID, ULONG}; 11 use winapi::shared::ntstatus::{STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS}; 12 13 const IOCTL_AFD_POLL: ULONG = 0x00012024; 14 15 /// Winsock2 AFD driver instance. 16 /// 17 /// 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. 18 #[derive(Debug)] 19 pub struct Afd { 20 fd: File, 21 } 22 23 #[repr(C)] 24 #[derive(Debug)] 25 pub struct AfdPollHandleInfo { 26 pub handle: HANDLE, 27 pub events: ULONG, 28 pub status: NTSTATUS, 29 } 30 31 unsafe impl Send for AfdPollHandleInfo {} 32 33 #[repr(C)] 34 pub struct AfdPollInfo { 35 pub timeout: LARGE_INTEGER, 36 // Can have only value 1. 37 pub number_of_handles: ULONG, 38 pub exclusive: ULONG, 39 pub handles: [AfdPollHandleInfo; 1], 40 } 41 42 impl fmt::Debug for AfdPollInfo { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result43 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 44 f.debug_struct("AfdPollInfo").finish() 45 } 46 } 47 48 impl Afd { 49 /// Poll `Afd` instance with `AfdPollInfo`. 50 /// 51 /// # Unsafety 52 /// 53 /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`). 54 /// `iosb` needs to be untouched after the call while operation is in effective at ALL TIME except for `cancel` method. 55 /// So be careful not to `poll` twice while polling. 56 /// User should deallocate there overlapped value when error to prevent memory leak. poll( &self, info: &mut AfdPollInfo, iosb: *mut IO_STATUS_BLOCK, overlapped: PVOID, ) -> io::Result<bool>57 pub unsafe fn poll( 58 &self, 59 info: &mut AfdPollInfo, 60 iosb: *mut IO_STATUS_BLOCK, 61 overlapped: PVOID, 62 ) -> io::Result<bool> { 63 let info_ptr: PVOID = info as *mut _ as PVOID; 64 (*iosb).u.Status = STATUS_PENDING; 65 let status = NtDeviceIoControlFile( 66 self.fd.as_raw_handle(), 67 null_mut(), 68 None, 69 overlapped, 70 iosb, 71 IOCTL_AFD_POLL, 72 info_ptr, 73 size_of::<AfdPollInfo>() as u32, 74 info_ptr, 75 size_of::<AfdPollInfo>() as u32, 76 ); 77 match status { 78 STATUS_SUCCESS => Ok(true), 79 STATUS_PENDING => Ok(false), 80 _ => Err(io::Error::from_raw_os_error( 81 RtlNtStatusToDosError(status) as i32 82 )), 83 } 84 } 85 86 /// Cancel previous polled request of `Afd`. 87 /// 88 /// iosb needs to be used by `poll` first for valid `cancel`. 89 /// 90 /// # Unsafety 91 /// 92 /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`). 93 /// Use it only with request is still being polled so that you have valid `IO_STATUS_BLOCK` to use. 94 /// User should NOT deallocate there overlapped value after the `cancel` to prevent double free. cancel(&self, iosb: *mut IO_STATUS_BLOCK) -> io::Result<()>95 pub unsafe fn cancel(&self, iosb: *mut IO_STATUS_BLOCK) -> io::Result<()> { 96 if (*iosb).u.Status != STATUS_PENDING { 97 return Ok(()); 98 } 99 100 let mut cancel_iosb = IO_STATUS_BLOCK { 101 u: IO_STATUS_BLOCK_u { Status: 0 }, 102 Information: 0, 103 }; 104 let status = NtCancelIoFileEx(self.fd.as_raw_handle(), iosb, &mut cancel_iosb); 105 if status == STATUS_SUCCESS || status == STATUS_NOT_FOUND { 106 return Ok(()); 107 } 108 Err(io::Error::from_raw_os_error( 109 RtlNtStatusToDosError(status) as i32 110 )) 111 } 112 } 113 114 cfg_io_source! { 115 use std::mem::zeroed; 116 use std::os::windows::io::{FromRawHandle, RawHandle}; 117 use std::sync::atomic::{AtomicUsize, Ordering}; 118 119 use miow::iocp::CompletionPort; 120 use ntapi::ntioapi::{NtCreateFile, FILE_OPEN}; 121 use winapi::shared::ntdef::{OBJECT_ATTRIBUTES, UNICODE_STRING, USHORT, WCHAR}; 122 use winapi::um::handleapi::INVALID_HANDLE_VALUE; 123 use winapi::um::winbase::{SetFileCompletionNotificationModes, FILE_SKIP_SET_EVENT_ON_HANDLE}; 124 use winapi::um::winnt::{SYNCHRONIZE, FILE_SHARE_READ, FILE_SHARE_WRITE}; 125 126 const AFD_HELPER_ATTRIBUTES: OBJECT_ATTRIBUTES = OBJECT_ATTRIBUTES { 127 Length: size_of::<OBJECT_ATTRIBUTES>() as ULONG, 128 RootDirectory: null_mut(), 129 ObjectName: &AFD_OBJ_NAME as *const _ as *mut _, 130 Attributes: 0, 131 SecurityDescriptor: null_mut(), 132 SecurityQualityOfService: null_mut(), 133 }; 134 135 const AFD_OBJ_NAME: UNICODE_STRING = UNICODE_STRING { 136 Length: (AFD_HELPER_NAME.len() * size_of::<WCHAR>()) as USHORT, 137 MaximumLength: (AFD_HELPER_NAME.len() * size_of::<WCHAR>()) as USHORT, 138 Buffer: AFD_HELPER_NAME.as_ptr() as *mut _, 139 }; 140 141 const AFD_HELPER_NAME: &[WCHAR] = &[ 142 '\\' as _, 143 'D' as _, 144 'e' as _, 145 'v' as _, 146 'i' as _, 147 'c' as _, 148 'e' as _, 149 '\\' as _, 150 'A' as _, 151 'f' as _, 152 'd' as _, 153 '\\' as _, 154 'M' as _, 155 'i' as _, 156 'o' as _ 157 ]; 158 159 static NEXT_TOKEN: AtomicUsize = AtomicUsize::new(0); 160 161 impl AfdPollInfo { 162 pub fn zeroed() -> AfdPollInfo { 163 unsafe { zeroed() } 164 } 165 } 166 167 impl Afd { 168 /// Create new Afd instance. 169 pub fn new(cp: &CompletionPort) -> io::Result<Afd> { 170 let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE; 171 let mut iosb = IO_STATUS_BLOCK { 172 u: IO_STATUS_BLOCK_u { Status: 0 }, 173 Information: 0, 174 }; 175 176 unsafe { 177 let status = NtCreateFile( 178 &mut afd_helper_handle as *mut _, 179 SYNCHRONIZE, 180 &AFD_HELPER_ATTRIBUTES as *const _ as *mut _, 181 &mut iosb, 182 null_mut(), 183 0 as ULONG, 184 FILE_SHARE_READ | FILE_SHARE_WRITE, 185 FILE_OPEN, 186 0 as ULONG, 187 null_mut(), 188 0 as ULONG, 189 ); 190 if status != STATUS_SUCCESS { 191 let raw_err = io::Error::from_raw_os_error( 192 RtlNtStatusToDosError(status) as i32 193 ); 194 let msg = format!("Failed to open \\Device\\Afd\\Mio: {}", raw_err); 195 return Err(io::Error::new(raw_err.kind(), msg)); 196 } 197 let fd = File::from_raw_handle(afd_helper_handle as RawHandle); 198 // Increment by 2 to reserve space for other types of handles. 199 // Non-AFD types (currently only NamedPipe), use odd numbered 200 // tokens. This allows the selector to differentate between them 201 // and dispatch events accordingly. 202 let token = NEXT_TOKEN.fetch_add(2, Ordering::Relaxed) + 2; 203 let afd = Afd { fd }; 204 cp.add_handle(token, &afd.fd)?; 205 match SetFileCompletionNotificationModes( 206 afd_helper_handle, 207 FILE_SKIP_SET_EVENT_ON_HANDLE, 208 ) { 209 0 => Err(io::Error::last_os_error()), 210 _ => Ok(afd), 211 } 212 } 213 } 214 } 215 } 216 217 pub const POLL_RECEIVE: u32 = 0b000_000_001; 218 pub const POLL_RECEIVE_EXPEDITED: u32 = 0b000_000_010; 219 pub const POLL_SEND: u32 = 0b000_000_100; 220 pub const POLL_DISCONNECT: u32 = 0b000_001_000; 221 pub const POLL_ABORT: u32 = 0b000_010_000; 222 pub const POLL_LOCAL_CLOSE: u32 = 0b000_100_000; 223 // Not used as it indicated in each event where a connection is connected, not 224 // just the first time a connection is established. 225 // Also see https://github.com/piscisaureus/wepoll/commit/8b7b340610f88af3d83f40fb728e7b850b090ece. 226 pub const POLL_CONNECT: u32 = 0b001_000_000; 227 pub const POLL_ACCEPT: u32 = 0b010_000_000; 228 pub const POLL_CONNECT_FAIL: u32 = 0b100_000_000; 229 230 pub const KNOWN_EVENTS: u32 = POLL_RECEIVE 231 | POLL_RECEIVE_EXPEDITED 232 | POLL_SEND 233 | POLL_DISCONNECT 234 | POLL_ABORT 235 | POLL_LOCAL_CLOSE 236 | POLL_ACCEPT 237 | POLL_CONNECT_FAIL; 238