vendor/mio-0.8.5/src/sys/windows/afd.rs - toolchain/rustc - Git at Google

 use std::ffi::c_void;
 use std::fmt;
 use std::fs::File;
 use std::io;
 use std::mem::size_of;
 use std::os::windows::io::AsRawHandle;

 use windows_sys::Win32::Foundation::{
     RtlNtStatusToDosError, HANDLE, NTSTATUS, STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS,
 };
 use windows_sys::Win32::System::WindowsProgramming::{
     NtDeviceIoControlFile, IO_STATUS_BLOCK, IO_STATUS_BLOCK_0,
 };

 const IOCTL_AFD_POLL: u32 = 0x00012024;

 #[link(name = "ntdll")]
 extern "system" {
     /// See <https://processhacker.sourceforge.io/doc/ntioapi_8h.html#a0d4d550cad4d62d75b76961e25f6550c>
     ///
     /// This is an undocumented API and as such not part of <https://github.com/microsoft/win32metadata>
     /// from which `windows-sys` is generated, and also unlikely to be added, so
     /// we manually declare it here
     fn NtCancelIoFileEx(
         FileHandle: HANDLE,
         IoRequestToCancel: *mut IO_STATUS_BLOCK,
         IoStatusBlock: *mut IO_STATUS_BLOCK,
     ) -> NTSTATUS;
 }
 /// 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: u32,
     pub status: NTSTATUS,
 }

 unsafe impl Send for AfdPollHandleInfo {}

 #[repr(C)]
 pub struct AfdPollInfo {
     pub timeout: i64,
     // Can have only value 1.
     pub number_of_handles: u32,
     pub exclusive: u32,
     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: *mut c_void,
     ) -> io::Result<bool> {
         let info_ptr = info as *mut _ as *mut c_void;
         (*iosb).Anonymous.Status = STATUS_PENDING;
         let status = NtDeviceIoControlFile(
             self.fd.as_raw_handle() as HANDLE,
             0,
             None,
             overlapped,
             iosb,
             IOCTL_AFD_POLL,
             info_ptr,
             size_of::<AfdPollInfo>() as u32,
             info_ptr,
             size_of::<AfdPollInfo>() 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).Anonymous.Status != STATUS_PENDING {
             return Ok(());
         }

         let mut cancel_iosb = IO_STATUS_BLOCK {
             Anonymous: IO_STATUS_BLOCK_0 { Status: 0 },
             Information: 0,
         };
         let status = NtCancelIoFileEx(self.fd.as_raw_handle() as 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::ptr::null_mut;
     use std::sync::atomic::{AtomicUsize, Ordering};

     use super::iocp::CompletionPort;
     use windows_sys::Win32::{
         Foundation::{UNICODE_STRING, INVALID_HANDLE_VALUE},
         System::WindowsProgramming::{
             OBJECT_ATTRIBUTES, FILE_SKIP_SET_EVENT_ON_HANDLE,
         },
         Storage::FileSystem::{FILE_OPEN, NtCreateFile, SetFileCompletionNotificationModes, SYNCHRONIZE, FILE_SHARE_READ, FILE_SHARE_WRITE},
     };

     const AFD_HELPER_ATTRIBUTES: OBJECT_ATTRIBUTES = OBJECT_ATTRIBUTES {
         Length: size_of::<OBJECT_ATTRIBUTES>() as u32,
         RootDirectory: 0,
         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::<u16>()) as u16,
         MaximumLength: (AFD_HELPER_NAME.len() * size_of::<u16>()) as u16,
         Buffer: AFD_HELPER_NAME.as_ptr() as *mut _,
     };

     const AFD_HELPER_NAME: &[u16] = &[
         '\\' 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(crate) fn new(cp: &CompletionPort) -> io::Result<Afd> {
             let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE;
             let mut iosb = IO_STATUS_BLOCK {
                 Anonymous: IO_STATUS_BLOCK_0 { 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 as u8 // This is just 2, so fits in u8
                 ) {
                     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;
	use std::ffi::c_void;
	use std::fmt;
	use std::fs::File;
	use std::io;
	use std::mem::size_of;
	use std::os::windows::io::AsRawHandle;

	use windows_sys::Win32::Foundation::{
	RtlNtStatusToDosError, HANDLE, NTSTATUS, STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS,
	};
	use windows_sys::Win32::System::WindowsProgramming::{
	NtDeviceIoControlFile, IO_STATUS_BLOCK, IO_STATUS_BLOCK_0,
	};

	const IOCTL_AFD_POLL: u32 = 0x00012024;

	#[link(name = "ntdll")]
	extern "system" {
	/// See <https://processhacker.sourceforge.io/doc/ntioapi_8h.html#a0d4d550cad4d62d75b76961e25f6550c>
	///
	/// This is an undocumented API and as such not part of <https://github.com/microsoft/win32metadata>
	/// from which `windows-sys` is generated, and also unlikely to be added, so
	/// we manually declare it here
	fn NtCancelIoFileEx(
	FileHandle: HANDLE,
	IoRequestToCancel: *mut IO_STATUS_BLOCK,
	IoStatusBlock: *mut IO_STATUS_BLOCK,
	) -> NTSTATUS;
	}
	/// 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: u32,
	pub status: NTSTATUS,
	}

	unsafe impl Send for AfdPollHandleInfo {}

	#[repr(C)]
	pub struct AfdPollInfo {
	pub timeout: i64,
	// Can have only value 1.
	pub number_of_handles: u32,
	pub exclusive: u32,
	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: *mut c_void,
	) -> io::Result<bool> {
	let info_ptr = info as mut _ as mut c_void;
	(*iosb).Anonymous.Status = STATUS_PENDING;
	let status = NtDeviceIoControlFile(
	self.fd.as_raw_handle() as HANDLE,
	0,
	None,
	overlapped,
	iosb,
	IOCTL_AFD_POLL,
	info_ptr,
	size_of::<AfdPollInfo>() as u32,
	info_ptr,
	size_of::<AfdPollInfo>() 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).Anonymous.Status != STATUS_PENDING {
	return Ok(());
	}

	let mut cancel_iosb = IO_STATUS_BLOCK {
	Anonymous: IO_STATUS_BLOCK_0 { Status: 0 },
	Information: 0,
	};
	let status = NtCancelIoFileEx(self.fd.as_raw_handle() as 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::ptr::null_mut;
	use std::sync::atomic::{AtomicUsize, Ordering};

	use super::iocp::CompletionPort;
	use windows_sys::Win32::{
	Foundation::{UNICODE_STRING, INVALID_HANDLE_VALUE},
	System::WindowsProgramming::{
	OBJECT_ATTRIBUTES, FILE_SKIP_SET_EVENT_ON_HANDLE,
	},
	Storage::FileSystem::{FILE_OPEN, NtCreateFile, SetFileCompletionNotificationModes, SYNCHRONIZE, FILE_SHARE_READ, FILE_SHARE_WRITE},
	};

	const AFD_HELPER_ATTRIBUTES: OBJECT_ATTRIBUTES = OBJECT_ATTRIBUTES {
	Length: size_of::<OBJECT_ATTRIBUTES>() as u32,
	RootDirectory: 0,
	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::<u16>()) as u16,
	MaximumLength: (AFD_HELPER_NAME.len() * size_of::<u16>()) as u16,
	Buffer: AFD_HELPER_NAME.as_ptr() as *mut _,
	};

	const AFD_HELPER_NAME: &[u16] = &[
	'\\' 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(crate) fn new(cp: &CompletionPort) -> io::Result<Afd> {
	let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE;
	let mut iosb = IO_STATUS_BLOCK {
	Anonymous: IO_STATUS_BLOCK_0 { 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 as u8 // This is just 2, so fits in u8
	) {
	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;