src/tls/server.rs

   1 use crate::tls::TlsConnectionMetadata;
   2 use crate::Connection;
   3 use async_std::net::{SocketAddr, TcpListener, ToSocketAddrs};
   4 use async_std::pin::Pin;
   5 use async_std::task::{Context, Poll};
   6 use async_tls::TlsAcceptor;
   7 use futures::{Stream, StreamExt};
   8 use log::*;
   9
  10 #[allow(dead_code)]
  11 pub struct TlsServer {
  12     local_addrs: SocketAddr,
  13     listener: TcpListener,
  14     acceptor: TlsAcceptor,
  15 }
  16
  17 impl TlsServer {
  18     pub async fn new<A: ToSocketAddrs + std::fmt::Display>(
  19         ip_addrs: A,
  20         acceptor: TlsAcceptor,
  21     ) -> anyhow::Result<Self> {
  22         let listener = TcpListener::bind(ip_addrs).await?;
  23         info!("Started TLS server at {}", listener.local_addr()?);
  24
  25         Ok(Self {
  26             local_addrs: listener.local_addr()?,
  27             listener,
  28             acceptor,
  29         })
  30     }
  31
  32     pub async fn accept(&self) -> anyhow::Result<Connection> {
  33         let (tcp_stream, peer_addr) = self.listener.accept().await?;
  34         debug!("Received connection attempt from {}", peer_addr);
  35
  36         match self.acceptor.accept(tcp_stream).await {
  37             Ok(tls_stream) => {
  38                 debug!("Completed TLS handshake with {}", peer_addr);
  39                 Ok(Connection::from(TlsConnectionMetadata::Server {
  40                     local_addr: self.local_addrs.clone(),
  41                     peer_addr,
  42                     stream: tls_stream,
  43                 }))
  44             }
  45
  46             Err(e) => {
  47                 warn!("Could not encrypt connection with TLS from {}", peer_addr);
  48                 Err(anyhow::Error::new(e))
  49             }
  50         }
  51     }
  52 }
  53
  54 impl Stream for TlsServer {
  55     type Item = Connection;
  56
  57     fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
  58         match futures::executor::block_on(self.listener.incoming().next()) {
  59             Some(Ok(tcp_stream)) => {
  60                 let peer_addr = tcp_stream
  61                     .peer_addr()
  62                     .expect("Could not retrieve peer IP address");
  63                 debug!("Received connection attempt from {}", peer_addr);
  64
  65                 match futures::executor::block_on(self.acceptor.accept(tcp_stream)) {
  66                     Ok(tls_stream) => {
  67                         debug!("Completed TLS handshake with {}", peer_addr);
  68                         Poll::Ready(Some(Connection::from(TlsConnectionMetadata::Server {
  69                             local_addr: self.local_addrs.clone(),
  70                             peer_addr,
  71                             stream: tls_stream,
  72                         })))
  73                     }
  74
  75                     Err(_e) => {
  76                         warn!("Could not encrypt connection with TLS from {}", peer_addr);
  77                         Poll::Pending
  78                     }
  79                 }
  80             }
  81
  82             Some(Err(e)) => {
  83                 error!(
  84                     "Encountered error when trying to accept new connection {}",
  85                     e
  86                 );
  87                 Poll::Pending
  88             }
  89
  90             None => Poll::Ready(None),
  91         }
  92     }
  93 }