mod healthcheck; use std::sync::mpsc; use std::time::Duration; use axum::extract::Request; use axum::middleware::{Next, from_fn}; use axum::response::Response; 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; 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("/test", 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") { println!("ADMIN_SECRET: {secret}"); match req.headers().get("Authorization") { Some(key) if secret == key.to_owned() => (), _ => { 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 }