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(&mut self, buf: &[u8]) -> io::Result<usize>`
17    fn send<'a, 'b>(&'a mut self, buf: &'b [u8]) -> PollSend<'a, 'b, Self> {
18        PollSend { socket: self, buf }
19    }
20
21    /// `async fn send(&mut self, buf: &[u8]) -> io::Result<usize>`
22    fn send_to<'a, 'b>(
23        &'a mut 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>(&mut self, buf: &mut [u8]) -> io::Result<()>`
35    fn recv<'a, 'b, B>(
36        &'a mut self,
37        buf: &'b mut B,
38    ) -> PollRecv<'a, 'b, Self, B>
39    where
40        B: bytes::BufMut,
41    {
42        PollRecv { socket: self, buf }
43    }
44
45    /// `async fn recv<B>(&mut self, buf: &mut [u8]) -> io::Result<SocketAddr>`
46    fn recv_from<'a, 'b, B>(
47        &'a mut self,
48        buf: &'b mut B,
49    ) -> PollRecvFrom<'a, 'b, Self, B>
50    where
51        B: bytes::BufMut,
52    {
53        PollRecvFrom { socket: self, buf }
54    }
55
56    /// `async fn recrecv_from_full(&mut self) -> io::Result<(Vec<u8>,
57    /// SocketAddr)>`
58    fn recv_from_full(&mut self) -> PollRecvFromFull<'_, Self> {
59        PollRecvFromFull { socket: self }
60    }
61}
62
63impl<S: AsyncSocket> AsyncSocketExt for S {}
64
65pub struct PollSend<'a, 'b, S> {
66    socket: &'a mut S,
67    buf: &'b [u8],
68}
69
70impl<S> Future for PollSend<'_, '_, S>
71where
72    S: AsyncSocket,
73{
74    type Output = io::Result<usize>;
75
76    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
77        let this: &mut Self = Pin::into_inner(self);
78        this.socket.poll_send(cx, this.buf)
79    }
80}
81
82pub struct PollSendTo<'a, 'b, S> {
83    socket: &'a mut S,
84    buf: &'b [u8],
85    addr: &'b SocketAddr,
86}
87
88impl<S> Future for PollSendTo<'_, '_, S>
89where
90    S: AsyncSocket,
91{
92    type Output = io::Result<usize>;
93
94    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
95        let this: &mut Self = Pin::into_inner(self);
96        this.socket.poll_send_to(cx, this.buf, this.addr)
97    }
98}
99
100pub struct PollRecv<'a, 'b, S, B> {
101    socket: &'a mut S,
102    buf: &'b mut B,
103}
104
105impl<S, B> Future for PollRecv<'_, '_, S, B>
106where
107    S: AsyncSocket,
108    B: bytes::BufMut,
109{
110    type Output = io::Result<()>;
111
112    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
113        let this: &mut Self = Pin::into_inner(self);
114        this.socket.poll_recv(cx, this.buf)
115    }
116}
117
118pub struct PollRecvFrom<'a, 'b, S, B> {
119    socket: &'a mut S,
120    buf: &'b mut B,
121}
122
123impl<S, B> Future for PollRecvFrom<'_, '_, S, B>
124where
125    S: AsyncSocket,
126    B: bytes::BufMut,
127{
128    type Output = io::Result<SocketAddr>;
129
130    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
131        let this: &mut Self = Pin::into_inner(self);
132        this.socket.poll_recv_from(cx, this.buf)
133    }
134}
135
136pub struct PollRecvFromFull<'a, S> {
137    socket: &'a mut S,
138}
139
140impl<S> Future for PollRecvFromFull<'_, S>
141where
142    S: AsyncSocket,
143{
144    type Output = io::Result<(Vec<u8>, SocketAddr)>;
145
146    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
147        let this: &mut Self = Pin::into_inner(self);
148        this.socket.poll_recv_from_full(cx)
149    }
150}