tokio_util/net/
mod.rs

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