compio_net\socket/
windows.rs

1//! Windows-specific per-socket state.
2
3use 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
31/// Strips the [`Option`] from the `LPFN_*` aliases of `windows-sys`: the ioctl
32/// writes a bare pointer, a null one meaning "unsupported".
33trait 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/// The winsock extension functions of one socket.
48#[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
89/// Query the extension function `guid` of the provider serving `socket`.
90///
91/// `Ok(None)` means the provider does not implement it.
92fn 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/// Per-socket state, holding the winsock extension cache.
115#[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}
compio_net\socket/windows.rs

compio_net\socket/
windows.rs
