Move server types into server module

Author: djc

src/lib.rs

@@ -38,20 +38,17 @@
 
 #![warn(unreachable_pub)]
 
-use std::future::Future;
 use std::io;
 #[cfg(unix)]
 use std::os::unix::io::{AsRawFd, RawFd};
 #[cfg(windows)]
 use std::os::windows::io::{AsRawSocket, RawSocket};
 use std::pin::Pin;
-use std::sync::Arc;
 use std::task::{Context, Poll};
 
 pub use rustls;
 
-use rustls::server::AcceptedAlert;
-use rustls::{CommonState, ServerConfig, ServerConnection};
+use rustls::CommonState;
 use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
 
 macro_rules! ready {
@@ -66,284 +63,8 @@ macro_rules! ready {
 pub mod client;
 pub use client::{Connect, FallibleConnect, TlsConnector, TlsConnectorWithAlpn};
 mod common;
-use common::{MidHandshake, TlsState};
 pub mod server;
-
-/// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method.
-#[derive(Clone)]
-pub struct TlsAcceptor {
-    inner: Arc<ServerConfig>,
-}
-
-impl From<Arc<ServerConfig>> for TlsAcceptor {
-    fn from(inner: Arc<ServerConfig>) -> TlsAcceptor {
-        TlsAcceptor { inner }
-    }
-}
-
-impl TlsAcceptor {
-    #[inline]
-    pub fn accept<IO>(&self, stream: IO) -> Accept<IO>
-    where
-        IO: AsyncRead + AsyncWrite + Unpin,
-    {
-        self.accept_with(stream, |_| ())
-    }
-
-    pub fn accept_with<IO, F>(&self, stream: IO, f: F) -> Accept<IO>
-    where
-        IO: AsyncRead + AsyncWrite + Unpin,
-        F: FnOnce(&mut ServerConnection),
-    {
-        let mut session = match ServerConnection::new(self.inner.clone()) {
-            Ok(session) => session,
-            Err(error) => {
-                return Accept(MidHandshake::Error {
-                    io: stream,
-                    // TODO(eliza): should this really return an `io::Error`?
-                    // Probably not...
-                    error: io::Error::new(io::ErrorKind::Other, error),
-                });
-            }
-        };
-        f(&mut session);
-
-        Accept(MidHandshake::Handshaking(server::TlsStream {
-            session,
-            io: stream,
-            state: TlsState::Stream,
-            need_flush: false,
-        }))
-    }
-
-    /// Get a read-only reference to underlying config
-    pub fn config(&self) -> &Arc<ServerConfig> {
-        &self.inner
-    }
-}
-
-pub struct LazyConfigAcceptor<IO> {
-    acceptor: rustls::server::Acceptor,
-    io: Option<IO>,
-    alert: Option<(rustls::Error, AcceptedAlert)>,
-}
-
-impl<IO> LazyConfigAcceptor<IO>
-where
-    IO: AsyncRead + AsyncWrite + Unpin,
-{
-    #[inline]
-    pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self {
-        Self {
-            acceptor,
-            io: Some(io),
-            alert: None,
-        }
-    }
-
-    /// Takes back the client connection. Will return `None` if called more than once or if the
-    /// connection has been accepted.
-    ///
-    /// # Example
-    ///
-    /// ```no_run
-    /// # fn choose_server_config(
-    /// #     _: rustls::server::ClientHello,
-    /// # ) -> std::sync::Arc<rustls::ServerConfig> {
-    /// #     unimplemented!();
-    /// # }
-    /// # #[allow(unused_variables)]
-    /// # async fn listen() {
-    /// use tokio::io::AsyncWriteExt;
-    /// let listener = tokio::net::TcpListener::bind("127.0.0.1:4443").await.unwrap();
-    /// let (stream, _) = listener.accept().await.unwrap();
-    ///
-    /// let acceptor = tokio_rustls::LazyConfigAcceptor::new(rustls::server::Acceptor::default(), stream);
-    /// tokio::pin!(acceptor);
-    ///
-    /// match acceptor.as_mut().await {
-    ///     Ok(start) => {
-    ///         let clientHello = start.client_hello();
-    ///         let config = choose_server_config(clientHello);
-    ///         let stream = start.into_stream(config).await.unwrap();
-    ///         // Proceed with handling the ServerConnection...
-    ///     }
-    ///     Err(err) => {
-    ///         if let Some(mut stream) = acceptor.take_io() {
-    ///             stream
-    ///                 .write_all(
-    ///                     format!("HTTP/1.1 400 Invalid Input\r\n\r\n\r\n{:?}\n", err)
-    ///                         .as_bytes()
-    ///                 )
-    ///                 .await
-    ///                 .unwrap();
-    ///         }
-    ///     }
-    /// }
-    /// # }
-    /// ```
-    pub fn take_io(&mut self) -> Option<IO> {
-        self.io.take()
-    }
-}
-
-impl<IO> Future for LazyConfigAcceptor<IO>
-where
-    IO: AsyncRead + AsyncWrite + Unpin,
-{
-    type Output = Result<StartHandshake<IO>, io::Error>;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = self.get_mut();
-        loop {
-            let io = match this.io.as_mut() {
-                Some(io) => io,
-                None => {
-                    return Poll::Ready(Err(io::Error::new(
-                        io::ErrorKind::Other,
-                        "acceptor cannot be polled after acceptance",
-                    )))
-                }
-            };
-
-            if let Some((err, mut alert)) = this.alert.take() {
-                match alert.write(&mut common::SyncWriteAdapter { io, cx }) {
-                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
-                        this.alert = Some((err, alert));
-                        return Poll::Pending;
-                    }
-                    Ok(0) | Err(_) => {
-                        return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err)))
-                    }
-                    Ok(_) => {
-                        this.alert = Some((err, alert));
-                        continue;
-                    }
-                };
-            }
-
-            let mut reader = common::SyncReadAdapter { io, cx };
-            match this.acceptor.read_tls(&mut reader) {
-                Ok(0) => return Err(io::ErrorKind::UnexpectedEof.into()).into(),
-                Ok(_) => {}
-                Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Poll::Pending,
-                Err(e) => return Err(e).into(),
-            }
-
-            match this.acceptor.accept() {
-                Ok(Some(accepted)) => {
-                    let io = this.io.take().unwrap();
-                    return Poll::Ready(Ok(StartHandshake { accepted, io }));
-                }
-                Ok(None) => {}
-                Err((err, alert)) => {
-                    this.alert = Some((err, alert));
-                }
-            }
-        }
-    }
-}
-
-/// An incoming connection received through [`LazyConfigAcceptor`].
-///
-/// This contains the generic `IO` asynchronous transport,
-/// [`ClientHello`](rustls::server::ClientHello) data,
-/// and all the state required to continue the TLS handshake (e.g. via
-/// [`StartHandshake::into_stream`]).
-#[derive(Debug)]
-pub struct StartHandshake<IO> {
-    accepted: rustls::server::Accepted,
-    io: IO,
-}
-
-impl<IO> StartHandshake<IO>
-where
-    IO: AsyncRead + AsyncWrite + Unpin,
-{
-    pub fn client_hello(&self) -> rustls::server::ClientHello<'_> {
-        self.accepted.client_hello()
-    }
-
-    pub fn into_stream(self, config: Arc<ServerConfig>) -> Accept<IO> {
-        self.into_stream_with(config, |_| ())
-    }
-
-    pub fn into_stream_with<F>(self, config: Arc<ServerConfig>, f: F) -> Accept<IO>
-    where
-        F: FnOnce(&mut ServerConnection),
-    {
-        let mut conn = match self.accepted.into_connection(config) {
-            Ok(conn) => conn,
-            Err((error, alert)) => {
-                return Accept(MidHandshake::SendAlert {
-                    io: self.io,
-                    alert,
-                    // TODO(eliza): should this really return an `io::Error`?
-                    // Probably not...
-                    error: io::Error::new(io::ErrorKind::InvalidData, error),
-                });
-            }
-        };
-        f(&mut conn);
-
-        Accept(MidHandshake::Handshaking(server::TlsStream {
-            session: conn,
-            io: self.io,
-            state: TlsState::Stream,
-            need_flush: false,
-        }))
-    }
-}
-
-/// Future returned from `TlsAcceptor::accept` which will resolve
-/// once the accept handshake has finished.
-pub struct Accept<IO>(MidHandshake<server::TlsStream<IO>>);
-
-/// Like [Accept], but returns `IO` on failure.
-pub struct FallibleAccept<IO>(MidHandshake<server::TlsStream<IO>>);
-
-impl<IO> Accept<IO> {
-    #[inline]
-    pub fn into_fallible(self) -> FallibleAccept<IO> {
-        FallibleAccept(self.0)
-    }
-
-    pub fn get_ref(&self) -> Option<&IO> {
-        match &self.0 {
-            MidHandshake::Handshaking(sess) => Some(sess.get_ref().0),
-            MidHandshake::SendAlert { io, .. } => Some(io),
-            MidHandshake::Error { io, .. } => Some(io),
-            MidHandshake::End => None,
-        }
-    }
-
-    pub fn get_mut(&mut self) -> Option<&mut IO> {
-        match &mut self.0 {
-            MidHandshake::Handshaking(sess) => Some(sess.get_mut().0),
-            MidHandshake::SendAlert { io, .. } => Some(io),
-            MidHandshake::Error { io, .. } => Some(io),
-            MidHandshake::End => None,
-        }
-    }
-}
-
-impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Accept<IO> {
-    type Output = io::Result<server::TlsStream<IO>>;
-
-    #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
-    }
-}
-
-impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleAccept<IO> {
-    type Output = Result<server::TlsStream<IO>, (io::Error, IO)>;
-
-    #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        Pin::new(&mut self.0).poll(cx)
-    }
-}
+pub use server::{Accept, FallibleAccept, LazyConfigAcceptor, StartHandshake, TlsAcceptor};
 
 /// Unified TLS stream type
 ///