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 self, buf: &'b [u8]) -> PollSend<'a, 'b, Self> {
18 PollSend { socket: self, buf }
19 }
20
21 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 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 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 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}