axol 0.2.0

Axol Web Framework
use std::{net::SocketAddr, sync::Arc};

use async_trait::async_trait;
use log::warn;
use rustls::{ServerConfig, server::Acceptor as RustlsAcceptor};
use tokio::net::TcpStream;
use tokio::sync::watch;
use tokio_rustls::{LazyConfigAcceptor, server::TlsStream};

use super::{Acceptor, LocalAddr, TcpAcceptor};

/// Accepts TLS connections, resolving the [`ServerConfig`] per connection.
///
/// The config lives behind a [`watch`] channel so certificates can be rotated at runtime;
/// connections arriving while the channel holds `None` are dropped.
pub struct TlsIncoming {
    tcp: TcpAcceptor,
    tls_config: watch::Receiver<Option<Arc<ServerConfig>>>,
}

impl TlsIncoming {
    pub async fn new(
        listen: SocketAddr,
        nodelay: bool,
        tls_config: watch::Receiver<Option<Arc<ServerConfig>>>,
    ) -> std::io::Result<Self> {
        let mut tcp = TcpAcceptor::bind(listen).await?;
        tcp.set_nodelay(nodelay);
        Ok(Self { tcp, tls_config })
    }

    pub async fn new_static(
        listen: SocketAddr,
        nodelay: bool,
        tls_config: ServerConfig,
    ) -> std::io::Result<Self> {
        Self::new(
            listen,
            nodelay,
            watch::channel(Some(Arc::new(tls_config))).1,
        )
        .await
    }

    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
        self.tcp.local_addr()
    }
}

impl LocalAddr for TlsIncoming {
    fn local_addr(&self) -> std::io::Result<SocketAddr> {
        TlsIncoming::local_addr(self)
    }
}

#[async_trait]
impl Acceptor for TlsIncoming {
    type Conn = TlsStream<TcpStream>;

    async fn accept(&mut self) -> std::io::Result<(Self::Conn, SocketAddr)> {
        loop {
            let (client, address) = self.tcp.accept().await?;

            let Some(server_config) = self.tls_config.borrow().clone() else {
                warn!(
                    "inbound TLS connection dropped (no certificates loaded, but were configured)"
                );
                continue;
            };

            // The handshake is driven here rather than in a spawned task; the accept loop in
            // `Server::serve_custom` is the only caller and a slow handshake would otherwise
            // block later connections. Callers wanting concurrent handshakes should wrap this.
            let accepted = match LazyConfigAcceptor::new(RustlsAcceptor::default(), client).await {
                Ok(x) => x,
                Err(e) => {
                    warn!("error during TLS init from {address}: {e}");
                    continue;
                }
            };
            match accepted.into_stream(server_config).await {
                Ok(stream) => return Ok((stream, address)),
                Err(e) => {
                    warn!("error during TLS handshake from {address}: {e}");
                    continue;
                }
            }
        }
    }
}