1 //! TCP/UDP/Unix helpers for tokio. 2 3 use crate::either::Either; 4 use std::future::Future; 5 use std::io::Result; 6 use std::pin::Pin; 7 use std::task::{Context, Poll}; 8 9 #[cfg(unix)] 10 pub mod unix; 11 12 /// A trait for a listener: `TcpListener` and `UnixListener`. 13 pub 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. poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>>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. accept(&mut self) -> ListenerAcceptFut<'_, Self> where Self: Sized,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. local_addr(&self) -> Result<Self::Addr>31 fn local_addr(&self) -> Result<Self::Addr>; 32 } 33 34 impl Listener for tokio::net::TcpListener { 35 type Io = tokio::net::TcpStream; 36 type Addr = std::net::SocketAddr; 37 poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>>38 fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>> { 39 Self::poll_accept(self, cx) 40 } 41 local_addr(&self) -> Result<Self::Addr>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"] 50 pub struct ListenerAcceptFut<'a, L> { 51 listener: &'a mut L, 52 } 53 54 impl<'a, L> Future for ListenerAcceptFut<'a, L> 55 where 56 L: Listener, 57 { 58 type Output = Result<(L::Io, L::Addr)>; 59 poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>60 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 61 self.listener.poll_accept(cx) 62 } 63 } 64 65 impl<L, R> Either<L, R> 66 where 67 L: Listener, 68 R: Listener, 69 { 70 /// Accepts a new incoming connection from this listener. accept(&mut self) -> Result<Either<(L::Io, L::Addr), (R::Io, R::Addr)>>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. local_addr(&self) -> Result<Either<L::Addr, R::Addr>>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