compio_net/socket/
mod.rs

1use std::{
2    future::Future,
3    io,
4    mem::{ManuallyDrop, MaybeUninit},
5    ops::Deref,
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10use compio_buf::{
11    BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut, SetLenExt, buf_try,
12};
13#[cfg(windows)]
14use compio_driver::op::Disconnect;
15#[cfg(unix)]
16use compio_driver::op::{Bind, CreateSocket, Listen, ShutdownSocket};
17use compio_driver::{
18    AsFd, AsRawFd, BorrowedFd, BufferRef, OpCode, RawFd, ResultTakeBuffer, SharedFd, TakeBuffer,
19    ToSharedFd,
20    op::{
21        Accept, BufResultExt, CloseSocket, Connect, Recv, RecvFlags, RecvFrom, RecvFromManaged,
22        RecvFromMulti, RecvFromMultiResult, RecvFromVectored, RecvManaged, RecvMsg, RecvMsgManaged,
23        RecvMsgMulti, RecvMsgMultiResult, RecvMulti, RecvResultExt, RecvVectored, ReturnFlags,
24        Send, SendFlags, SendMsg, SendMsgZc, SendTo, SendToVectored, SendToVectoredZc, SendToZc,
25        SendVectored, SendVectoredZc, SendZc, VecBufResultExt,
26    },
27    syscall,
28};
29use compio_runtime::{Attacher, Runtime, SubmitMulti, SubmitMultiStream, fd::PollFd};
30use futures_util::{Stream, StreamExt};
31use pin_project_lite::pin_project;
32use socket2::{Domain, Protocol, SockAddr, Socket as Socket2, Type};
33use sys::SocketState;
34
35use crate::Incoming;
36
37cfg_select! {
38    any(
39        target_os = "linux",
40        target_os = "android",
41        target_os = "hurd",
42        target_os = "dragonfly",
43        target_os = "freebsd",
44        target_os = "openbsd",
45        target_os = "netbsd",
46        target_os = "solaris",
47        target_os = "illumos",
48        target_os = "haiku",
49        target_os = "nto",
50        target_os = "cygwin"
51    ) => {
52        pub(crate) const MSG_NOSIGNAL: SendFlags =
53            SendFlags::from_bits_retain(libc::MSG_NOSIGNAL as _);
54    }
55    _ => {
56        pub(crate) const MSG_NOSIGNAL: SendFlags = SendFlags::empty();
57    }
58}
59
60cfg_select! {
61    windows => {
62        #[path = "windows.rs"]
63        mod sys;
64    }
65    target_os = "linux" => {
66        #[path = "linux.rs"]
67        mod sys;
68    }
69    _ => {
70        #[path = "stub.rs"]
71        mod sys;
72    }
73}
74
75/// A socket handle together with the state that belongs to it.
76#[derive(Debug)]
77pub struct SocketWithState {
78    inner: Socket2,
79    #[allow(dead_code)]
80    state: SocketState,
81}
82
83impl SocketWithState {
84    fn new(inner: Socket2) -> Self {
85        Self {
86            inner,
87            state: SocketState::new(),
88        }
89    }
90
91    fn into_inner(self) -> Socket2 {
92        self.inner
93    }
94}
95
96impl Deref for SocketWithState {
97    type Target = Socket2;
98
99    fn deref(&self) -> &Self::Target {
100        &self.inner
101    }
102}
103
104impl AsRawFd for SocketWithState {
105    fn as_raw_fd(&self) -> RawFd {
106        self.inner.as_raw_fd()
107    }
108}
109
110impl AsFd for SocketWithState {
111    fn as_fd(&self) -> BorrowedFd<'_> {
112        self.inner.as_fd()
113    }
114}
115
116#[derive(Debug, Clone)]
117pub struct Socket {
118    pub(crate) socket: Attacher<SocketWithState>,
119}
120
121impl Socket {
122    pub fn from_socket2(socket: Socket2) -> io::Result<Self> {
123        Ok(Self {
124            socket: Attacher::new(SocketWithState::new(socket))?,
125        })
126    }
127
128    pub fn peer_addr(&self) -> io::Result<SockAddr> {
129        self.socket.peer_addr()
130    }
131
132    pub fn local_addr(&self) -> io::Result<SockAddr> {
133        self.socket.local_addr()
134    }
135
136    pub fn to_poll_fd(&self) -> io::Result<PollFd<SocketWithState>> {
137        PollFd::from_shared_fd(self.to_shared_fd())
138    }
139
140    pub fn into_poll_fd(self) -> io::Result<PollFd<SocketWithState>> {
141        PollFd::from_shared_fd(self.socket.into_inner())
142    }
143
144    #[cfg(windows)]
145    pub async fn new(domain: Domain, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
146        Self::from_socket2(Socket2::new(domain, ty, protocol)?)
147    }
148
149    #[cfg(unix)]
150    pub async fn new(domain: Domain, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
151        let op = CreateSocket::new(
152            domain.into(),
153            ty.into(),
154            protocol.map(|p| p.into()).unwrap_or_default(),
155        );
156        let (_, op) = buf_try!(@try compio_runtime::submit(op).await);
157
158        Self::from_socket2(op.into_inner())
159    }
160
161    pub async fn bind(&self, addr: &SockAddr) -> io::Result<()> {
162        #[cfg(not(unix))]
163        self.socket.bind(addr)?;
164        #[cfg(unix)]
165        {
166            let op = Bind::new(self.to_shared_fd(), addr.clone());
167            compio_runtime::submit(op).await.0?;
168        }
169        Ok(())
170    }
171
172    pub async fn listen(&self, backlog: i32) -> io::Result<()> {
173        #[cfg(not(unix))]
174        {
175            self.socket.listen(backlog)
176        }
177        #[cfg(unix)]
178        {
179            let op = Listen::new(self.to_shared_fd(), backlog);
180            compio_runtime::submit(op).await.0?;
181            Ok(())
182        }
183    }
184
185    pub fn connect(&self, addr: &SockAddr) -> io::Result<()> {
186        self.socket.connect(addr)
187    }
188
189    pub async fn connect_async(&self, addr: &SockAddr) -> io::Result<()> {
190        let op = Connect::new(self.to_shared_fd(), addr.clone());
191        let (_, _op) = buf_try!(@try compio_runtime::submit(op).await);
192        #[cfg(windows)]
193        _op.update_context()?;
194        Ok(())
195    }
196
197    #[cfg(windows)]
198    pub async fn disconnect(&self) -> io::Result<()> {
199        let op = Disconnect::new(self.to_shared_fd());
200        compio_runtime::submit(op).await.0.map(|_| ())
201    }
202
203    #[cfg(unix)]
204    pub async fn accept(&self) -> io::Result<(Self, SockAddr)> {
205        let op = Accept::new(self.to_shared_fd());
206        let (BufResult(res, op), extra) = compio_runtime::submit(op).with_extra().await;
207        res?;
208        extra.update_socket_state(&*self.socket);
209        let (accept_sock, addr) = op.into_inner();
210        let accept_sock = Self::from_socket2(accept_sock)?;
211        Ok((accept_sock, addr))
212    }
213
214    #[cfg(windows)]
215    pub async fn accept(&self) -> io::Result<(Self, SockAddr)> {
216        let domain = self.local_addr()?.domain();
217        let ty = self.socket.r#type()?;
218        let protocol = self.socket.protocol()?;
219        let accept_sock = Socket2::new(domain, ty, protocol)?;
220        self.accept_with(Self::from_socket2(accept_sock)?).await
221    }
222
223    #[cfg(windows)]
224    pub async fn accept_with(&self, accept_sock: Self) -> io::Result<(Self, SockAddr)> {
225        let op = Accept::new(self.to_shared_fd(), accept_sock);
226        let (_, op) = buf_try!(@try compio_runtime::submit(op).await);
227        op.update_context()?;
228        let (accept_sock, addr) = op.into_addr()?;
229        Ok((accept_sock, addr))
230    }
231
232    pub fn incoming(&self) -> Incoming<'_> {
233        Incoming::new(self)
234    }
235
236    pub fn close(self) -> impl Future<Output = io::Result<()>> {
237        // Make sure that self won't be dropped after `close` called.
238        // Users may call this method and drop the future immediately. In that
239        // way the `close` should be cancelled.
240        let this = ManuallyDrop::new(self);
241        async move {
242            let fd = ManuallyDrop::into_inner(this)
243                .socket
244                .into_inner()
245                .take()
246                .await;
247            if let Some(fd) = fd {
248                let op = CloseSocket::new(fd.into_inner().into());
249                compio_runtime::submit(op).await.0?;
250            }
251            Ok(())
252        }
253    }
254
255    #[cfg(unix)]
256    async fn shutdown_impl(&self) -> io::Result<()> {
257        let op = ShutdownSocket::new(self.to_shared_fd(), std::net::Shutdown::Write);
258        compio_runtime::submit(op).await.0.map(|_| ())
259    }
260
261    #[cfg(windows)]
262    async fn shutdown_impl(&self) -> io::Result<()> {
263        self.socket.shutdown(std::net::Shutdown::Write)?;
264        Ok(())
265    }
266
267    pub async fn shutdown(&self) -> io::Result<()> {
268        match self.shutdown_impl().await {
269            // The socket is not connected, so we can ignore this error.
270            Err(e) if e.kind() == io::ErrorKind::NotConnected => Ok(()),
271            res => res,
272        }
273    }
274
275    pub async fn recv<B: IoBufMut>(&self, buffer: B, flags: RecvFlags) -> BufResult<usize, B> {
276        let fd = self.to_shared_fd();
277        let op = Recv::new(fd, buffer, flags);
278        let (res, extra) = compio_runtime::submit(op).with_extra().await;
279        extra.update_socket_state(&*self.socket);
280        let res = res.into_inner();
281        unsafe { res.map_advanced() }
282    }
283
284    pub async fn recv_vectored<V: IoVectoredBufMut>(
285        &self,
286        buffer: V,
287        flags: RecvFlags,
288    ) -> BufResult<usize, V> {
289        let fd = self.to_shared_fd();
290        let op = RecvVectored::new(fd, buffer, flags);
291        let (res, extra) = compio_runtime::submit(op).with_extra().await;
292        extra.update_socket_state(&*self.socket);
293        let res = res.into_inner();
294        unsafe { res.map_vec_advanced() }
295    }
296
297    pub async fn recv_managed(
298        &self,
299        len: usize,
300        flags: RecvFlags,
301    ) -> io::Result<Option<BufferRef>> {
302        let fd = self.to_shared_fd();
303        let (res, extra) = Runtime::with_current(|rt| {
304            let buffer_pool = rt.buffer_pool()?;
305            let op = RecvManaged::new(fd, &buffer_pool, len, flags)?;
306            io::Result::Ok(rt.submit(op).with_extra())
307        })?
308        .await;
309
310        extra.update_socket_state(&*self.socket);
311
312        unsafe { res.take_buffer() }
313    }
314
315    pub fn recv_multi(
316        &self,
317        len: usize,
318        flags: RecvFlags,
319    ) -> impl Stream<Item = io::Result<BufferRef>> {
320        let fd = self.to_shared_fd();
321        let rt = Runtime::current();
322        SubmitMultiStream::new(move || {
323            let buffer_pool = rt.buffer_pool()?;
324            let op = RecvMulti::new(fd.clone(), &buffer_pool, len, flags)?;
325            Ok(rt.submit_multi(op).into_managed(buffer_pool))
326        })
327    }
328
329    pub async fn send<T: IoBuf>(&self, buffer: T, flags: SendFlags) -> BufResult<usize, T> {
330        let fd = self.to_shared_fd();
331        let op = Send::new(fd, buffer, flags);
332        compio_runtime::submit(op).await.into_inner()
333    }
334
335    pub async fn send_vectored<T: IoVectoredBuf>(
336        &self,
337        buffer: T,
338        flags: SendFlags,
339    ) -> BufResult<usize, T> {
340        let fd = self.to_shared_fd();
341        let op = SendVectored::new(fd, buffer, flags);
342        compio_runtime::submit(op).await.into_inner()
343    }
344
345    pub async fn send_zerocopy<T: IoBuf>(
346        &self,
347        buf: T,
348        flags: SendFlags,
349    ) -> BufResult<usize, Zerocopy<SendZc<T, SharedFd<SocketWithState>>>> {
350        submit_zerocopy(SendZc::new(self.to_shared_fd(), buf, flags)).await
351    }
352
353    pub async fn send_zerocopy_vectored<T: IoVectoredBuf>(
354        &self,
355        buf: T,
356        flags: SendFlags,
357    ) -> BufResult<usize, Zerocopy<SendVectoredZc<T, SharedFd<SocketWithState>>>> {
358        submit_zerocopy(SendVectoredZc::new(self.to_shared_fd(), buf, flags)).await
359    }
360
361    pub async fn recv_from<T: IoBufMut>(
362        &self,
363        buffer: T,
364        flags: RecvFlags,
365    ) -> BufResult<(usize, Option<SockAddr>), T> {
366        let fd = self.to_shared_fd();
367        let op = RecvFrom::new(fd, buffer, flags);
368        let (res, extra) = compio_runtime::submit(op).with_extra().await;
369        extra.update_socket_state(&*self.socket);
370        let res = res.into_inner().map_addr();
371        unsafe { res.map_advanced() }
372    }
373
374    pub async fn recv_from_vectored<T: IoVectoredBufMut>(
375        &self,
376        buffer: T,
377        flags: RecvFlags,
378    ) -> BufResult<(usize, Option<SockAddr>), T> {
379        let fd = self.to_shared_fd();
380        let op = RecvFromVectored::new(fd, buffer, flags);
381        let (res, extra) = compio_runtime::submit(op).with_extra().await;
382        extra.update_socket_state(&*self.socket);
383        let res = res.into_inner().map_addr();
384        unsafe { res.map_vec_advanced() }
385    }
386
387    pub async fn recv_from_managed(
388        &self,
389        len: usize,
390        flags: RecvFlags,
391    ) -> io::Result<Option<(BufferRef, Option<SockAddr>)>> {
392        let fd = self.to_shared_fd();
393        let (inner, extra) = Runtime::with_current(|rt| {
394            let buffer_pool = rt.buffer_pool()?;
395            let op = RecvFromManaged::new(fd, &buffer_pool, len, flags)?;
396            io::Result::Ok(rt.submit(op).with_extra())
397        })?
398        .await;
399        extra.update_socket_state(&*self.socket);
400        let (len, op) = buf_try!(@try inner);
401        // Kernel returns 0 for the operation, drop the buffer and return
402        // Ok(None)
403        if len == 0 {
404            return Ok(None);
405        }
406        let Some((mut buf, addr)) = op.take_buffer() else {
407            return Err(io::Error::new(
408                io::ErrorKind::UnexpectedEof,
409                format!("Read {len} bytes, but no buffer was selected by kernel"),
410            ));
411        };
412        unsafe { buf.advance_to(len) };
413        Ok(Some((buf, addr)))
414    }
415
416    pub fn recv_from_multi(
417        &self,
418        flags: RecvFlags,
419    ) -> impl Stream<Item = io::Result<RecvFromMultiResult>> {
420        let fd = self.to_shared_fd();
421        let rt = Runtime::current();
422        SubmitMultiStream::new(move || {
423            let buffer_pool = rt.buffer_pool()?;
424            let op = RecvFromMulti::new(fd.clone(), &buffer_pool, flags)?;
425            Ok(rt.submit_multi(op).into_managed(buffer_pool))
426        })
427    }
428
429    pub async fn recv_msg<T: IoBufMut, C: IoBufMut>(
430        &self,
431        buffer: T,
432        control: C,
433        flags: RecvFlags,
434    ) -> BufResult<(usize, usize, Option<SockAddr>, ReturnFlags), (T, C)> {
435        self.recv_msg_vectored([buffer], control, flags)
436            .await
437            .map_buffer(|([buffer], control)| (buffer, control))
438    }
439
440    pub async fn recv_msg_vectored<T: IoVectoredBufMut, C: IoBufMut>(
441        &self,
442        buffer: T,
443        control: C,
444        flags: RecvFlags,
445    ) -> BufResult<(usize, usize, Option<SockAddr>, ReturnFlags), (T, C)> {
446        let fd = self.to_shared_fd();
447        let op = RecvMsg::new(fd, buffer, control, flags);
448        let (res, extra) = compio_runtime::submit(op).with_extra().await;
449        extra.update_socket_state(&*self.socket);
450        let res = res.into_inner().map_addr();
451        unsafe { res.map_vec_advanced() }
452    }
453
454    pub async fn recv_msg_managed<C: IoBufMut>(
455        &self,
456        len: usize,
457        control: C,
458        flags: RecvFlags,
459    ) -> io::Result<Option<(BufferRef, C, Option<SockAddr>, ReturnFlags)>> {
460        let fd = self.to_shared_fd();
461        let (inner, extra) = Runtime::with_current(|rt| {
462            let buffer_pool = rt.buffer_pool()?;
463            let op = RecvMsgManaged::new(fd, &buffer_pool, len, control, flags)?;
464            io::Result::Ok(rt.submit(op).with_extra())
465        })?
466        .await;
467        extra.update_socket_state(&*self.socket);
468        let (len, op) = buf_try!(@try inner);
469        // Kernel returns 0 for the operation, drop the buffer and return
470        // Ok(None)
471        if len == 0 {
472            return Ok(None);
473        }
474        let Some(((mut buf, mut control), addr, control_len, flags)) = op.take_buffer() else {
475            return Err(io::Error::new(
476                io::ErrorKind::UnexpectedEof,
477                format!("Read {len} bytes, but no buffer was selected by kernel"),
478            ));
479        };
480        unsafe { buf.advance_to(len) };
481        unsafe { control.advance_to(control_len) };
482        Ok(Some((buf, control, addr, flags)))
483    }
484
485    pub fn recv_msg_multi(
486        &self,
487        control_len: usize,
488        flags: RecvFlags,
489    ) -> impl Stream<Item = io::Result<RecvMsgMultiResult>> {
490        let fd = self.to_shared_fd();
491        let rt = Runtime::current();
492        SubmitMultiStream::new(move || {
493            let buffer_pool = rt.buffer_pool()?;
494            let op = RecvMsgMulti::new(fd.clone(), &buffer_pool, control_len, flags)?;
495            io::Result::Ok(
496                rt.submit_multi(op)
497                    .into_managed_with(buffer_pool, control_len),
498            )
499        })
500    }
501
502    pub async fn send_to<T: IoBuf>(
503        &self,
504        buffer: T,
505        addr: &SockAddr,
506        flags: SendFlags,
507    ) -> BufResult<usize, T> {
508        let fd = self.to_shared_fd();
509        let op = SendTo::new(fd, buffer, addr.clone(), flags);
510        compio_runtime::submit(op).await.into_inner()
511    }
512
513    pub async fn send_to_vectored<T: IoVectoredBuf>(
514        &self,
515        buffer: T,
516        addr: &SockAddr,
517        flags: SendFlags,
518    ) -> BufResult<usize, T> {
519        let fd = self.to_shared_fd();
520        let op = SendToVectored::new(fd, buffer, addr.clone(), flags);
521        compio_runtime::submit(op).await.into_inner()
522    }
523
524    pub async fn send_to_zerocopy<T: IoBuf>(
525        &self,
526        buffer: T,
527        addr: &SockAddr,
528        flags: SendFlags,
529    ) -> BufResult<usize, Zerocopy<SendToZc<T, SharedFd<SocketWithState>>>> {
530        let op = SendToZc::new(self.to_shared_fd(), buffer, addr.clone(), flags);
531        submit_zerocopy(op).await
532    }
533
534    pub async fn send_to_zerocopy_vectored<T: IoVectoredBuf>(
535        &self,
536        buffer: T,
537        addr: &SockAddr,
538        flags: SendFlags,
539    ) -> BufResult<usize, Zerocopy<SendToVectoredZc<T, SharedFd<SocketWithState>>>> {
540        let op = SendToVectoredZc::new(self.to_shared_fd(), buffer, addr.clone(), flags);
541        submit_zerocopy(op).await
542    }
543
544    pub async fn send_msg<T: IoBuf, C: IoBuf>(
545        &self,
546        buffer: T,
547        control: C,
548        addr: Option<&SockAddr>,
549        flags: SendFlags,
550    ) -> BufResult<usize, (T, C)> {
551        self.send_msg_vectored([buffer], control, addr, flags)
552            .await
553            .map_buffer(|([buffer], control)| (buffer, control))
554    }
555
556    pub async fn send_msg_vectored<T: IoVectoredBuf, C: IoBuf>(
557        &self,
558        buffer: T,
559        control: C,
560        addr: Option<&SockAddr>,
561        flags: SendFlags,
562    ) -> BufResult<usize, (T, C)> {
563        let fd = self.to_shared_fd();
564        let op = SendMsg::new(fd, buffer, control, addr.cloned(), flags);
565        compio_runtime::submit(op).await.into_inner()
566    }
567
568    pub async fn send_msg_zerocopy<T: IoBuf, C: IoBuf>(
569        &self,
570        buffer: T,
571        control: C,
572        addr: Option<&SockAddr>,
573        flags: SendFlags,
574    ) -> BufResult<usize, Extract<Zerocopy<SendMsgZc<[T; 1], C, SharedFd<SocketWithState>>>, T, C>>
575    {
576        let BufResult(res, fut) = self
577            .send_msg_zerocopy_vectored([buffer], control, addr, flags)
578            .await;
579
580        BufResult(res, Extract::new(fut))
581    }
582
583    pub async fn send_msg_zerocopy_vectored<T: IoVectoredBuf, C: IoBuf>(
584        &self,
585        buffer: T,
586        control: C,
587        addr: Option<&SockAddr>,
588        flags: SendFlags,
589    ) -> BufResult<usize, Zerocopy<SendMsgZc<T, C, SharedFd<SocketWithState>>>> {
590        let fd = self.to_shared_fd();
591        let op = SendMsgZc::new(fd, buffer, control, addr.cloned(), flags);
592        submit_zerocopy(op).await
593    }
594
595    #[cfg(unix)]
596    pub unsafe fn get_socket_option<T: Copy>(&self, level: i32, name: i32) -> io::Result<T> {
597        let mut value: MaybeUninit<T> = MaybeUninit::uninit();
598        let mut len = size_of::<T>() as libc::socklen_t;
599        syscall!(libc::getsockopt(
600            self.socket.as_raw_fd(),
601            level,
602            name,
603            value.as_mut_ptr() as _,
604            &mut len
605        ))
606        .map(|_| {
607            debug_assert_eq!(len as usize, size_of::<T>());
608            // SAFETY: The value is initialized by `getsockopt`.
609            unsafe { value.assume_init() }
610        })
611    }
612
613    #[cfg(windows)]
614    pub unsafe fn get_socket_option<T: Copy>(&self, level: i32, name: i32) -> io::Result<T> {
615        let mut value: MaybeUninit<T> = MaybeUninit::uninit();
616        let mut len = size_of::<T>() as i32;
617        syscall!(
618            SOCKET,
619            windows_sys::Win32::Networking::WinSock::getsockopt(
620                self.socket.as_raw_fd() as _,
621                level,
622                name,
623                value.as_mut_ptr() as _,
624                &mut len
625            )
626        )
627        .map(|_| {
628            debug_assert_eq!(len as usize, size_of::<T>());
629            // SAFETY: The value is initialized by `getsockopt`.
630            unsafe { value.assume_init() }
631        })
632    }
633
634    #[cfg(unix)]
635    pub unsafe fn set_socket_option<T: Copy>(
636        &self,
637        level: i32,
638        name: i32,
639        value: &T,
640    ) -> io::Result<()> {
641        syscall!(libc::setsockopt(
642            self.socket.as_raw_fd(),
643            level,
644            name,
645            value as *const _ as _,
646            std::mem::size_of::<T>() as _
647        ))
648        .map(|_| ())
649    }
650
651    #[cfg(windows)]
652    pub unsafe fn set_socket_option<T: Copy>(
653        &self,
654        level: i32,
655        name: i32,
656        value: &T,
657    ) -> io::Result<()> {
658        syscall!(
659            SOCKET,
660            windows_sys::Win32::Networking::WinSock::setsockopt(
661                self.socket.as_raw_fd() as _,
662                level,
663                name,
664                value as *const _ as _,
665                std::mem::size_of::<T>() as _
666            )
667        )
668        .map(|_| ())
669    }
670}
671
672impl AsRawFd for Socket {
673    fn as_raw_fd(&self) -> RawFd {
674        self.socket.as_raw_fd()
675    }
676}
677
678impl AsFd for Socket {
679    fn as_fd(&self) -> BorrowedFd<'_> {
680        self.socket.as_fd()
681    }
682}
683
684#[cfg(unix)]
685impl std::os::fd::FromRawFd for SocketWithState {
686    unsafe fn from_raw_fd(fd: RawFd) -> Self {
687        Self::new(unsafe { std::os::fd::FromRawFd::from_raw_fd(fd) })
688    }
689}
690
691#[cfg(windows)]
692impl std::os::windows::io::FromRawSocket for SocketWithState {
693    unsafe fn from_raw_socket(sock: std::os::windows::io::RawSocket) -> Self {
694        Self::new(unsafe { std::os::windows::io::FromRawSocket::from_raw_socket(sock) })
695    }
696}
697
698#[cfg(unix)]
699impl std::os::fd::FromRawFd for Socket {
700    unsafe fn from_raw_fd(fd: RawFd) -> Self {
701        Self {
702            socket: unsafe { std::os::fd::FromRawFd::from_raw_fd(fd) },
703        }
704    }
705}
706
707impl compio_driver::ToSharedFd<SocketWithState> for Socket {
708    fn to_shared_fd(&self) -> compio_driver::SharedFd<SocketWithState> {
709        self.socket.to_shared_fd()
710    }
711}
712
713#[cfg(windows)]
714impl std::os::windows::io::FromRawSocket for Socket {
715    unsafe fn from_raw_socket(sock: std::os::windows::io::RawSocket) -> Self {
716        Self {
717            socket: unsafe { std::os::windows::io::FromRawSocket::from_raw_socket(sock) },
718        }
719    }
720}
721
722#[cfg(windows)]
723impl std::os::windows::io::AsSocket for Socket {
724    fn as_socket(&self) -> std::os::windows::io::BorrowedSocket<'_> {
725        self.socket.as_socket()
726    }
727}
728
729#[cfg(windows)]
730impl std::os::windows::io::AsRawSocket for Socket {
731    fn as_raw_socket(&self) -> std::os::windows::io::RawSocket {
732        self.socket.as_raw_socket()
733    }
734}
735
736pub struct Zerocopy<T: OpCode> {
737    stream: Option<SubmitMulti<T>>,
738}
739
740impl<T: OpCode + IntoInner + 'static> Zerocopy<T> {
741    pub(crate) fn new(stream: SubmitMulti<T>) -> Self {
742        Self {
743            stream: Some(stream),
744        }
745    }
746}
747
748impl<T: OpCode + IntoInner + 'static> Future for Zerocopy<T> {
749    type Output = T::Inner;
750
751    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
752        let this = self.get_mut();
753
754        let stream = this.stream.as_mut().expect("stream should be Some");
755
756        let _ = futures_util::ready!(stream.poll_next_unpin(cx));
757
758        let stream = this.stream.take().expect("stream should be Some");
759
760        Poll::Ready(
761            stream
762                .try_take()
763                .map_err(|_| ())
764                .expect("Cannot retrieve buffer")
765                .into_inner(),
766        )
767    }
768}
769
770pin_project! {
771    pub struct Extract<F, T, C>
772    where
773        F: Future<Output = ([T; 1], C)>,
774        T: IoBuf,
775        C: IoBuf,
776    {
777        #[pin]
778        future: F,
779    }
780}
781
782impl<F, T, C> Extract<F, T, C>
783where
784    F: Future<Output = ([T; 1], C)>,
785    T: IoBuf,
786    C: IoBuf,
787{
788    pub(crate) fn new(future: F) -> Self {
789        Self { future }
790    }
791}
792
793impl<F, T, C> Future for Extract<F, T, C>
794where
795    F: Future<Output = ([T; 1], C)>,
796    T: IoBuf,
797    C: IoBuf,
798{
799    type Output = (T, C);
800
801    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
802        let this = self.project();
803
804        match this.future.poll(cx) {
805            Poll::Pending => Poll::Pending,
806            Poll::Ready(([buf], control)) => Poll::Ready((buf, control)),
807        }
808    }
809}
810
811async fn submit_zerocopy<T: OpCode + IntoInner + 'static>(op: T) -> BufResult<usize, Zerocopy<T>> {
812    let mut stream = compio_runtime::submit_multi(op);
813    let res = stream
814        .next()
815        .await
816        .expect("SubmitMulti should yield at least one item")
817        .0;
818
819    let fut = Zerocopy::new(stream);
820
821    BufResult(res, fut)
822}
compio_net/socket/mod.rs

compio_net/socket/
mod.rs
