General Upgrade, NEW: WebServers, Better Docs

mtp-webserver/src/server.rs

@@ -0,0 +1,577 @@
+use crate::{
+    Http3Request, Http3Response, Router, RouterError, WebMTPConnection, WebServerError,
+    transport::accept_web_connection,
+};
+use bytes::{Buf, Bytes};
+use http::{Method, Request, Response, StatusCode};
+use mtp_common::CommunicationError;
+use mtp_host::HostConfig;
+use rustls::pki_types::{PrivateKeyDer, pem::PemObject};
+use std::{net::SocketAddr, sync::Arc, time::Duration};
+use tokio::sync::{Semaphore, watch};
+
+/// Observability hooks for the web server.
+///
+/// Implement this trait to receive metrics about connections, requests, and
+/// errors.  All methods have default no-op implementations so callers only
+/// need to override the hooks they care about.
+pub trait WebServerMetrics: Send + Sync {
+    fn connection_accepted(&self) {}
+    fn connection_closed(&self, _duration: Duration, _reason: &str) {}
+    fn request_started(&self, _path: &str) {}
+    fn request_completed(&self, _path: &str, _status: u16, _duration: Duration) {}
+    fn error_occurred(&self, _error: &WebServerError) {}
+}
+
+/// Configuration for the HTTP/3 server and MTP routing.
+///
+/// Use the builder methods to customise behaviour.  All fields have sensible
+/// defaults so `WebServerConfig::new()` gives a usable production-ready
+/// configuration.
+#[derive(Clone)]
+pub struct WebServerConfig {
+    pub(crate) router: Router,
+    pub(crate) mtp_path: String,
+    pub max_request_body: usize,
+    pub max_connections: usize,
+    pub request_timeout: Duration,
+    pub drain_timeout: Duration,
+    pub(crate) metrics: Option<Arc<dyn WebServerMetrics>>,
+}
+
+impl Default for WebServerConfig {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl WebServerConfig {
+    pub fn new() -> Self {
+        Self {
+            router: Router::new(),
+            mtp_path: "/".to_string(),
+            max_request_body: 4 * 1024 * 1024,
+            max_connections: 256,
+            request_timeout: Duration::from_secs(30),
+            drain_timeout: Duration::from_secs(10),
+            metrics: None,
+        }
+    }
+
+    pub fn route<F, Fut>(mut self, path: impl Into<String>, handler: F) -> Result<Self, RouterError>
+    where
+        F: Fn(Http3Request, Http3Response) -> Fut + Send + Sync + 'static,
+        Fut: std::future::Future<Output = Http3Response> + Send + 'static,
+    {
+        self.router = self.router.route(path, handler)?;
+        Ok(self)
+    }
+
+    pub fn route_method<F, Fut>(
+        mut self,
+        method: Method,
+        path: impl Into<String>,
+        handler: F,
+    ) -> Result<Self, RouterError>
+    where
+        F: Fn(Http3Request, Http3Response) -> Fut + Send + Sync + 'static,
+        Fut: std::future::Future<Output = Http3Response> + Send + 'static,
+    {
+        self.router = self.router.route_method(method, path, handler)?;
+        Ok(self)
+    }
+
+    pub fn fallback<F, Fut>(mut self, handler: F) -> Result<Self, RouterError>
+    where
+        F: Fn(Http3Request, Http3Response) -> Fut + Send + Sync + 'static,
+        Fut: std::future::Future<Output = Http3Response> + Send + 'static,
+    {
+        self.router = self.router.fallback(handler)?;
+        Ok(self)
+    }
+
+    pub fn mtp_path(mut self, path: impl Into<String>) -> Self {
+        self.mtp_path = path.into();
+        self
+    }
+
+    pub fn max_request_body(mut self, bytes: usize) -> Self {
+        self.max_request_body = bytes;
+        self
+    }
+
+    pub fn max_connections(mut self, max: usize) -> Self {
+        self.max_connections = max;
+        self
+    }
+
+    pub fn request_timeout(mut self, timeout: Duration) -> Self {
+        self.request_timeout = timeout;
+        self
+    }
+
+    pub fn drain_timeout(mut self, timeout: Duration) -> Self {
+        self.drain_timeout = timeout;
+        self
+    }
+
+    pub fn with_metrics(mut self, metrics: Arc<dyn WebServerMetrics>) -> Self {
+        self.metrics = Some(metrics);
+        self
+    }
+}
+
+/// An HTTP/3 server bound to MTP's configured address and certificate.
+///
+/// One task owns the Quinn endpoint and dispatches all HTTP/3 requests.  This
+/// is the required ownership model for adding WebTransport MTP sessions on the
+/// same UDP socket without competing endpoint accept loops.
+pub struct MTPWebServer {
+    endpoint: quinn::Endpoint,
+    mtp_incoming: tokio::sync::mpsc::Receiver<Result<WebMTPConnection, mtp_host::AcceptError>>,
+    shutdown_tx: watch::Sender<()>,
+    _driver: Option<tokio::task::JoinHandle<()>>,
+}
+
+impl MTPWebServer {
+    pub async fn new(
+        host_config: HostConfig,
+        web_config: WebServerConfig,
+    ) -> Result<Self, CommunicationError> {
+        #[cfg(feature = "crypto")]
+        if !matches!(
+            host_config.authentication_policy,
+            mtp_host::AuthenticationPolicy::Unauthenticated
+        ) {
+            return Err(CommunicationError::Other(
+                "web authentication is not supported yet; use Unauthenticated".into(),
+            ));
+        }
+        let endpoint = build_endpoint(&host_config)?;
+        let driver_endpoint = endpoint.clone();
+        let (mtp_tx, mtp_incoming) = tokio::sync::mpsc::channel(16);
+        let (shutdown_tx, shutdown_rx) = watch::channel(());
+        let connection_semaphore = Arc::new(Semaphore::new(web_config.max_connections));
+        let driver_config = DriverConfig {
+            router: web_config.router,
+            mtp_path: web_config.mtp_path,
+            max_request_body: web_config.max_request_body,
+            request_timeout: web_config.request_timeout,
+            drain_timeout: web_config.drain_timeout,
+            send_pongs: host_config.send_pongs,
+            policy: host_config.policy,
+            metrics: web_config.metrics,
+        };
+        let driver = tokio::spawn(run_driver(
+            driver_endpoint,
+            driver_config,
+            mtp_tx,
+            connection_semaphore,
+            shutdown_rx,
+        ));
+        Ok(Self {
+            endpoint,
+            mtp_incoming,
+            shutdown_tx,
+            _driver: Some(driver),
+        })
+    }
+
+    pub fn local_addr(&self) -> SocketAddr {
+        self.endpoint
+            .local_addr()
+            .expect("endpoint has a local address")
+    }
+
+    pub async fn accept(&mut self) -> Result<Option<WebMTPConnection>, mtp_host::AcceptError> {
+        match self.mtp_incoming.recv().await {
+            Some(result) => result.map(Some),
+            None => Ok(None),
+        }
+    }
+
+    /// Signal the server to shut down gracefully.
+    ///
+    /// Stops accepting new QUIC connections, waits for the driver task to
+    /// finish its accept loop, then allows a configurable drain period for
+    /// in-flight requests to complete before closing the endpoint.
+    pub async fn shutdown(mut self) {
+        let _ = self.shutdown_tx.send(());
+        if let Some(driver) = self._driver.take() {
+            let _ = driver.await;
+        }
+        self.endpoint
+            .close(quinn::VarInt::from_u32(0), b"mtp-webserver shutdown");
+    }
+
+    /// Stop accepting new QUIC connections and wait briefly for the driver to stop.
+    pub async fn close(mut self) {
+        self.endpoint
+            .close(quinn::VarInt::from_u32(0), b"mtp-webserver shutdown");
+        if let Some(driver) = self._driver.take() {
+            driver.abort();
+        }
+    }
+}
+
+impl Drop for MTPWebServer {
+    fn drop(&mut self) {
+        if let Some(driver) = self._driver.take() {
+            driver.abort();
+        }
+    }
+}
+
+struct DriverConfig {
+    router: Router,
+    mtp_path: String,
+    max_request_body: usize,
+    request_timeout: Duration,
+    drain_timeout: Duration,
+    send_pongs: bool,
+    policy: mtp_transport::Policy,
+    metrics: Option<Arc<dyn WebServerMetrics>>,
+}
+
+fn build_endpoint(config: &HostConfig) -> Result<quinn::Endpoint, CommunicationError> {
+    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
+    let certificates = rustls::pki_types::CertificateDer::pem_slice_iter(&config.tls_fullchain)
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|_| CommunicationError::CertificateLoadFailed)?;
+    let key = PrivateKeyDer::from_pem_slice(&config.tls_key)
+        .map_err(|_| CommunicationError::CertificateParseFailed)?;
+
+    let mut tls = rustls::ServerConfig::builder()
+        .with_no_client_auth()
+        .with_single_cert(certificates, key)
+        .map_err(|_| CommunicationError::CertificateLoadFailed)?;
+    tls.alpn_protocols = vec![b"h3".to_vec()];
+
+    let server = quinn::ServerConfig::with_crypto(Arc::new(
+        quinn::crypto::rustls::QuicServerConfig::try_from(tls)
+            .map_err(|error| CommunicationError::Other(error.to_string()))?,
+    ));
+    quinn::Endpoint::server(server, SocketAddr::new(config.ip, config.port))
+        .map_err(|error| CommunicationError::Other(error.to_string()))
+}
+
+async fn run_driver(
+    endpoint: quinn::Endpoint,
+    config: DriverConfig,
+    mtp_tx: tokio::sync::mpsc::Sender<Result<WebMTPConnection, mtp_host::AcceptError>>,
+    connection_semaphore: Arc<Semaphore>,
+    mut shutdown_rx: watch::Receiver<()>,
+) {
+    let DriverConfig {
+        router,
+        mtp_path,
+        max_request_body,
+        request_timeout,
+        drain_timeout,
+        send_pongs,
+        policy,
+        metrics,
+    } = config;
+    let mut connection_tasks = tokio::task::JoinSet::new();
+    loop {
+        tokio::select! {
+            biased;
+            _ = shutdown_rx.changed() => {
+                break;
+            }
+            incoming = endpoint.accept() => {
+                let Some(incoming) = incoming else {
+                    break;
+                };
+                let permit = match connection_semaphore.clone().acquire_owned().await {
+                    Ok(permit) => permit,
+                    Err(_) => break,
+                };
+                let router = router.clone();
+                let mtp_path = mtp_path.clone();
+                let mtp_tx = mtp_tx.clone();
+                let metrics = metrics.clone();
+                connection_tasks.spawn(async move {
+                    let _permit = permit;
+                    let connect_start = std::time::Instant::now();
+                    let connection = match incoming.await {
+                        Ok(connection) => connection,
+                        Err(error) => {
+                            tracing::debug!(%error, "QUIC connection failed during handshake");
+                            return;
+                        }
+                    };
+                    let mut builder = h3::server::builder();
+                    builder.enable_extended_connect(true);
+                    builder.enable_webtransport(true);
+                    builder.enable_datagram(true);
+                    builder.max_webtransport_sessions(16);
+                    let mut h3 = match builder
+                        .build(h3_quinn::Connection::new(connection.clone()))
+                        .await
+                    {
+                        Ok(connection) => connection,
+                        Err(error) => {
+                            tracing::debug!(%error, "HTTP/3 connection setup failed");
+                            return;
+                        }
+                    };
+
+                    if let Some(ref m) = metrics {
+                        m.connection_accepted();
+                    }
+
+                    let mut tasks = tokio::task::JoinSet::new();
+                    loop {
+                        let resolver = match h3.accept().await {
+                            Ok(Some(resolver)) => resolver,
+                            Ok(None) => break,
+                            Err(error) => {
+                                tracing::debug!(%error, "HTTP/3 request accept failed");
+                                break;
+                            }
+                        };
+                        let (request, mut stream) = match resolver.resolve_request().await {
+                            Ok(request) => request,
+                            Err(error) => {
+                                tracing::debug!(%error, "HTTP/3 request parse failed");
+                                continue;
+                            }
+                        };
+                        if request.method() == http::Method::CONNECT && request.uri().path() == mtp_path {
+                            if request.extensions().get::<h3::ext::Protocol>()
+                                != Some(&h3::ext::Protocol::WEB_TRANSPORT)
+                            {
+                                let _ = stream
+                                    .send_response(
+                                        Response::builder()
+                                            .status(StatusCode::METHOD_NOT_ALLOWED)
+                                            .body(())
+                                            .unwrap(),
+                                    )
+                                    .await;
+                                let _ = stream.finish().await;
+                                continue;
+                            }
+                            let session = match h3_webtransport::server::WebTransportSession::accept(
+                                request, stream, h3,
+                            )
+                            .await
+                            {
+                                Ok(session) => Arc::new(session),
+                                Err(error) => {
+                                    tracing::debug!(%error, "WebTransport session accept failed");
+                                    return;
+                                }
+                            };
+                            tasks.spawn(run_session_requests(
+                                session.clone(),
+                                router.clone(),
+                                max_request_body,
+                                request_timeout,
+                                metrics.clone(),
+                            ));
+                            let result =
+                                accept_web_connection(session, mtp_path, connection, send_pongs, policy)
+                                    .await;
+                            let _ = mtp_tx.send(result).await;
+                            return;
+                        }
+                        let router = router.clone();
+                        let metrics = metrics.clone();
+                        tasks.spawn(async move {
+                            let path = request.uri().path().to_string();
+                            if let Some(ref m) = metrics {
+                                m.request_started(&path);
+                            }
+                            let req_start = std::time::Instant::now();
+                            let (response, status) = match tokio::time::timeout(
+                                request_timeout,
+                                handle_http_request(request, &mut stream, &router, max_request_body),
+                            )
+                            .await
+                            {
+                                Ok(Ok((resp, status))) => (resp, status),
+                                Ok(Err(_)) => (
+                                    Http3Response::new(StatusCode::BAD_GATEWAY),
+                                    StatusCode::BAD_GATEWAY,
+                                ),
+                                Err(_) => (
+                                    Http3Response::new(StatusCode::REQUEST_TIMEOUT),
+                                    StatusCode::REQUEST_TIMEOUT,
+                                ),
+                            };
+                            let _ = write_response(&mut stream, response).await;
+                            if let Some(ref m) = metrics {
+                                m.request_completed(&path, status.as_u16(), req_start.elapsed());
+                            }
+                        });
+                    }
+                    tasks.join_all().await;
+                    if let Some(ref m) = metrics {
+                        m.connection_closed(connect_start.elapsed(), "normal");
+                    }
+                });
+            }
+        }
+    }
+
+    // --- Drain phase: wait for in-flight connections ---
+
+    endpoint.close(quinn::VarInt::from_u32(0), b"mtp-webserver shutdown");
+
+    let drain_start = std::time::Instant::now();
+    while !connection_tasks.is_empty() {
+        tokio::select! {
+            Some(result) = connection_tasks.join_next() => {
+                if let Err(e) = result {
+                    tracing::warn!("Connection task panicked: {}", e);
+                }
+            }
+            _ = tokio::time::sleep(drain_timeout.saturating_sub(drain_start.elapsed())) => {
+                tracing::warn!(
+                    "Drain timeout expired with {} connections still in flight",
+                    connection_tasks.len()
+                );
+                break;
+            }
+        }
+    }
+
+    connection_tasks.shutdown().await;
+}
+
+async fn handle_http_request<S>(
+    request: Request<()>,
+    stream: &mut h3::server::RequestStream<S, Bytes>,
+    router: &Router,
+    max_request_body: usize,
+) -> Result<(Http3Response, StatusCode), WebServerError>
+where
+    S: h3::quic::BidiStream<Bytes>,
+{
+    let (request, too_large) = read_request(request, stream, max_request_body)
+        .await
+        .map_err(|e| WebServerError::Http(format!("request body read failed: {e}")))?;
+    if too_large {
+        return Ok((
+            Http3Response::new(StatusCode::PAYLOAD_TOO_LARGE),
+            StatusCode::PAYLOAD_TOO_LARGE,
+        ));
+    }
+    let path = request.uri.path().to_string();
+    let handler = router
+        .handler(&request.method, &path)
+        .or_else(|| router.fallback_handler());
+    match handler {
+        Some(handler) => {
+            let response = handler(request, Http3Response::default()).await;
+            let status = response.status;
+            Ok((response, status))
+        }
+        None => Ok((
+            Http3Response::new(StatusCode::NOT_FOUND),
+            StatusCode::NOT_FOUND,
+        )),
+    }
+}
+
+async fn run_session_requests(
+    session: Arc<h3_webtransport::server::WebTransportSession<h3_quinn::Connection, Bytes>>,
+    router: Router,
+    max_request_body: usize,
+    request_timeout: Duration,
+    metrics: Option<Arc<dyn WebServerMetrics>>,
+) {
+    loop {
+        match session.accept_bi().await {
+            Ok(Some(h3_webtransport::server::AcceptedBi::Request(request, mut stream))) => {
+                let router = router.clone();
+                let metrics = metrics.clone();
+                tokio::spawn(async move {
+                    let path = request.uri().path().to_string();
+                    if let Some(ref m) = metrics {
+                        m.request_started(&path);
+                    }
+                    let req_start = std::time::Instant::now();
+                    let (response, status) = match tokio::time::timeout(
+                        request_timeout,
+                        handle_http_request(request, &mut stream, &router, max_request_body),
+                    )
+                    .await
+                    {
+                        Ok(Ok((resp, status))) => (resp, status),
+                        Ok(Err(_)) => (
+                            Http3Response::new(StatusCode::BAD_GATEWAY),
+                            StatusCode::BAD_GATEWAY,
+                        ),
+                        Err(_) => (
+                            Http3Response::new(StatusCode::REQUEST_TIMEOUT),
+                            StatusCode::REQUEST_TIMEOUT,
+                        ),
+                    };
+                    let _ = write_response(&mut stream, response).await;
+                    if let Some(ref m) = metrics {
+                        m.request_completed(&path, status.as_u16(), req_start.elapsed());
+                    }
+                });
+            }
+            Ok(Some(h3_webtransport::server::AcceptedBi::BidiStream(_, _))) => {}
+            Ok(None) | Err(_) => break,
+        }
+    }
+}
+
+async fn read_request<S>(
+    request: Request<()>,
+    stream: &mut h3::server::RequestStream<S, Bytes>,
+    max_body: usize,
+) -> Result<(Http3Request, bool), h3::error::StreamError>
+where
+    S: h3::quic::BidiStream<Bytes>,
+{
+    let (parts, _) = request.into_parts();
+    let mut body = Vec::new();
+    let mut too_large = false;
+    while let Some(chunk) = stream.recv_data().await? {
+        if body.len().saturating_add(chunk.remaining()) > max_body {
+            too_large = true;
+            break;
+        }
+        body.extend_from_slice(chunk.chunk());
+    }
+    Ok((
+        Http3Request {
+            method: parts.method,
+            uri: parts.uri,
+            headers: parts.headers,
+            body: (!body.is_empty()).then(|| Bytes::from(body)),
+        },
+        too_large,
+    ))
+}
+
+async fn write_response<S>(
+    stream: &mut h3::server::RequestStream<S, Bytes>,
+    response: Http3Response,
+) -> Result<(), h3::error::StreamError>
+where
+    S: h3::quic::BidiStream<Bytes>,
+{
+    let mut builder = Response::builder().status(response.status);
+    for (name, value) in &response.headers {
+        builder = builder.header(name, value);
+    }
+    stream
+        .send_response(builder.body(()).expect("valid HTTP response"))
+        .await?;
+    for chunk in response.body {
+        stream.send_data(chunk).await?;
+    }
+    if let Some(mut chunks) = response.stream {
+        while let Some(chunk) = chunks.recv().await {
+            stream.send_data(chunk).await?;
+        }
+    }
+    stream.finish().await
+}