defter 0.3.3

//! Serve the router over plain HTTP or, with a certificate, over HTTPS.
//!
//! HTTPS lets defter sit behind something that passes TLS through rather than
//! terminating it — such as dragoman's `--sni-route`, which hands connections
//! for defter's hostname, still encrypted, to defter on a loopback port. The
//! certificate is then one the CDN in front trusts, e.g. a Cloudflare origin
//! certificate covering defter's hostname.

use std::future::Future;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::time::Duration;

use axum::Router;
use axum_server::Handle;
use axum_server::tls_rustls::RustlsConfig;

/// A PEM certificate (chain) and its private key.
#[derive(Debug, Clone)]
pub struct Tls {
    pub cert: PathBuf,
    pub key: PathBuf,
}

/// How long open connections get to finish when shutting down.
const DRAIN: Duration = Duration::from_secs(30);

/// Serve `app` on `addr` until `shutdown` completes. `ready` receives the
/// bound address (useful with port 0).
pub async fn serve(
    addr: SocketAddr,
    app: Router,
    tls: Option<Tls>,
    shutdown: impl Future<Output = ()> + Send + 'static,
    ready: impl FnOnce(SocketAddr) + Send + 'static,
) -> Result<(), String> {
    let service = app.into_make_service();
    let handle: Handle<SocketAddr> = Handle::new();
    let on_shutdown = handle.clone();
    tokio::spawn(async move {
        shutdown.await;
        on_shutdown.graceful_shutdown(Some(DRAIN));
    });
    let on_ready = handle.clone();
    tokio::spawn(async move {
        if let Some(addr) = on_ready.listening().await {
            ready(addr);
        }
    });

    match tls {
        None => axum_server::bind(addr)
            .handle(handle)
            .serve(service)
            .await
            .map_err(|e| format!("serving on {addr}: {e}")),
        Some(tls) => {
            let config = RustlsConfig::from_pem_file(&tls.cert, &tls.key)
                .await
                .map_err(|e| {
                    format!(
                        "loading the certificate {} and key {}: {e}",
                        tls.cert.display(),
                        tls.key.display()
                    )
                })?;
            axum_server::bind_rustls(addr, config)
                .handle(handle)
                .serve(service)
                .await
                .map_err(|e| format!("serving HTTPS on {addr}: {e}"))
        }
    }
}

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

    /// Start `app` on a free port; returns its address and a stopper.
    async fn start(
        app: Router,
        tls: Option<Tls>,
    ) -> (SocketAddr, tokio::sync::oneshot::Sender<()>) {
        let (stop, stopped) = tokio::sync::oneshot::channel::<()>();
        let (bound, address) = tokio::sync::oneshot::channel();
        tokio::spawn(async move {
            let addr = "127.0.0.1:0".parse().unwrap();
            serve(
                addr,
                app,
                tls,
                async {
                    let _ = stopped.await;
                },
                move |addr| {
                    let _ = bound.send(addr);
                },
            )
            .await
            .unwrap();
        });
        (address.await.unwrap(), stop)
    }

    fn app() -> Router {
        Router::new().route("/", get(|| async { "defter" }))
    }

    fn client() -> reqwest::ClientBuilder {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
        reqwest::Client::builder().no_proxy()
    }

    #[tokio::test]
    async fn serves_plain_http() {
        let (addr, _stop) = start(app(), None).await;
        let body = client()
            .build()
            .unwrap()
            .get(format!("http://{addr}/"))
            .send()
            .await
            .unwrap()
            .text()
            .await
            .unwrap();
        assert_eq!(body, "defter");
    }

    #[tokio::test]
    async fn serves_https_with_the_given_certificate() {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
        let dir = tempfile::tempdir().unwrap();
        let cert = rcgen::generate_simple_self_signed(vec!["defter.test".into()]).unwrap();
        let tls = Tls {
            cert: dir.path().join("cert.pem"),
            key: dir.path().join("key.pem"),
        };
        std::fs::write(&tls.cert, cert.cert.pem()).unwrap();
        std::fs::write(&tls.key, cert.signing_key.serialize_pem()).unwrap();

        let (addr, _stop) = start(app(), Some(tls)).await;
        let trusting = client()
            .tls_certs_only([reqwest::Certificate::from_pem(cert.cert.pem().as_bytes()).unwrap()])
            .resolve("defter.test", addr)
            .build()
            .unwrap();
        let response = trusting
            .get(format!("https://defter.test:{}/", addr.port()))
            .send()
            .await
            .unwrap();
        assert_eq!(response.version(), reqwest::Version::HTTP_2);
        assert_eq!(response.text().await.unwrap(), "defter");

        // A client that does not trust the certificate gets no answer.
        let untrusting = client().resolve("defter.test", addr).build().unwrap();
        assert!(
            untrusting
                .get(format!("https://defter.test:{}/", addr.port()))
                .send()
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn a_missing_certificate_is_an_error() {
        let tls = Tls {
            cert: "/nonexistent/cert.pem".into(),
            key: "/nonexistent/key.pem".into(),
        };
        let addr = "127.0.0.1:0".parse().unwrap();
        let error = serve(addr, app(), Some(tls), std::future::pending(), |_| {})
            .await
            .unwrap_err();
        assert!(error.contains("/nonexistent/cert.pem"), "{error}");
    }
}