iota_http/

lib.rs

// Copyright (c) Mysten Labs, Inc.
// Modifications Copyright (c) 2025 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

use std::{collections::HashMap, sync::Arc, time::Duration};

use connection_handler::{OnConnectionClose, sleep_or_pending};
pub use http;
use http::{Request, Response};
use hyper_util::service::TowerToHyperService;
use io::ServerIo;
use tokio::task::JoinSet;
use tokio_rustls::{TlsAcceptor, rustls};
use tower::{Service, ServiceBuilder, ServiceExt};
use tracing::trace;

use self::{body::BoxBody, connection_info::ActiveConnections};

pub mod body;
mod config;
mod connection_handler;
mod connection_info;
mod fuse;
mod io;
mod listener;

pub use config::Config;
pub use connection_info::{ConnectInfo, ConnectionId, ConnectionInfo, PeerCertificates};
pub use listener::{Listener, ListenerExt};

pub(crate) type BoxError = Box<dyn std::error::Error + Send + Sync>;
/// h2 alpn in plain format for rustls.
const ALPN_H2: &[u8] = b"h2";
/// h1 alpn in plain format for rustls.
const ALPN_H1: &[u8] = b"http/1.1";

#[derive(Default)]
pub struct Builder {
    config: Config,
    tls_config: Option<rustls::ServerConfig>,
}

impl Builder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn config(mut self, config: Config) -> Self {
        self.config = config;
        self
    }

    // Convenience method for configuring TLS with a single server cert
    //
    // Attempts to load PEM formatted files for the certificate chain and private
    // key material from the provided file system paths.
    pub fn tls_single_cert(
        self,
        cert_file: impl AsRef<std::path::Path>,
        private_key_file: impl AsRef<std::path::Path>,
    ) -> Result<Self, BoxError> {
        let tls_config =
            iota_tls::create_rustls_server_config_from_pem(cert_file, private_key_file)?;
        Ok(self.tls_config(tls_config))
    }

    pub fn tls_config(mut self, tls_config: rustls::ServerConfig) -> Self {
        self.tls_config = Some(tls_config);
        self
    }

    pub fn serve<A, S, ResponseBody>(
        self,
        addr: A,
        service: S,
    ) -> Result<ServerHandle<std::net::SocketAddr>, BoxError>
    where
        A: std::net::ToSocketAddrs,
        S: Service<
                Request<BoxBody>,
                Response = Response<ResponseBody>,
                Error: Into<BoxError>,
                Future: Send,
            > + Clone
            + Send
            + 'static,
        ResponseBody: http_body::Body<Data = bytes::Bytes, Error: Into<BoxError>> + Send + 'static,
    {
        let listener = listener::TcpListenerWithOptions::new(
            addr,
            self.config.tcp_nodelay,
            self.config.tcp_keepalive,
        )?;

        Self::serve_with_listener(self, listener, service)
    }

    fn serve_with_listener<L, S, ResponseBody>(
        self,
        listener: L,
        service: S,
    ) -> Result<ServerHandle<L::Addr>, BoxError>
    where
        L: Listener,
        S: Service<
                Request<BoxBody>,
                Response = Response<ResponseBody>,
                Error: Into<BoxError>,
                Future: Send,
            > + Clone
            + Send
            + 'static,
        ResponseBody: http_body::Body<Data = bytes::Bytes, Error: Into<BoxError>> + Send + 'static,
    {
        self.config.validate()?;

        let local_addr = listener.local_addr()?;
        let graceful_shutdown_token = tokio_util::sync::CancellationToken::new();
        let connections = ActiveConnections::default();

        let tls_config = self.tls_config.map(|mut tls| {
            tls.alpn_protocols.push(ALPN_H2.into());
            if self.config.accept_http1 {
                tls.alpn_protocols.push(ALPN_H1.into());
            }
            Arc::new(tls)
        });

        let (watch_sender, watch_receiver) = tokio::sync::watch::channel(());
        let server = Server {
            config: self.config,
            tls_config,
            listener,
            local_addr: local_addr.clone(),
            service: ServiceBuilder::new()
                .layer(tower::util::BoxCloneService::layer())
                .map_response(|response: Response<ResponseBody>| response.map(body::boxed))
                .map_err(Into::into)
                .service(service),
            pending_connections: JoinSet::new(),
            connection_handlers: JoinSet::new(),
            connections: connections.clone(),
            graceful_shutdown_token: graceful_shutdown_token.clone(),
            _watch_receiver: watch_receiver,
        };

        let handle = ServerHandle(Arc::new(HandleInner {
            local_addr,
            connections,
            graceful_shutdown_token,
            watch_sender,
        }));

        tokio::spawn(server.serve());

        Ok(handle)
    }
}

#[derive(Debug)]
pub struct ServerHandle<A = std::net::SocketAddr>(Arc<HandleInner<A>>);

#[derive(Debug)]
struct HandleInner<A = std::net::SocketAddr> {
    /// The local address of the server.
    local_addr: A,
    connections: ActiveConnections<A>,
    graceful_shutdown_token: tokio_util::sync::CancellationToken,
    watch_sender: tokio::sync::watch::Sender<()>,
}

impl<A> ServerHandle<A> {
    /// Returns the local address of the server
    pub fn local_addr(&self) -> &A {
        &self.0.local_addr
    }

    /// Trigger a graceful shutdown of the server, but don't wait till the
    /// server has completed shutting down
    pub fn trigger_shutdown(&self) {
        self.0.graceful_shutdown_token.cancel();
    }

    /// Completes once the network has been shutdown.
    ///
    /// This explicitly *does not* trigger the network to shutdown, see
    /// `trigger_shutdown` or `shutdown` if you want to trigger shutting
    /// down the server.
    pub async fn wait_for_shutdown(&self) {
        self.0.watch_sender.closed().await
    }

    /// Triggers a shutdown of the server and waits for it to complete shutting
    /// down.
    pub async fn shutdown(&self) {
        self.trigger_shutdown();
        self.wait_for_shutdown().await;
    }

    /// Checks if the Server has been shutdown.
    pub fn is_shutdown(&self) -> bool {
        self.0.watch_sender.is_closed()
    }

    pub fn connections(
        &self,
    ) -> std::sync::RwLockReadGuard<'_, HashMap<ConnectionId, ConnectionInfo<A>>> {
        self.0.connections.read().unwrap()
    }

    /// Returns the number of active connections the server is handling
    pub fn number_of_connections(&self) -> usize {
        self.connections().len()
    }
}

impl<A> Clone for ServerHandle<A> {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }
}

type ConnectingOutput<Io, Addr> = Result<(ServerIo<Io>, Addr), crate::BoxError>;

struct Server<L: Listener> {
    config: Config,
    tls_config: Option<Arc<rustls::ServerConfig>>,

    listener: L,
    local_addr: L::Addr,
    service: tower::util::BoxCloneService<Request<BoxBody>, Response<BoxBody>, crate::BoxError>,

    pending_connections: JoinSet<ConnectingOutput<L::Io, L::Addr>>,
    connection_handlers: JoinSet<()>,
    connections: ActiveConnections<L::Addr>,
    graceful_shutdown_token: tokio_util::sync::CancellationToken,
    // Used to signal to a ServerHandle when the server has completed shutting down
    _watch_receiver: tokio::sync::watch::Receiver<()>,
}

impl<L> Server<L>
where
    L: Listener,
{
    async fn serve(mut self) -> Result<(), BoxError> {
        loop {
            tokio::select! {
                _ = self.graceful_shutdown_token.cancelled() => {
                    trace!("signal received, shutting down");
                    break;
                },
                // While the handshake limit is reached, leave new connections in the kernel
                // backlog rather than accepting them into an unbounded set of pending tasks.
                (io, remote_addr) = self.listener.accept(), if self.accepts_more_connections() => {
                    self.handle_incoming(io, remote_addr);
                },
                // A failed task affects only its own connection, so the loop keeps serving
                // the others.
                Some(maybe_connection) = self.pending_connections.join_next() => {
                    let (io, remote_addr) = match maybe_connection {
                        Ok(Ok((io, remote_addr))) => {
                            (io, remote_addr)
                        }
                        Ok(Err(e)) => {
                            tracing::debug!(error = %e, "error accepting connection");
                            continue;
                        }
                        Err(e) => {
                            tracing::error!(error = %e, "connection handshake task failed");
                            continue;
                        }
                    };

                    trace!("connection accepted");
                    self.handle_connection(io, remote_addr);
                },
                Some(connection_handler_output) = self.connection_handlers.join_next() => {
                    if let Err(e) = connection_handler_output {
                        tracing::error!(error = %e, "connection task failed");
                    }
                },
            }
        }

        // Shutting down, wait for all connection handlers to finish
        self.shutdown().await;

        Ok(())
    }

    /// Whether another connection may be accepted, or the limit on concurrent
    /// TLS handshakes is currently reached.
    fn accepts_more_connections(&self) -> bool {
        self.config
            .max_pending_connections
            .is_none_or(|max| self.pending_connections.len() < max)
    }

    fn handle_incoming(&mut self, io: L::Io, remote_addr: L::Addr) {
        if let Some(tls) = self.tls_config.clone() {
            let tls_acceptor = TlsAcceptor::from(tls);
            let allow_insecure = self.config.allow_insecure;
            let handshake_timeout = self.config.handshake_timeout;
            self.pending_connections.spawn(async move {
                tokio::select! {
                    result = handshake(io, remote_addr, tls_acceptor, allow_insecure) => result,
                    // Dropping the handshake closes the connection, releasing its file descriptor.
                    _ = sleep_or_pending(handshake_timeout) => Err(std::io::Error::new(
                        std::io::ErrorKind::TimedOut,
                        "TLS handshake timed out",
                    )
                    .into()),
                }
            });
        } else {
            self.handle_connection(ServerIo::new_io(io), remote_addr);
        }
    }

    fn handle_connection(&mut self, io: ServerIo<L::Io>, remote_addr: L::Addr) {
        let connection_shutdown_token = self.graceful_shutdown_token.child_token();
        let connection_info = ConnectionInfo::new(
            remote_addr,
            io.peer_certs(),
            connection_shutdown_token.clone(),
        );
        let connection_id = connection_info.id();
        let connect_info = connection_info::ConnectInfo {
            local_addr: self.local_addr.clone(),
            remote_addr: connection_info.remote_address().clone(),
        };
        let peer_certificates = connection_info.peer_certificates().cloned();
        let hyper_io = hyper_util::rt::TokioIo::new(io);

        let hyper_svc = TowerToHyperService::new(self.service.clone().map_request(
            move |mut request: Request<hyper::body::Incoming>| {
                request.extensions_mut().insert(connect_info.clone());
                if let Some(peer_certificates) = peer_certificates.clone() {
                    request.extensions_mut().insert(peer_certificates);
                }

                request.map(body::boxed)
            },
        ));

        self.connections
            .write()
            .unwrap()
            .insert(connection_id, connection_info);
        let on_connection_close = OnConnectionClose::new(connection_id, self.connections.clone());

        self.connection_handlers
            .spawn(connection_handler::serve_connection(
                hyper_io,
                hyper_svc,
                self.config.connection_builder(),
                connection_shutdown_token,
                self.config.max_connection_age,
                on_connection_close,
            ));
    }

    async fn shutdown(mut self) {
        // The time we are willing to wait for a connection to get gracefully shutdown
        // before we attempt to forcefully shutdown all active connections
        const CONNECTION_SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(1);

        // Just to be careful make sure the token is canceled
        self.graceful_shutdown_token.cancel();

        // Terminate any in-progress pending connections
        self.pending_connections.shutdown().await;

        // Wait for all connection handlers to terminate
        trace!(
            "waiting for {} connections to close",
            self.connection_handlers.len()
        );

        let graceful_shutdown =
            async { while self.connection_handlers.join_next().await.is_some() {} };

        if tokio::time::timeout(CONNECTION_SHUTDOWN_GRACE_PERIOD, graceful_shutdown)
            .await
            .is_err()
        {
            tracing::warn!(
                "Failed to stop all connection handlers in {:?}. Forcing shutdown.",
                CONNECTION_SHUTDOWN_GRACE_PERIOD
            );
            self.connection_handlers.shutdown().await;
        }
    }
}

/// Runs the pre-authentication part of accepting a connection: sniffing
/// insecure traffic when it is allowed, then the TLS handshake.
async fn handshake<Io, Addr>(
    io: Io,
    remote_addr: Addr,
    tls_acceptor: TlsAcceptor,
    allow_insecure: bool,
) -> ConnectingOutput<Io, Addr>
where
    Io: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
    if allow_insecure {
        // XXX: If we want to allow for supporting insecure traffic from other types of
        // io, we'll need to implement a generic peekable IO type
        if let Some(tcp) = <dyn std::any::Any>::downcast_ref::<tokio::net::TcpStream>(&io) {
            // Determine whether new connection is TLS.
            let mut buf = [0; 1];
            // `peek` blocks until at least some data is available, so if there is no error
            // then it must return the one byte we are requesting.
            tcp.peek(&mut buf).await?;
            // First byte of a TLS handshake is 0x16, so if it isn't 0x16 then its
            // insecure
            if buf != [0x16] {
                tracing::trace!("accepting insecure connection");
                return Ok((ServerIo::new_io(io), remote_addr));
            }
        } else {
            tracing::warn!(
                "'allow_insecure' is configured but io type is not 'tokio::net::TcpStream'"
            );
        }
    }

    tracing::trace!("accepting TLS connection");
    let io = tls_acceptor.accept(io).await?;
    Ok((ServerIo::new_tls_io(io), remote_addr))
}

#[cfg(test)]
mod tests {
    use axum::Router;

    use super::*;

    #[tokio::test]
    async fn simple() {
        const MESSAGE: &str = "Hello, World!";

        let app = Router::new().route("/", axum::routing::get(|| async { MESSAGE }));

        let handle = Builder::new().serve(("localhost", 0), app).unwrap();

        let url = format!("http://{}", handle.local_addr());

        let response = reqwest::get(url).await.unwrap().bytes().await.unwrap();

        assert_eq!(response, MESSAGE.as_bytes());
    }

    #[tokio::test]
    async fn shutdown() {
        const MESSAGE: &str = "Hello, World!";

        let app = Router::new().route("/", axum::routing::get(|| async { MESSAGE }));

        let handle = Builder::new().serve(("localhost", 0), app).unwrap();

        let url = format!("http://{}", handle.local_addr());

        let response = reqwest::get(url).await.unwrap().bytes().await.unwrap();

        // a request was just made so we should have 1 active connection
        assert_eq!(handle.connections().len(), 1);

        assert_eq!(response, MESSAGE.as_bytes());

        assert!(!handle.is_shutdown());

        handle.shutdown().await;

        assert!(handle.is_shutdown());

        // Now that the network has been shutdown there should be zero connections
        assert_eq!(handle.connections().len(), 0);
    }

    const SERVER_NAME: &str = "iota-http-test";

    /// A server config and a client config that trusts it, without client
    /// authentication.
    fn test_tls_configs() -> (rustls::ServerConfig, rustls::ClientConfig) {
        use fastcrypto::{
            ed25519::{Ed25519KeyPair, Ed25519PrivateKey},
            traits::{KeyPair, ToFromBytes},
        };

        let keypair = Ed25519KeyPair::from(Ed25519PrivateKey::from_bytes(&[42; 32]).unwrap());
        let public_key = keypair.public().to_owned();
        (
            iota_tls::create_rustls_server_config(keypair.private(), SERVER_NAME.to_string()),
            iota_tls::create_rustls_client_config(public_key, SERVER_NAME.to_string(), None),
        )
    }

    /// The peer is unauthenticated until its handshake completes, so a peer
    /// that never starts one must not hold the connection open indefinitely.
    #[tokio::test]
    async fn silent_peer_is_dropped_after_the_handshake_timeout() {
        use tokio::io::AsyncReadExt as _;

        const HANDSHAKE_TIMEOUT: Duration = Duration::from_millis(200);

        let (server_tls_config, _) = test_tls_configs();
        let handle = Builder::new()
            .config(Config::default().handshake_timeout(Some(HANDSHAKE_TIMEOUT)))
            .tls_config(server_tls_config)
            .serve(("localhost", 0), Router::new())
            .unwrap();

        // Connect, then never send a ClientHello.
        let mut connection = tokio::net::TcpStream::connect(handle.local_addr())
            .await
            .unwrap();

        let mut buf = [0u8; 1];
        let read = tokio::time::timeout(HANDSHAKE_TIMEOUT * 25, connection.read(&mut buf))
            .await
            .expect("the server must not wait for the handshake past the deadline");

        assert!(
            matches!(read, Ok(0) | Err(_)),
            "the server must close the connection, got {read:?}"
        );
    }

    /// While the pending limit is reached the server stops accepting, and
    /// resumes once a slot frees.
    #[tokio::test]
    async fn pending_connections_are_capped() {
        let (server_tls_config, client_tls_config) = test_tls_configs();
        let handle = Builder::new()
            .config(
                Config::default()
                    // Long enough that the stalled peer below is released by the test
                    // rather than by the deadline.
                    .handshake_timeout(Some(Duration::from_secs(60)))
                    .max_pending_connections(Some(1)),
            )
            .tls_config(server_tls_config)
            .serve(("localhost", 0), Router::new())
            .unwrap();
        let addr = *handle.local_addr();

        // Occupy the only handshake slot with a peer that never speaks.
        let stalled = tokio::net::TcpStream::connect(addr).await.unwrap();
        tokio::time::sleep(Duration::from_millis(200)).await;

        let connector = tokio_rustls::TlsConnector::from(Arc::new(client_tls_config));
        let server_name = rustls::pki_types::ServerName::try_from(SERVER_NAME).unwrap();
        let handshake = async {
            let io = tokio::net::TcpStream::connect(addr).await.unwrap();
            connector.connect(server_name, io).await
        };
        tokio::pin!(handshake);

        // The connection reaches the kernel backlog but is not accepted. Without the
        // limit the handshake completes in milliseconds, so a wide window here only
        // guards against a slow machine, it does not weaken the assertion.
        assert!(
            tokio::time::timeout(Duration::from_secs(2), &mut handshake)
                .await
                .is_err(),
            "the server must not handshake while the pending limit is reached"
        );

        drop(stalled);

        tokio::time::timeout(Duration::from_secs(10), &mut handshake)
            .await
            .expect("the server must resume accepting once a slot frees")
            .expect("the handshake must succeed");
    }

    /// A limit the accept loop can never fall below, and a limit whose slots
    /// nothing releases, both leave the server unable to accept.
    #[tokio::test]
    async fn unrecoverable_limits_are_rejected() {
        let served = |config| {
            Builder::new()
                .config(config)
                .tls_config(test_tls_configs().0)
                .serve(("localhost", 0), Router::new())
        };

        assert!(
            served(Config::default().max_pending_connections(Some(0))).is_err(),
            "a zero limit must be rejected"
        );
        assert!(
            served(
                Config::default()
                    .handshake_timeout(None)
                    .max_pending_connections(Some(8))
            )
            .is_err(),
            "a limit without a handshake deadline must be rejected"
        );
        assert!(
            served(
                Config::default()
                    .handshake_timeout(None)
                    .max_pending_connections(None)
            )
            .is_ok(),
            "removing both bounds stays allowed"
        );
    }
}