kglite-bolt-server 0.21.3

//! The Bolt accept loop, owned here instead of delegated to `boltr`'s
//! `BoltServer::serve`.
//!
//! `boltr` 0.3 binds its own listener and never sets `TCP_NODELAY` on the
//! streams it accepts, and exposes no hook to do so. Bolt is request/response
//! with several small messages per transaction, so on Linux every small reply
//! waits in Nagle's buffer for the client's delayed ACK (~44 ms per exchange,
//! ~132 ms per explicit transaction; ~1 ms / ~2 ms with the option set).
//!
//! Everything a connection needs after `accept` is public in `boltr`, so this
//! module re-creates only the listener side: bind, session reaper, accept loop,
//! shutdown, optional TLS. Per-connection work is `boltr`'s own
//! `server_handshake` + `Connection::run`, run behind the pump in
//! `connection.rs`.

use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use crate::coalesce::CoalescingWriter;
use crate::connection::{serve_connection, ConnectionContext};
use crate::inflight::InflightCancel;
use boltr::error::BoltError;
use boltr::server::handshake::server_handshake;
use boltr::server::{AuthValidator, BoltBackend, SessionHandle, SessionManager};
use rustls_pki_types::pem::PemObject;
use rustls_pki_types::{CertificateDer, PrivateKeyDer};
use tokio::net::{TcpListener, TcpStream};
use tokio_rustls::TlsAcceptor;

/// What the accept loop needs; the counterpart of `boltr::server::BoltServer`'s
/// builder fields.
pub struct BoltListener<B: BoltBackend> {
    pub backend: B,
    pub auth: Option<Arc<dyn AuthValidator>>,
    pub tls: Option<TlsAcceptor>,
    pub idle_timeout: Option<Duration>,
    /// Accept to LOGON, and LOGOFF to LOGON; `None` waits forever.
    pub handshake_timeout: Option<Duration>,
    pub max_sessions: usize,
    pub max_message_size: usize,
    pub shutdown: Pin<Box<dyn Future<Output = ()> + Send>>,
}

fn invalid_tls<E: Into<Box<dyn std::error::Error + Send + Sync>>>(e: E) -> BoltError {
    BoltError::Io(std::io::Error::new(std::io::ErrorKind::InvalidData, e))
}

/// Build a TLS acceptor from PEM cert chain and key bytes.
///
/// Needs a process-wide rustls crypto provider to be installed first.
pub fn tls_acceptor_from_pem(cert_pem: &[u8], key_pem: &[u8]) -> Result<TlsAcceptor, BoltError> {
    let certs: Vec<CertificateDer<'static>> = CertificateDer::pem_slice_iter(cert_pem)
        .collect::<Result<_, _>>()
        .map_err(invalid_tls)?;
    let key = PrivateKeyDer::from_pem_slice(key_pem).map_err(invalid_tls)?;
    let config = tokio_rustls::rustls::ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(certs, key)
        .map_err(invalid_tls)?;
    Ok(TlsAcceptor::from(Arc::new(config)))
}

/// Per-socket setup, applied to the raw TCP stream before any TLS handshake so
/// plain and TLS connections both get it. A failure is logged, not fatal: the
/// connection still works, only slower.
fn tune_stream(stream: &TcpStream, peer_addr: SocketAddr) {
    if let Err(e) = stream.set_nodelay(true) {
        tracing::debug!(%peer_addr, error = %e, "could not set TCP_NODELAY");
    }
}

impl<B: BoltBackend + InflightCancel + 'static> BoltListener<B> {
    /// Bind `addr` and serve until the shutdown future resolves.
    ///
    /// Like `boltr`'s loop, in-flight connection tasks are detached and are
    /// not awaited at shutdown.
    pub async fn serve(self, addr: SocketAddr) -> Result<(), BoltError> {
        let listener = TcpListener::bind(addr).await?;
        let backend = Arc::new(self.backend);
        let sessions = Arc::new(SessionManager::new(Some(self.max_sessions)));
        let ctx = Arc::new(ConnectionContext {
            backend: backend.clone(),
            sessions: sessions.clone(),
            auth: self.auth.clone(),
            max_message_size: self.max_message_size,
            handshake_timeout: self.handshake_timeout,
        });

        let reaper = self.idle_timeout.map(|timeout| {
            let sessions = sessions.clone();
            let backend = backend.clone();
            tokio::spawn(async move {
                let mut interval = tokio::time::interval(timeout / 2);
                loop {
                    interval.tick().await;
                    for id in sessions.reap_idle(timeout) {
                        let _ = backend.close_session(&SessionHandle(id.clone())).await;
                        tracing::debug!(session_id = %id, "reaped idle Bolt session");
                    }
                }
            })
        });

        let tls_label = if self.tls.is_some() { " (TLS)" } else { "" };
        tracing::info!(%addr, "Bolt server listening{}", tls_label);

        let tls = self.tls.map(Arc::new);
        let mut shutdown = self.shutdown;
        loop {
            tokio::select! {
                accepted = listener.accept() => match accepted {
                    Ok((stream, peer_addr)) => {
                        tune_stream(&stream, peer_addr);
                        tokio::spawn(run_connection(stream, peer_addr, ctx.clone(), tls.clone()));
                    }
                    Err(e) => tracing::warn!(error = %e, "accept error"),
                },
                () = &mut shutdown => {
                    tracing::info!("Bolt server shutting down");
                    break;
                }
            }
        }

        if let Some(handle) = reaper {
            handle.abort();
        }
        Ok(())
    }
}

async fn run_connection<B: BoltBackend + InflightCancel + 'static>(
    stream: TcpStream,
    peer_addr: SocketAddr,
    ctx: Arc<ConnectionContext<B>>,
    tls: Option<Arc<TlsAcceptor>>,
) {
    let accepted_at = tokio::time::Instant::now();
    match tls {
        Some(acceptor) => {
            let accepted =
                within(ctx.handshake_timeout, accepted_at, acceptor.accept(stream)).await;
            match accepted {
                Some(Ok(stream)) => handshake_and_run(stream, peer_addr, &ctx, accepted_at).await,
                Some(Err(e)) => tracing::debug!(%peer_addr, error = %e, "TLS handshake failed"),
                None => tracing::debug!(%peer_addr, "TLS handshake timed out"),
            }
        }
        None => handshake_and_run(stream, peer_addr, &ctx, accepted_at).await,
    }
}

/// `future`'s output, or `None` when `timeout` after `started` passes first.
/// `None` for the timeout waits as long as the future takes.
async fn within<T>(
    timeout: Option<Duration>,
    started: tokio::time::Instant,
    future: impl Future<Output = T>,
) -> Option<T> {
    match timeout {
        Some(timeout) => tokio::time::timeout_at(started + timeout, future)
            .await
            .ok(),
        None => Some(future.await),
    }
}

async fn handshake_and_run<S, B>(
    mut stream: S,
    peer_addr: SocketAddr,
    ctx: &ConnectionContext<B>,
    accepted_at: tokio::time::Instant,
) where
    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
    B: BoltBackend + InflightCancel + 'static,
{
    match within(
        ctx.handshake_timeout,
        accepted_at,
        server_handshake(&mut stream),
    )
    .await
    {
        Some(Ok(version)) => {
            tracing::debug!(%peer_addr, ?version, "Bolt handshake complete");
            let (reader, writer) = tokio::io::split(stream);
            let writer = CoalescingWriter::new(writer);
            serve_connection(reader, writer, ctx, peer_addr, accepted_at).await;
        }
        Some(Err(e)) => tracing::debug!(%peer_addr, error = %e, "Bolt handshake failed"),
        None => tracing::debug!(%peer_addr, "Bolt handshake timed out"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A freshly accepted socket has Nagle on; `tune_stream` is what turns it
    /// off, for every connection the loop accepts (issue #201).
    #[tokio::test]
    async fn accepted_streams_get_tcp_nodelay() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let _client = TcpStream::connect(addr).await.unwrap();
        let (server_side, peer) = listener.accept().await.unwrap();
        assert!(
            !server_side.nodelay().unwrap(),
            "premise: Nagle is on by default"
        );
        tune_stream(&server_side, peer);
        assert!(server_side.nodelay().unwrap());
    }
}