1//! TCP/UDP/Unix helpers for tokio.
2
3use crate::either::Either;
4use std::future::Future;
5use std::io::Result;
6use std::pin::Pin;
7use std::task::{Context, Poll};
8
9#[cfg(unix)]
10pub mod unix;
11
12/// A trait for a listener: `TcpListener` and `UnixListener`.
13pub trait Listener {
14 /// The stream's type of this listener.
15 type Io: tokio::io::AsyncRead + tokio::io::AsyncWrite;
16 /// The socket address type of this listener.
17 type Addr;
18
19 /// Polls to accept a new incoming connection to this listener.
20 fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>>;
21
22 /// Accepts a new incoming connection from this listener.
23 fn accept(&mut self) -> ListenerAcceptFut<'_, Self>
24 where
25 Self: Sized,
26 {
27 ListenerAcceptFut { listener: self }
28 }
29
30 /// Returns the local address that this listener is bound to.
31 fn local_addr(&self) -> Result<Self::Addr>;
32}
33
34impl Listener for tokio::net::TcpListener {
35 type Io = tokio::net::TcpStream;
36 type Addr = std::net::SocketAddr;
37
38 fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>> {
39 Self::poll_accept(self, cx)
40 }
41
42 fn local_addr(&self) -> Result<Self::Addr> {
43 self.local_addr().map(Into::into)
44 }
45}
46
47/// Future for accepting a new connection from a listener.
48#[derive(Debug)]
49#[must_use = "futures do nothing unless you `.await` or poll them"]
50pub struct ListenerAcceptFut<'a, L> {
51 listener: &'a mut L,
52}
53
54impl<'a, L> Future for ListenerAcceptFut<'a, L>
55where
56 L: Listener,
57{
58 type Output = Result<(L::Io, L::Addr)>;
59
60 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
61 self.listener.poll_accept(cx)
62 }
63}
64
65impl<L, R> Either<L, R>
66where
67 L: Listener,
68 R: Listener,
69{
70 /// Accepts a new incoming connection from this listener.
71 pub async fn accept(&mut self) -> Result<Either<(L::Io, L::Addr), (R::Io, R::Addr)>> {
72 match self {
73 Either::Left(listener) => {
74 let (stream, addr) = listener.accept().await?;
75 Ok(Either::Left((stream, addr)))
76 }
77 Either::Right(listener) => {
78 let (stream, addr) = listener.accept().await?;
79 Ok(Either::Right((stream, addr)))
80 }
81 }
82 }
83
84 /// Returns the local address that this listener is bound to.
85 pub fn local_addr(&self) -> Result<Either<L::Addr, R::Addr>> {
86 match self {
87 Either::Left(listener) => {
88 let addr = listener.local_addr()?;
89 Ok(Either::Left(addr))
90 }
91 Either::Right(listener) => {
92 let addr = listener.local_addr()?;
93 Ok(Either::Right(addr))
94 }
95 }
96 }
97}
98