compio-tls 0.10.0-rc.2

use std::{fmt::Debug, io};

use compio_io::{AsyncRead, AsyncWrite, compat::AsyncStream, util::Splittable};

use crate::TlsStream;

#[derive(Clone)]
enum TlsConnectorInner {
    #[cfg(feature = "native-tls")]
    NativeTls(crate::native::TlsConnector),
    #[cfg(feature = "rustls")]
    Rustls(futures_rustls::TlsConnector),
    #[cfg(feature = "py-dynamic-openssl")]
    PyDynamicOpenSsl(crate::py_ossl::TlsConnector),
    #[cfg(not(any(
        feature = "native-tls",
        feature = "rustls",
        feature = "py-dynamic-openssl"
    )))]
    None(std::convert::Infallible),
}

impl Debug for TlsConnectorInner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            #[cfg(feature = "native-tls")]
            Self::NativeTls(_) => f.debug_tuple("NativeTls").finish(),
            #[cfg(feature = "rustls")]
            Self::Rustls(_) => f.debug_tuple("Rustls").finish(),
            #[cfg(feature = "py-dynamic-openssl")]
            Self::PyDynamicOpenSsl(_) => f.debug_tuple("PyDynamicOpenSsl").finish(),
            #[cfg(not(any(
                feature = "native-tls",
                feature = "rustls",
                feature = "py-dynamic-openssl"
            )))]
            Self::None(f) => match *f {},
        }
    }
}

/// A wrapper around a [`native_tls::TlsConnector`] or [`rustls::ClientConfig`],
/// providing an async `connect` method.
#[derive(Debug, Clone)]
pub struct TlsConnector(TlsConnectorInner);

#[cfg(feature = "native-tls")]
impl From<native_tls::TlsConnector> for TlsConnector {
    fn from(value: native_tls::TlsConnector) -> Self {
        Self(TlsConnectorInner::NativeTls(value.into()))
    }
}

#[cfg(feature = "rustls")]
impl From<std::sync::Arc<rustls::ClientConfig>> for TlsConnector {
    fn from(value: std::sync::Arc<rustls::ClientConfig>) -> Self {
        Self(TlsConnectorInner::Rustls(value.into()))
    }
}

#[cfg(feature = "py-dynamic-openssl")]
#[doc(hidden)]
impl From<compio_py_dynamic_openssl::SSLContext> for TlsConnector {
    fn from(value: compio_py_dynamic_openssl::SSLContext) -> Self {
        Self(TlsConnectorInner::PyDynamicOpenSsl(value.into()))
    }
}

impl TlsConnector {
    /// Connects the provided stream with this connector, assuming the provided
    /// domain.
    ///
    /// This function will internally call `TlsConnector::connect` to connect
    /// the stream and returns a future representing the resolution of the
    /// connection operation. The returned future will resolve to either
    /// `TlsStream<S>` or `Error` depending if it's successful or not.
    ///
    /// This is typically used for clients who have already established, for
    /// example, a TCP connection to a remote server. That stream is then
    /// provided here to perform the client half of a connection to a
    /// TLS-powered server.
    pub async fn connect<S: Splittable + 'static>(
        &self,
        domain: &str,
        stream: S,
    ) -> io::Result<TlsStream<S>>
    where
        S::ReadHalf: AsyncRead + Unpin,
        S::WriteHalf: AsyncWrite + Unpin,
    {
        self.connect_compat(domain, AsyncStream::new(stream)).await
    }

    /// Similar to `connect` but accepts an [`AsyncStream`] instead of a raw
    /// stream. Users are free to adjust the inner buffer sizes and limits.
    pub async fn connect_compat<S: Splittable + 'static>(
        &self,
        domain: &str,
        stream: AsyncStream<S>,
    ) -> io::Result<TlsStream<S>>
    where
        S::ReadHalf: AsyncRead + Unpin,
        S::WriteHalf: AsyncWrite + Unpin,
    {
        match &self.0 {
            #[cfg(feature = "native-tls")]
            TlsConnectorInner::NativeTls(c) => {
                let client = c.connect(domain, Box::pin(stream)).await?;
                Ok(TlsStream::from(client))
            }
            #[cfg(feature = "rustls")]
            TlsConnectorInner::Rustls(c) => {
                let client = c
                    .connect(
                        domain.to_string().try_into().map_err(io::Error::other)?,
                        Box::pin(stream),
                    )
                    .await?;
                Ok(TlsStream::from(client))
            }
            #[cfg(feature = "py-dynamic-openssl")]
            TlsConnectorInner::PyDynamicOpenSsl(c) => {
                let client = c.connect(domain, Box::pin(stream)).await?;
                Ok(TlsStream::from(client))
            }
            #[cfg(not(any(
                feature = "native-tls",
                feature = "rustls",
                feature = "py-dynamic-openssl"
            )))]
            TlsConnectorInner::None(f) => match *f {},
        }
    }
}

#[derive(Clone)]
enum TlsAcceptorInner {
    #[cfg(feature = "native-tls")]
    NativeTls(crate::native::TlsAcceptor),
    #[cfg(feature = "rustls")]
    Rustls(futures_rustls::TlsAcceptor),
    #[cfg(feature = "py-dynamic-openssl")]
    PyDynamicOpenSsl(crate::py_ossl::TlsAcceptor),
    #[cfg(not(any(
        feature = "native-tls",
        feature = "rustls",
        feature = "py-dynamic-openssl"
    )))]
    None(std::convert::Infallible),
}

/// A wrapper around a [`native_tls::TlsAcceptor`] or [`rustls::ServerConfig`],
/// providing an async `accept` method.
///
/// [`native_tls::TlsAcceptor`]: https://docs.rs/native-tls/latest/native_tls/struct.TlsAcceptor.html
/// [`rustls::ServerConfig`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConfig.html
#[derive(Clone)]
pub struct TlsAcceptor(TlsAcceptorInner);

#[cfg(feature = "native-tls")]
impl From<native_tls::TlsAcceptor> for TlsAcceptor {
    fn from(value: native_tls::TlsAcceptor) -> Self {
        Self(TlsAcceptorInner::NativeTls(value.into()))
    }
}

#[cfg(feature = "rustls")]
impl From<std::sync::Arc<rustls::ServerConfig>> for TlsAcceptor {
    fn from(value: std::sync::Arc<rustls::ServerConfig>) -> Self {
        Self(TlsAcceptorInner::Rustls(value.into()))
    }
}

#[cfg(feature = "py-dynamic-openssl")]
impl From<compio_py_dynamic_openssl::SSLContext> for TlsAcceptor {
    fn from(value: compio_py_dynamic_openssl::SSLContext) -> Self {
        Self(TlsAcceptorInner::PyDynamicOpenSsl(value.into()))
    }
}

impl TlsAcceptor {
    /// Accepts a new client connection with the provided stream.
    ///
    /// This function will internally call `TlsAcceptor::accept` to connect
    /// the stream and returns a future representing the resolution of the
    /// connection operation. The returned future will resolve to either
    /// `TlsStream<S>` or `Error` depending if it's successful or not.
    ///
    /// This is typically used after a new socket has been accepted from a
    /// `TcpListener`. That socket is then passed to this function to perform
    /// the server half of accepting a client connection.
    pub async fn accept<S: Splittable + 'static>(&self, stream: S) -> io::Result<TlsStream<S>>
    where
        S::ReadHalf: AsyncRead + Unpin,
        S::WriteHalf: AsyncWrite + Unpin,
    {
        self.accept_compat(AsyncStream::new(stream)).await
    }

    /// Similar to `accept` but accepts an [`AsyncStream`] instead of a raw
    /// stream. Users are free to adjust the inner buffer sizes and limits.
    pub async fn accept_compat<S: Splittable + 'static>(
        &self,
        stream: AsyncStream<S>,
    ) -> io::Result<TlsStream<S>>
    where
        S::ReadHalf: AsyncRead + Unpin,
        S::WriteHalf: AsyncWrite + Unpin,
    {
        match &self.0 {
            #[cfg(feature = "native-tls")]
            TlsAcceptorInner::NativeTls(c) => {
                let server = c.accept(Box::pin(stream)).await?;
                Ok(TlsStream::from(server))
            }
            #[cfg(feature = "rustls")]
            TlsAcceptorInner::Rustls(c) => {
                let server = c.accept(Box::pin(stream)).await?;
                Ok(TlsStream::from(server))
            }
            #[cfg(feature = "py-dynamic-openssl")]
            TlsAcceptorInner::PyDynamicOpenSsl(a) => {
                let server = a.accept(Box::pin(stream)).await?;
                Ok(TlsStream::from(server))
            }
            #[cfg(not(any(
                feature = "native-tls",
                feature = "rustls",
                feature = "py-dynamic-openssl"
            )))]
            TlsAcceptorInner::None(f) => match *f {},
        }
    }
}