tokio/tokio-util/src/net/mod.rs

//! TCP/UDP/Unix helpers for tokio.

use crate::either::Either;
use std::future::Future;
use std::io::Result;
use std::pin::Pin;
use std::task::{Context, Poll};

#[cfg(unix)]
pub mod unix;

/// A trait for a listener: `TcpListener` and `UnixListener`.
pub trait Listener {
    /// The stream's type of this listener.
    type Io: tokio::io::AsyncRead + tokio::io::AsyncWrite;
    /// The socket address type of this listener.
    type Addr;

    /// Polls to accept a new incoming connection to this listener.
    fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>>;

    /// Accepts a new incoming connection from this listener.
    fn accept(&mut self) -> ListenerAcceptFut<'_, Self>
    where
        Self: Sized,
    {
        ListenerAcceptFut { listener: self }
    }

    /// Returns the local address that this listener is bound to.
    fn local_addr(&self) -> Result<Self::Addr>;
}

impl Listener for tokio::net::TcpListener {
    type Io = tokio::net::TcpStream;
    type Addr = std::net::SocketAddr;

    fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>> {
        Self::poll_accept(self, cx)
    }

    fn local_addr(&self) -> Result<Self::Addr> {
        self.local_addr().map(Into::into)
    }
}

/// Future for accepting a new connection from a listener.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct ListenerAcceptFut<'a, L> {
    listener: &'a mut L,
}

impl<'a, L> Future for ListenerAcceptFut<'a, L>
where
    L: Listener,
{
    type Output = Result<(L::Io, L::Addr)>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.listener.poll_accept(cx)
    }
}

impl<L, R> Either<L, R>
where
    L: Listener,
    R: Listener,
{
    /// Accepts a new incoming connection from this listener.
    pub async fn accept(&mut self) -> Result<Either<(L::Io, L::Addr), (R::Io, R::Addr)>> {
        match self {
            Either::Left(listener) => {
                let (stream, addr) = listener.accept().await?;
                Ok(Either::Left((stream, addr)))
            }
            Either::Right(listener) => {
                let (stream, addr) = listener.accept().await?;
                Ok(Either::Right((stream, addr)))
            }
        }
    }

    /// Returns the local address that this listener is bound to.
    pub fn local_addr(&self) -> Result<Either<L::Addr, R::Addr>> {
        match self {
            Either::Left(listener) => {
                let addr = listener.local_addr()?;
                Ok(Either::Left(addr))
            }
            Either::Right(listener) => {
                let addr = listener.local_addr()?;
                Ok(Either::Right(addr))
            }
        }
    }
}