ai-app/server/src/auth.rs

//! Bearer-token auth for the entire HTTP surface.
//!
//! This server's API *is* remote code execution, so the token gates every route
//! with zero unauthenticated endpoints -- the middleware is applied once around
//! the whole router (including the fallback) in `main.rs`, never per-route, so a
//! new route can't forget it. See PLAN.md's security section for the threat
//! model.
//!
//! Nothing in this module -- and nothing anywhere else -- may log the
//! Authorization header or the token; the test below is a tripwire against a
//! logging change silently starting to. It is one test covering both gating and
//! logging on purpose -- see the note in it.

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use axum::extract::{ConnectInfo, Request, State};
use axum::http::{StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use wg_app_link::enroll::{take_pending, token_hash_hex, token_matches};

use crate::config::TokenEntry;
use crate::session::SessionManager;

/// Applied to every rejection. Not against brute force -- infeasible at 256
/// bits -- but so a scanner probing the port shows up as a slow, loggable
/// drip rather than a fast one.
const REJECT_DELAY: Duration = Duration::from_millis(300);

pub async fn require_token(
    State(manager): State<Arc<SessionManager>>,
    request: Request,
    next: Next,
) -> Response {
    let presented = request
        .headers()
        .get(header::AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.strip_prefix("Bearer "));
    if let Some(token) = presented {
        let hashes: Vec<String> = manager
            .tokens()
            .into_iter()
            .map(|entry| entry.sha256)
            .collect();
        if token_matches(token, &hashes) {
            return next.run(request).await;
        }
        // A link minted by `--enroll-link` while this server was running:
        // the entry moves from the spool into the config here, on first
        // use, and is an ordinary token from then on.
        match take_pending(&manager.pending_enrollments_dir(), token) {
            Ok(Some(name)) => {
                let entry = TokenEntry {
                    name,
                    sha256: token_hash_hex(token),
                };
                match manager.add_token(entry) {
                    Ok(()) => {
                        tracing::info!("adopted an enrollment minted with --enroll-link");
                        return next.run(request).await;
                    }
                    Err(error) => tracing::error!("couldn't adopt the enrollment: {error:#}"),
                }
            }
            Ok(None) => {}
            Err(error) => tracing::warn!("couldn't read the enrollment spool: {error:#}"),
        }
    }

    // Peer address only -- never the header value. Absent when there is no
    // real socket (tests driving the router directly).
    let peer = request
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|ConnectInfo(addr)| addr.to_string())
        .unwrap_or_else(|| "unknown peer".to_string());
    tracing::warn!("rejected request from {peer}: missing or invalid bearer token");
    tokio::time::sleep(REJECT_DELAY).await;
    (StatusCode::UNAUTHORIZED, "missing or invalid bearer token").into_response()
}

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

    use axum::Router;
    use axum::body::Body;
    use axum::routing::get;
    use tower::ServiceExt;

    use wg_app_link::enroll::{generate_token, spool_pending};

    fn manager_with_token(dir: &std::path::Path, token: &str) -> Arc<SessionManager> {
        let manager = Arc::new(
            SessionManager::new(
                dir.join("config.ron"),
                dir.join("sessions"),
                dir.join("models"),
            )
            .expect("manager"),
        );
        manager
            .set_tokens(vec![TokenEntry {
                name: "phone".to_string(),
                sha256: token_hash_hex(token),
            }])
            .expect("set token");
        manager
    }

    fn guarded_router(manager: Arc<SessionManager>) -> Router {
        Router::new()
            .route("/probe", get(|| async { "ok" }))
            .fallback(|| async { StatusCode::NOT_FOUND })
            .layer(axum::middleware::from_fn_with_state(manager, require_token))
    }

    fn request(path: &str, auth: Option<&str>) -> Request {
        let mut builder = axum::http::Request::builder().uri(path);
        if let Some(auth) = auth {
            builder = builder.header(header::AUTHORIZATION, auth);
        }
        builder.body(Body::empty()).expect("request")
    }

    /// One test rather than separate gating and logging tests,
    /// deliberately: tracing caches callsite interest process-wide, so a
    /// test that hits the rejection path with no subscriber installed can
    /// poison the interest cache for the one that captures logs. Keeping
    /// every exercise of the middleware under the capturing subscriber
    /// makes the log assertions deterministic.
    #[tokio::test]
    async fn gates_every_route_and_never_logs_the_token() {
        #[derive(Clone, Default)]
        struct Capture(Arc<Mutex<Vec<u8>>>);
        impl std::io::Write for Capture {
            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
                self.0.lock().unwrap().extend_from_slice(buf);
                Ok(buf.len())
            }
            fn flush(&mut self) -> std::io::Result<()> {
                Ok(())
            }
        }
        impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for Capture {
            type Writer = Capture;
            fn make_writer(&'a self) -> Capture {
                self.clone()
            }
        }

        let capture = Capture::default();
        let subscriber = tracing_subscriber::fmt()
            .with_max_level(tracing::Level::TRACE)
            .with_writer(capture.clone())
            .finish();
        let _guard = tracing::subscriber::set_default(subscriber);

        let dir = tempfile::tempdir().expect("tempdir");
        let token = generate_token();
        let router = guarded_router(manager_with_token(dir.path(), &token));

        // No header, wrong token, wrong scheme: 401 everywhere, including
        // paths that don't exist -- a scanner learns nothing.
        for (path, auth) in [
            ("/probe", None),
            ("/probe", Some("Bearer wrong".to_string())),
            ("/probe", Some(format!("Basic {token}"))),
            ("/no-such-route", None),
        ] {
            let response = router
                .clone()
                .oneshot(request(path, auth.as_deref()))
                .await
                .expect("response");
            assert_eq!(
                response.status(),
                StatusCode::UNAUTHORIZED,
                "{path} {auth:?}"
            );
        }

        let ok = router
            .clone()
            .oneshot(request("/probe", Some(&format!("Bearer {token}"))))
            .await
            .expect("response");
        assert_eq!(ok.status(), StatusCode::OK);

        // The tripwire that keeps a future logging change (e.g. logging
        // request headers) from silently leaking credentials.
        let logged = String::from_utf8_lossy(&capture.0.lock().unwrap()).into_owned();
        assert!(
            !logged.contains(&token),
            "the bearer token leaked into the logs: {logged}"
        );
        // The rejections themselves do get logged (that's the point).
        assert!(logged.contains("missing or invalid bearer token"));
    }

    /// A token spooled by `--enroll-link` is refused by nothing: the first
    /// request carrying it is served, and from then on it is in the config
    /// like any other.
    #[tokio::test]
    async fn a_spooled_enrollment_is_adopted_on_first_use() {
        // Under a subscriber, like every other exercise of this middleware.
        // `tracing` caches a callsite's interest process-wide the first time it
        // is reached, so the refusal at the end of this test -- reached with no
        // subscriber on this thread -- could cache the rejection warning as
        // never-enabled and make the tripwire above see an empty log. That
        // failed about one full-suite run in ten, in the test that exists to
        // notice a credential leak, which is the worst place for a flake.
        let _guard = tracing::subscriber::set_default(
            tracing_subscriber::fmt()
                .with_max_level(tracing::Level::TRACE)
                .with_writer(std::io::sink)
                .finish(),
        );

        let dir = tempfile::tempdir().expect("tempdir");
        let manager = manager_with_token(dir.path(), "first");
        let spooled = generate_token();
        spool_pending(&manager.pending_enrollments_dir(), "tablet", &spooled).unwrap();

        let router = guarded_router(Arc::clone(&manager));
        let bearer = format!("Bearer {spooled}");
        let served = router
            .clone()
            .oneshot(request("/probe", Some(&bearer)))
            .await
            .unwrap();
        assert_eq!(served.status(), StatusCode::OK);
        let names: Vec<String> = manager.tokens().into_iter().map(|t| t.name).collect();
        assert_eq!(names, vec!["phone".to_string(), "tablet".to_string()]);
        assert!(
            std::fs::read_dir(manager.pending_enrollments_dir())
                .unwrap()
                .next()
                .is_none(),
            "the spool entry is consumed"
        );

        let again = router
            .oneshot(request("/probe", Some(&bearer)))
            .await
            .unwrap();
        assert_eq!(again.status(), StatusCode::OK, "now an ordinary token");
        let stranger = guarded_router(manager)
            .oneshot(request("/probe", Some("Bearer nope")))
            .await
            .unwrap();
        assert_eq!(stranger.status(), StatusCode::UNAUTHORIZED);
    }
}