1use std::{
4 ffi::c_void,
5 io,
6 mem::{size_of, size_of_val},
7 os::windows::io::AsRawSocket,
8 ptr::null_mut,
9 task::Poll,
10};
11
12use compio_driver::{AsSocketState, win32_result, winsock_result};
13use once_cell::sync::OnceCell as OnceLock;
14use socket2::socklen_t;
15use windows_sys::{
16 Win32::{
17 Networking::WinSock::{
18 LPFN_ACCEPTEX, LPFN_CONNECTEX, LPFN_DISCONNECTEX, LPFN_GETACCEPTEXSOCKADDRS,
19 LPFN_WSARECVMSG, LPWSAOVERLAPPED_COMPLETION_ROUTINE,
20 SIO_GET_EXTENSION_FUNCTION_POINTER, SOCKADDR, SOCKET, SOCKET_ERROR, WSAGetLastError,
21 WSAID_ACCEPTEX, WSAID_CONNECTEX, WSAID_DISCONNECTEX, WSAID_GETACCEPTEXSOCKADDRS,
22 WSAID_WSARECVMSG, WSAIoctl, WSAMSG,
23 },
24 System::IO::OVERLAPPED,
25 },
26 core::GUID,
27};
28
29use super::SocketWithState;
30
31trait FnPtr {
34 type Bare;
35}
36
37impl<T> FnPtr for Option<T> {
38 type Bare = T;
39}
40
41type AcceptEx = <LPFN_ACCEPTEX as FnPtr>::Bare;
42type GetAcceptExSockaddrs = <LPFN_GETACCEPTEXSOCKADDRS as FnPtr>::Bare;
43type ConnectEx = <LPFN_CONNECTEX as FnPtr>::Bare;
44type DisconnectEx = <LPFN_DISCONNECTEX as FnPtr>::Bare;
45type RecvMsg = <LPFN_WSARECVMSG as FnPtr>::Bare;
46
47#[derive(Debug, Default)]
49struct WsaFns {
50 accept_ex: OnceLock<Option<AcceptEx>>,
51 get_accept_ex_sockaddrs: OnceLock<Option<GetAcceptExSockaddrs>>,
52 connect_ex: OnceLock<Option<ConnectEx>>,
53 disconnect_ex: OnceLock<Option<DisconnectEx>>,
54 recv_msg: OnceLock<Option<RecvMsg>>,
55}
56
57impl WsaFns {
58 fn accept_ex(&self, socket: SOCKET) -> io::Result<LPFN_ACCEPTEX> {
59 Ok(*self
60 .accept_ex
61 .get_or_try_init(|| get_wsa_fn::<AcceptEx>(socket, WSAID_ACCEPTEX))?)
62 }
63
64 fn get_accept_ex_sockaddrs(&self, socket: SOCKET) -> io::Result<LPFN_GETACCEPTEXSOCKADDRS> {
65 Ok(*self.get_accept_ex_sockaddrs.get_or_try_init(|| {
66 get_wsa_fn::<GetAcceptExSockaddrs>(socket, WSAID_GETACCEPTEXSOCKADDRS)
67 })?)
68 }
69
70 fn connect_ex(&self, socket: SOCKET) -> io::Result<LPFN_CONNECTEX> {
71 Ok(*self
72 .connect_ex
73 .get_or_try_init(|| get_wsa_fn::<ConnectEx>(socket, WSAID_CONNECTEX))?)
74 }
75
76 fn disconnect_ex(&self, socket: SOCKET) -> io::Result<LPFN_DISCONNECTEX> {
77 Ok(*self
78 .disconnect_ex
79 .get_or_try_init(|| get_wsa_fn::<DisconnectEx>(socket, WSAID_DISCONNECTEX))?)
80 }
81
82 fn recv_msg(&self, socket: SOCKET) -> io::Result<LPFN_WSARECVMSG> {
83 Ok(*self
84 .recv_msg
85 .get_or_try_init(|| get_wsa_fn::<RecvMsg>(socket, WSAID_WSARECVMSG))?)
86 }
87}
88
89fn get_wsa_fn<F>(socket: SOCKET, guid: GUID) -> io::Result<Option<F>> {
93 let mut fptr = None;
94 let mut returned = 0;
95 let res = unsafe {
96 WSAIoctl(
97 socket,
98 SIO_GET_EXTENSION_FUNCTION_POINTER,
99 std::ptr::addr_of!(guid).cast(),
100 size_of_val(&guid) as _,
101 std::ptr::addr_of_mut!(fptr).cast(),
102 size_of::<F>() as _,
103 &mut returned,
104 null_mut(),
105 None,
106 )
107 };
108 if res == SOCKET_ERROR {
109 return Err(io::Error::from_raw_os_error(unsafe { WSAGetLastError() }));
110 }
111 Ok(fptr)
112}
113
114#[derive(Debug, Default)]
116pub(super) struct SocketState {
117 wsa: Box<WsaFns>,
118}
119
120impl SocketState {
121 pub(super) fn new() -> Self {
122 SocketState::default()
123 }
124}
125
126impl AsSocketState for SocketWithState {
127 unsafe fn wsa_accept_ex(
128 &self,
129 accept_socket: SOCKET,
130 output_buffer: *mut c_void,
131 receive_data_length: u32,
132 local_address_length: u32,
133 remote_address_length: u32,
134 bytes_received: *mut u32,
135 overlapped: *mut OVERLAPPED,
136 ) -> Poll<io::Result<usize>> {
137 let socket = self.inner.as_raw_socket() as SOCKET;
138 let f = match self.state.wsa.accept_ex(socket) {
139 Ok(Some(f)) => f,
140 Ok(None) => return Poll::Ready(Err(unsupported())),
141 Err(e) => return Poll::Ready(Err(e)),
142 };
143 let res = unsafe {
144 f(
145 socket,
146 accept_socket,
147 output_buffer,
148 receive_data_length,
149 local_address_length,
150 remote_address_length,
151 bytes_received,
152 overlapped,
153 )
154 };
155 win32_result(res, unsafe { *bytes_received })
156 }
157
158 unsafe fn wsa_get_accept_ex_sockaddrs(
159 &self,
160 output_buffer: *const c_void,
161 receive_data_length: u32,
162 local_address_length: u32,
163 remote_address_length: u32,
164 local_sockaddr: *mut *mut SOCKADDR,
165 local_sockaddr_length: *mut socklen_t,
166 remote_sockaddr: *mut *mut SOCKADDR,
167 remote_sockaddr_length: *mut socklen_t,
168 ) -> io::Result<()> {
169 let Some(f) = self
170 .state
171 .wsa
172 .get_accept_ex_sockaddrs(self.inner.as_raw_socket() as SOCKET)?
173 else {
174 return Err(unsupported());
175 };
176 unsafe {
177 f(
178 output_buffer.cast_mut(),
179 receive_data_length,
180 local_address_length,
181 remote_address_length,
182 local_sockaddr,
183 local_sockaddr_length,
184 remote_sockaddr,
185 remote_sockaddr_length,
186 )
187 };
188 Ok(())
189 }
190
191 unsafe fn wsa_connect_ex(
192 &self,
193 name: *const SOCKADDR,
194 name_length: socklen_t,
195 send_buffer: *const c_void,
196 send_data_length: u32,
197 bytes_sent: *mut u32,
198 overlapped: *mut OVERLAPPED,
199 ) -> Poll<io::Result<usize>> {
200 let socket = self.inner.as_raw_socket() as SOCKET;
201 let f = match self.state.wsa.connect_ex(socket) {
202 Ok(Some(f)) => f,
203 Ok(None) => return Poll::Ready(Err(unsupported())),
204 Err(e) => return Poll::Ready(Err(e)),
205 };
206 let res = unsafe {
207 f(
208 socket,
209 name,
210 name_length,
211 send_buffer,
212 send_data_length,
213 bytes_sent,
214 overlapped,
215 )
216 };
217 win32_result(res, unsafe { *bytes_sent })
218 }
219
220 unsafe fn wsa_disconnect_ex(
221 &self,
222 overlapped: *mut OVERLAPPED,
223 flags: u32,
224 reserved: u32,
225 ) -> Poll<io::Result<usize>> {
226 let socket = self.inner.as_raw_socket() as SOCKET;
227 let f = match self.state.wsa.disconnect_ex(socket) {
228 Ok(Some(f)) => f,
229 Ok(None) => return Poll::Ready(Err(unsupported())),
230 Err(e) => return Poll::Ready(Err(e)),
231 };
232 let res = unsafe { f(socket, overlapped, flags, reserved) };
233 win32_result(res, 0)
234 }
235
236 unsafe fn wsa_recv_msg(
237 &self,
238 msg: *mut WSAMSG,
239 bytes_received: *mut u32,
240 overlapped: *mut OVERLAPPED,
241 completion_routine: LPWSAOVERLAPPED_COMPLETION_ROUTINE,
242 ) -> Poll<io::Result<usize>> {
243 let socket = self.inner.as_raw_socket() as SOCKET;
244 let f = match self.state.wsa.recv_msg(socket) {
245 Ok(Some(f)) => f,
246 Ok(None) => return Poll::Ready(Err(unsupported())),
247 Err(e) => return Poll::Ready(Err(e)),
248 };
249 let res = unsafe { f(socket, msg, bytes_received, overlapped, completion_routine) };
250 winsock_result(res, unsafe { *bytes_received })
251 }
252}
253
254fn unsupported() -> io::Error {
255 io::Error::new(
256 io::ErrorKind::Unsupported,
257 "the socket does not provide this winsock extension",
258 )
259}