defter 0.3.8

//! 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, or, for a hostname no CDN is in front
//! of, one from Let's Encrypt.

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;
use futures::StreamExt;
use rustls_acme::AcmeConfig;
use rustls_acme::caches::DirCache;

/// Where the HTTPS certificate comes from.
#[derive(Debug, Clone)]
pub enum Tls {
    /// A PEM certificate (chain) and its private key.
    Files { cert: PathBuf, key: PathBuf },
    /// Let's Encrypt, which is asked for one and for renewals through the
    /// TLS-ALPN-01 challenge. That needs only connections for `domain` on port
    /// 443 to reach this server, as they do through dragoman's `--sni-route`.
    Acme {
        domain: String,
        email: String,
        /// Keeps the account and certificates between restarts.
        cache: PathBuf,
        /// Let's Encrypt's staging environment, whose certificates browsers
        /// do not trust.
        staging: bool,
    },
}

/// 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::Acme {
            domain,
            email,
            cache,
            staging,
        }) => {
            std::fs::create_dir_all(&cache)
                .map_err(|e| format!("creating the certificate cache {}: {e}", cache.display()))?;
            let mut acme = AcmeConfig::new([domain.as_str()])
                .contact([format!("mailto:{email}").as_str()])
                .cache(DirCache::new(cache))
                .directory_lets_encrypt(!staging)
                .state();
            let acceptor = acme.axum_acceptor(acme.default_rustls_config());
            // Orders and renews for as long as the server runs.
            tokio::spawn(async move {
                while let Some(event) = acme.next().await {
                    match event {
                        Ok(ok) => tracing::info!(event = ?ok, "ACME"),
                        Err(error) => tracing::error!(%error, "ACME"),
                    }
                }
            });
            axum_server::bind(addr)
                .acceptor(acceptor)
                .handle(handle)
                .serve(service)
                .await
                .map_err(|e| format!("serving HTTPS on {addr}: {e}"))
        }
        Some(Tls::Files { cert, key }) => {
            let config = RustlsConfig::from_pem_file(&cert, &key)
                .await
                .map_err(|e| {
                    format!(
                        "loading the certificate {} and key {}: {e}",
                        cert.display(),
                        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 (cert_path, key_path) = (dir.path().join("cert.pem"), dir.path().join("key.pem"));
        let tls = Tls::Files {
            cert: cert_path.clone(),
            key: key_path.clone(),
        };
        std::fs::write(&cert_path, cert.cert.pem()).unwrap();
        std::fs::write(&key_path, 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::Files {
            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}");
    }
}