Axum + Tower middleware pattern

Axum bagusnya pakai Tower untuk middleware — pattern Service trait yang composable. Tapi custom middleware-nya rumit dibanding Express. Snippet ini implement request ID propagation, auth check, dan simple rate limit dengan idiomatic Rust. Plus extractor untuk inject ke handler.

Kode

// Cargo.toml
// [dependencies]
// axum = "0.7"
// tower = { version = "0.5", features = ["util"] }
// tower-http = { version = "0.5", features = ["trace", "request-id", "cors"] }
// tokio = { version = "1", features = ["full"] }
// uuid = { version = "1.10", features = ["v4"] }
// http = "1"

// main.rs
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use axum::{
    body::Body,
    extract::{Extension, Request, State},
    http::{HeaderName, HeaderValue, StatusCode},
    middleware::{self, Next},
    response::{IntoResponse, Response},
    routing::get,
    Json, Router,
};
use serde_json::json;
use tower_http::trace::TraceLayer;
use uuid::Uuid;

/// Shared state untuk rate limiter
#[derive(Clone)]
struct AppState {
    rate_limiter: Arc<Mutex<HashMap<String, Vec<Instant>>>>,
}

/// Authenticated user — di-extract dari token middleware
#[derive(Clone, Debug)]
struct AuthUser {
    id: String,
    email: String,
    roles: Vec<String>,
}

// ==========================================
// Middleware 1: Request ID propagation
// ==========================================

const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");

#[derive(Clone, Debug)]
struct RequestId(String);

async fn request_id_middleware(mut req: Request, next: Next) -> Response {
    // Gunakan request-id dari header kalau ada (upstream gateway), else generate
    let id = req
        .headers()
        .get(&REQUEST_ID_HEADER)
        .and_then(|h| h.to_str().ok())
        .map(String::from)
        .unwrap_or_else(|| Uuid::new_v4().to_string());

    // Inject ke extensions supaya handler bisa akses
    req.extensions_mut().insert(RequestId(id.clone()));

    let mut response = next.run(req).await;

    // Echo balik di response header
    if let Ok(header_val) = HeaderValue::from_str(&id) {
        response.headers_mut().insert(REQUEST_ID_HEADER, header_val);
    }

    response
}

// ==========================================
// Middleware 2: Auth (validate Bearer token)
// ==========================================

async fn auth_middleware(mut req: Request, next: Next) -> Result<Response, AppError> {
    let token = req
        .headers()
        .get(http::header::AUTHORIZATION)
        .and_then(|h| h.to_str().ok())
        .and_then(|s| s.strip_prefix("Bearer "))
        .ok_or_else(|| AppError {
            status: StatusCode::UNAUTHORIZED,
            message: "Token tidak ada".into(),
        })?;

    let user = verify_token(token).map_err(|e| AppError {
        status: StatusCode::UNAUTHORIZED,
        message: format!("Token tidak valid: {e}"),
    })?;

    // Inject AuthUser ke extensions
    req.extensions_mut().insert(user);

    Ok(next.run(req).await)
}

fn verify_token(token: &str) -> Result<AuthUser, String> {
    // Dummy — di production: jwt verify, check expiration, dll
    if token.starts_with("valid-") {
        Ok(AuthUser {
            id: "user-123".to_string(),
            email: "[email protected]".to_string(),
            roles: vec!["user".to_string()],
        })
    } else {
        Err("token format salah".into())
    }
}

// ==========================================
// Middleware 3: Rate limit per-IP / per-user
// ==========================================

async fn rate_limit_middleware(
    State(state): State<AppState>,
    req: Request,
    next: Next,
) -> Result<Response, AppError> {
    let key = req
        .extensions()
        .get::<AuthUser>()
        .map(|u| format!("user:{}", u.id))
        .unwrap_or_else(|| {
            req.headers()
                .get("x-forwarded-for")
                .and_then(|h| h.to_str().ok())
                .map(|s| format!("ip:{s}"))
                .unwrap_or_else(|| "ip:unknown".to_string())
        });

    let allowed = {
        let mut limiter = state.rate_limiter.lock().unwrap();
        let now = Instant::now();
        let window = Duration::from_secs(60);

        let timestamps = limiter.entry(key).or_default();
        timestamps.retain(|t| now.duration_since(*t) < window);

        if timestamps.len() < 60 {
            timestamps.push(now);
            true
        } else {
            false
        }
    };

    if !allowed {
        return Err(AppError {
            status: StatusCode::TOO_MANY_REQUESTS,
            message: "Rate limit exceeded (60 req/menit)".into(),
        });
    }

    Ok(next.run(req).await)
}

// ==========================================
// Error type
// ==========================================

struct AppError {
    status: StatusCode,
    message: String,
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let body = Json(json!({ "error": self.message }));
        (self.status, body).into_response()
    }
}

// ==========================================
// Handlers
// ==========================================

async fn public_handler(Extension(req_id): Extension<RequestId>) -> impl IntoResponse {
    Json(json!({
        "message": "Halo dari Indonesia",
        "request_id": req_id.0,
    }))
}

async fn protected_handler(
    Extension(user): Extension<AuthUser>,
    Extension(req_id): Extension<RequestId>,
) -> impl IntoResponse {
    Json(json!({
        "user": {
            "id": user.id,
            "email": user.email,
            "roles": user.roles,
        },
        "request_id": req_id.0,
    }))
}

// ==========================================
// Main — composition Tower layers
// ==========================================

#[tokio::main]
async fn main() {
    let state = AppState {
        rate_limiter: Arc::new(Mutex::new(HashMap::new())),
    };

    // Route public — cuma request_id middleware
    let public_routes = Router::new()
        .route("/", get(public_handler))
        .route("/healthz", get(|| async { "OK" }));

    // Route protected — auth + rate limit + request_id
    let protected_routes = Router::new()
        .route("/me", get(protected_handler))
        .route_layer(middleware::from_fn(auth_middleware))
        .route_layer(middleware::from_fn_with_state(
            state.clone(),
            rate_limit_middleware,
        ));

    let app = Router::new()
        .merge(public_routes)
        .nest("/api", protected_routes)
        .layer(middleware::from_fn(request_id_middleware))
        .layer(TraceLayer::new_for_http()) // tower-http: log request/response
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    println!("Listening :3000");
    axum::serve(listener, app).await.unwrap();
}

Pemakaian

# Test public endpoint
curl -i http://localhost:3000/
# HTTP/1.1 200 OK
# x-request-id: 5f3b8a2c-1d9e-4f6a-7b8c-9d0e1f2a3b4c
# content-type: application/json
#
# {"message":"Halo dari Indonesia","request_id":"5f3b8a2c-..."}

# Test protected tanpa token — 401
curl http://localhost:3000/api/me
# {"error":"Token tidak ada"}

# Test dengan token valid
curl http://localhost:3000/api/me -H "Authorization: Bearer valid-abc"
# {"user":{...},"request_id":"..."}

# Test rate limit
for i in $(seq 1 70); do
  curl -s -o /dev/null -w "%{http_code}\n" \
    http://localhost:3000/api/me \
    -H "Authorization: Bearer valid-abc"
done
# 200 ... (60 kali) ... 429 (sisanya)

# Test propagate request_id dari upstream
curl http://localhost:3000/ -H "x-request-id: upstream-abc-123"
# Response: x-request-id: upstream-abc-123 (di-echo balik)

// Pattern: middleware untuk extract custom user
// Bikin extractor pakai FromRequestParts trait — lebih ergonomis
use axum::{async_trait, extract::FromRequestParts, http::request::Parts};

#[async_trait]
impl<S> FromRequestParts<S> for AuthUser
where
    S: Send + Sync,
{
    type Rejection = AppError;

    async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Self::Rejection> {
        parts
            .extensions
            .get::<AuthUser>()
            .cloned()
            .ok_or_else(|| AppError {
                status: StatusCode::UNAUTHORIZED,
                message: "Tidak authenticated".into(),
            })
    }
}

// Sekarang handler bisa pakai AuthUser langsung (lebih clean):
async fn me_v2(user: AuthUser) -> Json<serde_json::Value> {
    Json(json!({"id": user.id, "email": user.email}))
}

Kapan dipakai

• Production Axum HTTP API.
• gRPC service dengan tonic (Tower middleware compatible).
• Service yang butuh request ID propagation untuk distributed tracing.
• API gateway pattern di Rust.

Catatan

• middleware::from_fn untuk one-off middleware. Untuk reusable cross-project, implement Tower Service / Layer trait penuh.
• route_layer vs layer — route_layer apply hanya ke route dalam, layer apply ke seluruh router. Order matters.
• Extensions — request-scoped data. Type-safe karena keyed by Rust type, bukan string.
• Tower Service trait lebih powerful tapi rumit. Untuk 90% kasus, from_fn cukup.
• Error handling via custom type + IntoResponse — idiomatic Rust pattern. Hindari panic di middleware.
• State sharing — from_fn_with_state inject state. Untuk Mutex seperti contoh, OK untuk demo. Production pakai DashMap atau Redis.
• TraceLayer built-in cukup bagus. Combine dengan tracing-subscriber untuk structured log.

Mutex di async context bisa starve worker. Untuk rate limiter production, pakai Redis atau crate khusus seperti tower-governor yang lock-free.