tftio-kb 2.5.3

//! Bearer-token authentication middleware (T06-03).
//!
//! Reads the expected token at startup and rejects requests without a
//! matching `Authorization: Bearer <token>` header. Token comparison uses
//! constant-time equality.
//!
//! Ported from Haskell `KB/Auth.hs`.

use axum::{
    extract::Request,
    http::{StatusCode, header},
    middleware::Next,
    response::{IntoResponse, Response},
};

/// The expected bearer token, held in memory. Never logged.
#[derive(Clone)]
pub struct BearerToken {
    expected: Vec<u8>,
}

impl BearerToken {
    /// Build a `BearerToken` from a string value (typically from `KB_TOKEN` env).
    #[must_use]
    pub fn new(token: &str) -> Self {
        BearerToken {
            expected: token.as_bytes().to_vec(),
        }
    }

    /// Check whether a request carries the expected bearer token.
    fn check(&self, req: &Request) -> bool {
        let Some(auth_header) = req.headers().get(header::AUTHORIZATION) else {
            return false;
        };
        let Ok(value) = auth_header.to_str() else {
            return false;
        };
        let Some(presented) = value.strip_prefix("Bearer ") else {
            return false;
        };
        const_time_eq(&self.expected, presented.as_bytes())
    }
}

/// Axum middleware that gates requests on a matching bearer token.
pub async fn bearer_auth(
    token: axum::extract::State<BearerToken>,
    req: Request,
    next: Next,
) -> Response {
    if token.check(&req) {
        next.run(req).await
    } else {
        unauthorized()
    }
}

fn unauthorized() -> Response {
    (
        StatusCode::UNAUTHORIZED,
        [
            (header::WWW_AUTHENTICATE, "Bearer"),
            (header::CONTENT_TYPE, "text/plain; charset=utf-8"),
        ],
        "unauthorized\n",
    )
        .into_response()
}

/// Constant-time byte slice equality. Length check short-circuits (length
/// is not secret), but per-byte comparison walks the full input.
fn const_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.iter()
        .zip(b.iter())
        .fold(0u8, |acc, (x, y)| acc | (x ^ y))
        == 0
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{Router, body::Body, middleware, routing::get as axum_get};
    use tower::ServiceExt;

    fn test_app(token: &str) -> Router {
        let bt = BearerToken::new(token);
        Router::new()
            .route("/", axum_get(|| async { "ok" }))
            .layer(middleware::from_fn_with_state(bt, bearer_auth))
    }

    async fn send_req(app: &Router, auth: Option<&str>) -> (StatusCode, String) {
        let mut req = Request::builder().uri("/");
        if let Some(token) = auth {
            req = req.header("Authorization", format!("Bearer {token}"));
        }
        let resp = app
            .clone()
            .oneshot(req.body(Body::empty()).unwrap())
            .await
            .unwrap();
        let status = resp.status();
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        (status, String::from_utf8_lossy(&body).to_string())
    }

    #[tokio::test]
    async fn correct_token_passes() {
        let app = test_app("secret");
        let (status, body) = send_req(&app, Some("secret")).await;
        assert_eq!(status, StatusCode::OK);
        assert_eq!(body, "ok");
    }

    #[tokio::test]
    async fn wrong_token_rejected() {
        let app = test_app("secret");
        let (status, _) = send_req(&app, Some("wrong")).await;
        assert_eq!(status, StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn missing_token_rejected() {
        let app = test_app("secret");
        let (status, _) = send_req(&app, None).await;
        assert_eq!(status, StatusCode::UNAUTHORIZED);
    }

    #[test]
    fn const_time_eq_matches() {
        assert!(const_time_eq(b"abc", b"abc"));
        assert!(!const_time_eq(b"abc", b"abd"));
        assert!(!const_time_eq(b"abc", b"ab"));
        assert!(!const_time_eq(b"", b"a"));
        assert!(const_time_eq(b"", b""));
    }
}