netlink_sys/
async_socket_ext.rs

1// SPDX-License-Identifier: MIT
2
3use std::{
4    future::Future,
5    io,
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10use crate::{AsyncSocket, SocketAddr};
11
12/// Support trait for [`AsyncSocket`]
13///
14/// Provides awaitable variants of the poll functions from [`AsyncSocket`].
15pub trait AsyncSocketExt: AsyncSocket {
16    /// `async fn send(&self, buf: &[u8]) -> io::Result<usize>`
17    fn send<'a, 'b>(&'a self, buf: &'b [u8]) -> PollSend<'a, 'b, Self> {
18        PollSend { socket: self, buf }
19    }
20
21    /// `async fn send(&self, buf: &[u8]) -> io::Result<usize>`
22    fn send_to<'a, 'b>(
23        &'a self,
24        buf: &'b [u8],
25        addr: &'b SocketAddr,
26    ) -> PollSendTo<'a, 'b, Self> {
27        PollSendTo {
28            socket: self,
29            buf,
30            addr,
31        }
32    }
33
34    /// `async fn recv<B>(&self, buf: &mut [u8]) -> io::Result<()>`
35    fn recv<'a, 'b, B>(&'a self, buf: &'b mut B) -> PollRecv<'a, 'b, Self, B>
36    where
37        B: bytes::BufMut,
38    {
39        PollRecv { socket: self, buf }
40    }
41
42    /// `async fn recv<B>(&self, buf: &mut [u8]) -> io::Result<SocketAddr>`
43    fn recv_from<'a, 'b, B>(
44        &'a self,
45        buf: &'b mut B,
46    ) -> PollRecvFrom<'a, 'b, Self, B>
47    where
48        B: bytes::BufMut,
49    {
50        PollRecvFrom { socket: self, buf }
51    }
52
53    /// `async fn recrecv_from_full(&self) -> io::Result<(Vec<u8>,
54    /// SocketAddr)>`
55    fn recv_from_full(&self) -> PollRecvFromFull<'_, Self> {
56        PollRecvFromFull { socket: self }
57    }
58}
59
60impl<S: AsyncSocket> AsyncSocketExt for S {}
61
62pub struct PollSend<'a, 'b, S> {
63    socket: &'a S,
64    buf: &'b [u8],
65}
66
67impl<S> Future for PollSend<'_, '_, S>
68where
69    S: AsyncSocket,
70{
71    type Output = io::Result<usize>;
72
73    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
74        let this: &mut Self = Pin::into_inner(self);
75        this.socket.poll_send(cx, this.buf)
76    }
77}
78
79pub struct PollSendTo<'a, 'b, S> {
80    socket: &'a S,
81    buf: &'b [u8],
82    addr: &'b SocketAddr,
83}
84
85impl<S> Future for PollSendTo<'_, '_, S>
86where
87    S: AsyncSocket,
88{
89    type Output = io::Result<usize>;
90
91    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
92        let this: &mut Self = Pin::into_inner(self);
93        this.socket.poll_send_to(cx, this.buf, this.addr)
94    }
95}
96
97pub struct PollRecv<'a, 'b, S, B> {
98    socket: &'a S,
99    buf: &'b mut B,
100}
101
102impl<S, B> Future for PollRecv<'_, '_, S, B>
103where
104    S: AsyncSocket,
105    B: bytes::BufMut,
106{
107    type Output = io::Result<()>;
108
109    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
110        let this: &mut Self = Pin::into_inner(self);
111        this.socket.poll_recv(cx, this.buf)
112    }
113}
114
115pub struct PollRecvFrom<'a, 'b, S, B> {
116    socket: &'a S,
117    buf: &'b mut B,
118}
119
120impl<S, B> Future for PollRecvFrom<'_, '_, S, B>
121where
122    S: AsyncSocket,
123    B: bytes::BufMut,
124{
125    type Output = io::Result<SocketAddr>;
126
127    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
128        let this: &mut Self = Pin::into_inner(self);
129        this.socket.poll_recv_from(cx, this.buf)
130    }
131}
132
133pub struct PollRecvFromFull<'a, S> {
134    socket: &'a S,
135}
136
137impl<S> Future for PollRecvFromFull<'_, S>
138where
139    S: AsyncSocket,
140{
141    type Output = io::Result<(Vec<u8>, SocketAddr)>;
142
143    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
144        let this: &mut Self = Pin::into_inner(self);
145        this.socket.poll_recv_from_full(cx)
146    }
147}