1use std::{
4 future::Future,
5 io,
6 pin::Pin,
7 task::{Context, Poll},
8};
9
10use crate::{AsyncSocket, SocketAddr};
11
12pub trait AsyncSocketExt: AsyncSocket {
16 fn send<'a, 'b>(&'a mut self, buf: &'b [u8]) -> PollSend<'a, 'b, Self> {
18 PollSend { socket: self, buf }
19 }
20
21 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 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 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 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}