mod healthcheck; use std::sync::mpsc; use std::time::Duration; use axum::extract::Request; use axum::middleware::{Next, from_fn}; use axum::response::Response; #[allow(unused_imports)] use axum::routing::{get, post}; use axum::{Router, body::Body}; use http::StatusCode; use tower::ServiceBuilder; use tower_http::timeout::TimeoutLayer; use tower_http::trace::TraceLayer; use tracing::{Level, warn}; use uuid::Uuid; pub fn app() -> (Router, mpsc::Receiver) { let (tx, rx) = mpsc::channel(); ( Router::new() .route("/healthcheck", get(healthcheck::healthcheck)) .route("/forever", get(std::future::pending::<()>)) .nest("/admin", admin(tx)) .layer( ServiceBuilder::new() .layer( TraceLayer::new_for_http().make_span_with(|req: &Request| { tracing::span!( Level::DEBUG, "request", trace_id = Uuid::now_v7().to_string(), method = format!("{}", req.method()), uri = format!("{}", req.uri()), ) }), ) .layer(TimeoutLayer::new(Duration::from_secs(10))), ), rx, ) } fn admin(tx: mpsc::Sender) -> Router { let r = Router::new().route("/", get(async || StatusCode::OK)); let r = add_shutdown_endpoint(r, tx); r.layer(from_fn(async |req: Request, next: Next| { if let Ok(secret) = std::env::var("ADMIN_SECRET") { match req.headers().get("Authorization") { Some(key) if secret == *key => (), Some(key) => { warn!("Unauthorized request with key: {key:?}"); return Response::builder() .status(StatusCode::UNAUTHORIZED) .body(Body::empty()) .unwrap(); } _ => { warn!("Unauthorized request no key given"); return Response::builder() .status(StatusCode::UNAUTHORIZED) .body(Body::empty()) .unwrap(); } } } next.run(req).await })) } #[cfg(feature = "shutdown")] fn add_shutdown_endpoint(r: Router, tx: mpsc::Sender) -> Router { r.route( "/shutdown", post(async move || { let res = tx.send(true); if res.is_ok() { StatusCode::OK } else { StatusCode::INTERNAL_SERVER_ERROR } }), ) } #[cfg(not(feature = "shutdown"))] fn add_shutdown_endpoint(r: Router, _: mpsc::Sender) -> Router { r } #[cfg(test)] mod tests { use super::*; use tower::{self, ServiceExt}; #[tokio::test] async fn test_authorization_disables_when_no_env_var_set() { let (app, _) = app(); let resp = app .oneshot( axum::http::Request::builder() .uri("/admin") .body(Body::empty()) .unwrap(), ) .await .unwrap(); assert_eq!(resp.status(), StatusCode::OK); } }