axum 0.8 and tower Middleware

[dependencies]
axum = { version = "0.8", features = ["macros"] }   # "macros" for #[debug_handler]; default features include json, query, tokio, http1
tokio = { version = "1", features = ["macros", "rt-multi-thread", "net", "signal", "time"] }
tower = { version = "0.5", features = ["util", "limit", "load-shed", "timeout"] }
tower-http = { version = "0.6", features = ["trace", "timeout", "cors", "compression-gzip", "limit",
                                            "request-id", "catch-panic", "sensitive-headers"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
thiserror = "2"
tracing = "0.1"

use axum::{Router, routing::{get, post}};

fn api_router() -> Router<AppStateRef> {
    Router::new()
        .route("/orders", get(list_orders).post(create_order))       // several methods, one path
        .route("/orders/{id}", get(get_order).delete(delete_order))   // capture
        .route("/files/{*path}", get(serve_file))                     // wildcard (rest of path)
        .route("/users/{user_id}/orders/{order_id}", get(user_order)) // several captures
}

fn app(state: AppStateRef) -> Router {
    Router::new()
        .nest("/api/v1", api_router())   // prefix
        .route("/healthz", get(|| async { "ok" }))
        .fallback(not_found)             // 404 for everything else
        .with_state(state)
}

use axum::{Json, extract::{Path, Query, State}, http::StatusCode};

#[derive(serde::Deserialize)]
struct Paging { limit: Option<u32>, cursor: Option<String> }

#[axum::debug_handler]
async fn create_order(
    State(state): State<AppStateRef>,
    Path(tenant): Path<String>,
    Query(paging): Query<Paging>,
    Json(input): Json<NewOrder>,      // last
) -> Result<(StatusCode, Json<Order>), AppError> {
    let order = state.repo.create(&tenant, input).await?;
    let _ = paging;
    Ok((StatusCode::CREATED, Json(order)))
}

use std::sync::Arc;

pub struct AppState {
    pub db: sqlx::PgPool,
    pub http: reqwest::Client,
    pub config: Config,
}
pub type AppStateRef = Arc<AppState>;

use axum::extract::FromRef;

#[derive(Clone, FromRef)]
struct AppState2 {
    db: sqlx::PgPool,
    http: reqwest::Client,
}

async fn ping_db(axum::extract::State(db): axum::extract::State<sqlx::PgPool>) -> &'static str {
    let _ = db;
    "ok"
}

use axum::{
    extract::FromRequestParts,
    http::{StatusCode, request::Parts},
};

pub struct ApiKey(pub String);

impl<S: Send + Sync> FromRequestParts<S> for ApiKey {
    type Rejection = (StatusCode, &'static str);

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        parts
            .headers
            .get("x-api-key")
            .and_then(|v| v.to_str().ok())
            .map(|v| ApiKey(v.to_owned()))
            .ok_or((StatusCode::UNAUTHORIZED, "missing x-api-key"))
    }
}

use axum::{Json, http::StatusCode, response::{IntoResponse, Response}};
use serde_json::json;

#[derive(Debug, thiserror::Error)]
pub enum AppError {
    #[error("{0} not found")]
    NotFound(&'static str),
    #[error("validation failed: {0}")]
    Validation(String),
    #[error("conflict: {0}")]
    Conflict(String),
    #[error("upstream timed out")]
    UpstreamTimeout,
    #[error(transparent)]
    Database(#[from] sqlx::Error),
    #[error(transparent)]
    Upstream(#[from] reqwest::Error),
    #[error(transparent)]
    Other(#[from] anyhow::Error),
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let (status, code) = match &self {
            AppError::NotFound(_) => (StatusCode::NOT_FOUND, "not_found"),
            AppError::Validation(_) => (StatusCode::UNPROCESSABLE_ENTITY, "validation"),
            AppError::Conflict(_) => (StatusCode::CONFLICT, "conflict"),
            AppError::UpstreamTimeout => (StatusCode::GATEWAY_TIMEOUT, "upstream_timeout"),
            AppError::Upstream(e) if e.is_timeout() => (StatusCode::GATEWAY_TIMEOUT, "upstream_timeout"),
            AppError::Upstream(_) => (StatusCode::BAD_GATEWAY, "upstream"),
            AppError::Database(sqlx::Error::RowNotFound) => (StatusCode::NOT_FOUND, "not_found"),
            AppError::Database(sqlx::Error::PoolTimedOut) => (StatusCode::SERVICE_UNAVAILABLE, "busy"),
            AppError::Database(_) | AppError::Other(_) => (StatusCode::INTERNAL_SERVER_ERROR, "internal"),
        };
        let message = if status.is_server_error() {
            tracing::error!(error = ?self, "request failed"); // full detail in logs only
            "internal error".to_owned()
        } else {
            self.to_string()
        };
        (status, Json(json!({ "error": { "code": code, "message": message } }))).into_response()
    }
}

use axum::{extract::{FromRequest, rejection::JsonRejection}, response::{IntoResponse, Response}};

#[derive(FromRequest)]
#[from_request(via(axum::Json), rejection(ApiError))]
pub struct ApiJson<T>(pub T);

pub struct ApiError(JsonRejection);

impl From<JsonRejection> for ApiError {
    fn from(r: JsonRejection) -> Self { ApiError(r) }
}

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let status = self.0.status();
        (status, axum::Json(serde_json::json!({ "error": { "code": "bad_request", "message": self.0.body_text() } })))
            .into_response()
    }
}

use std::time::Duration;
use axum::http::{HeaderName, HeaderValue, Method, StatusCode, header};
use tower::ServiceBuilder;
use tower_http::{
    catch_panic::CatchPanicLayer,
    compression::CompressionLayer,
    cors::CorsLayer,
    request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer},
    sensitive_headers::SetSensitiveRequestHeadersLayer,
    timeout::TimeoutLayer,
    trace::TraceLayer,
};

fn with_middleware(router: axum::Router) -> axum::Router {
    let x_request_id = HeaderName::from_static("x-request-id");
    let cors = CorsLayer::new()
        .allow_origin("https://app.example.com".parse::<HeaderValue>().expect("valid origin"))
        .allow_methods([Method::GET, Method::POST, Method::DELETE])
        .allow_headers([header::CONTENT_TYPE, header::AUTHORIZATION])
        .max_age(Duration::from_secs(600));

    router.layer(
        ServiceBuilder::new()
            .layer(SetSensitiveRequestHeadersLayer::new([header::AUTHORIZATION, header::COOKIE]))
            .layer(SetRequestIdLayer::new(x_request_id.clone(), MakeRequestUuid))
            .layer(TraceLayer::new_for_http())
            .layer(PropagateRequestIdLayer::new(x_request_id))
            .layer(CatchPanicLayer::new())
            .layer(TimeoutLayer::with_status_code(StatusCode::REQUEST_TIMEOUT, Duration::from_secs(10)))
            .layer(cors)
            .layer(CompressionLayer::new()),
    )
}

request -> SetSensitiveHeaders -> SetRequestId -> Trace -> PropagateRequestId -> CatchPanic -> Timeout
        -> Cors -> Compression -> handler

use axum::{BoxError, error_handling::HandleErrorLayer, http::StatusCode};
use tower::{ServiceBuilder, limit::GlobalConcurrencyLimitLayer, load_shed::LoadShedLayer};

fn with_overload_protection(router: axum::Router) -> axum::Router {
    router.layer(
        ServiceBuilder::new()
            .layer(HandleErrorLayer::new(|_err: BoxError| async { StatusCode::SERVICE_UNAVAILABLE }))
            .layer(LoadShedLayer::new())                    // reject when the inner service is not ready
            .layer(GlobalConcurrencyLimitLayer::new(512)),  // at most 512 requests in handlers
    )
}

use axum::{Router, extract::DefaultBodyLimit, routing::post};
use tower_http::limit::RequestBodyLimitLayer;

fn upload_routes() -> Router<AppStateRef> {
    Router::new()
        .route("/uploads", post(upload))
        .layer(DefaultBodyLimit::disable())
        .layer(RequestBodyLimitLayer::new(250 * 1024 * 1024))
}

use axum::{extract::{Request, State}, http::StatusCode, middleware::{self, Next}, response::Response};

async fn require_api_key(
    State(state): State<AppStateRef>,
    req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let ok = req
        .headers()
        .get("x-api-key")
        .and_then(|v| v.to_str().ok())
        .is_some_and(|k| state.config.api_keys.iter().any(|known| known == k));
    if !ok {
        return Err(StatusCode::UNAUTHORIZED);
    }
    Ok(next.run(req).await)
}

fn protected(state: AppStateRef) -> axum::Router<AppStateRef> {
    axum::Router::new()
        .route("/admin/stats", axum::routing::get(stats))
        .route_layer(middleware::from_fn_with_state(state, require_api_key))
}

use std::{sync::atomic::Ordering, time::Duration};
use axum::{extract::State, http::StatusCode};

async fn readyz(State(state): State<AppStateRef>) -> StatusCode {
    if state.shutting_down.load(Ordering::Relaxed) {
        return StatusCode::SERVICE_UNAVAILABLE;
    }
    let ping = sqlx::query("SELECT 1").execute(&state.db);
    match tokio::time::timeout(Duration::from_millis(500), ping).await {
        Ok(Ok(_)) => StatusCode::OK,
        _ => StatusCode::SERVICE_UNAVAILABLE,
    }
}

use std::net::SocketAddr;
use axum::extract::ConnectInfo;

async fn whoami(ConnectInfo(addr): ConnectInfo<SocketAddr>) -> String {
    addr.to_string()
}

async fn run(app: axum::Router) -> anyhow::Result<()> {
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?;
    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
        .with_graceful_shutdown(async { let _ = tokio::signal::ctrl_c().await; })
        .await?;
    Ok(())
}