tokio_util/net/
mod.rs

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}
tokio_util/net/mod.rs

tokio_util/net/
mod.rs
