Middleware¶

use tower_http::trace::TraceLayer;

let mut app = App::new();
app.layer(TraceLayer::new_for_http());

use axum::{
    extract::Request,
    http::{header, StatusCode},
    middleware::Next,
    response::Response,
    Extension,
};
use jsonwebtoken::{decode, DecodingKey, Validation};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Claims {
    pub sub: String,
    pub roles: Vec<String>,
    pub exp: usize,
}

#[derive(Debug, Clone)]
pub struct AuthContext {
    pub user_id: String,
    pub roles: Vec<String>,
}

pub async fn auth_guard(
    mut request: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    // Extract token from Authorization header
    let auth_header = request
        .headers()
        .get(header::AUTHORIZATION)
        .and_then(|h| h.to_str().ok())
        .ok_or(StatusCode::UNAUTHORIZED)?;

    if !auth_header.starts_with("Bearer ") {
        return Err(StatusCode::UNAUTHORIZED);
    }

    let token = &auth_header[7..]; // Strip "Bearer "

    // Verify and decode JWT
    let token_data = decode::<Claims>(
        token,
        &DecodingKey::from_secret(b"your-secret-key"),
        &Validation::default(),
    )
    .map_err(|_| StatusCode::UNAUTHORIZED)?;

    // Enrich context with authenticated user
    let auth_ctx = AuthContext {
        user_id: token_data.claims.sub,
        roles: token_data.claims.roles,
    };

    request.extensions_mut().insert(auth_ctx);

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

// Usage with App:
// app.layer(axum::middleware::from_fn(auth_guard));

use axum::{
    extract::Request,
    http::HeaderValue,
    middleware::Next,
    response::Response,
};
use tracing::{info, span, Instrument, Level};
use uuid::Uuid;

pub async fn observability_middleware(
    request: Request,
    next: Next,
) -> Response {
    // Generate or propagate request ID
    let request_id = request
        .headers()
        .get("x-request-id")
        .and_then(|h| h.to_str().ok())
        .map(String::from)
        .unwrap_or_else(|| Uuid::new_v4().to_string());

    // Create span for distributed tracing
    let span = span!(
        Level::INFO,
        "http_request",
        request_id = %request_id,
        method = %request.method(),
        path = %request.uri().path(),
        user_agent = request
            .headers()
            .get("user-agent")
            .and_then(|h| h.to_str().ok())
            .unwrap_or("unknown"),
    );

    // Log request start
    info!(parent: &span, "request_started");

    let start = std::time::Instant::now();

    // Execute request within span
    let mut response = next.run(request).instrument(span.clone()).await;

    let duration_ms = start.elapsed().as_millis();

    // Log request completion
    info!(
        parent: &span,
        status = %response.status().as_u16(),
        duration_ms = %duration_ms,
        "request_completed"
    );

    // Propagate request ID in response headers
    response.headers_mut().insert(
        "x-request-id",
        HeaderValue::from_str(&request_id).unwrap_or_else(|_| HeaderValue::from_static("")),
    );

    response
}

// Usage with App:
// app.layer(axum::middleware::from_fn(observability_middleware));

use axum::{
    body::Body,
    extract::Request,
    http::{header, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
    Extension,
};
use governor::{Quota, RateLimiter};
use nonzero_ext::nonzero;
use std::{net::IpAddr, sync::Arc};

#[derive(Debug, Clone)]
pub struct RequestContext {
    pub tenant: String,
    pub features: Vec<String>,
}

type SharedRateLimiter = Arc<RateLimiter<IpAddr, governor::state::keyed::DefaultKeyedStateStore<IpAddr>, governor::clock::DefaultClock>>;

pub fn create_rate_limiter() -> SharedRateLimiter {
    // 100 requests per minute per IP
    Arc::new(RateLimiter::keyed(Quota::per_minute(nonzero!(100u32))))
}

pub async fn request_shaper(
    Extension(limiter): Extension<SharedRateLimiter>,
    mut request: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    // 1. Rate limiting
    let client_ip: IpAddr = request
        .headers()
        .get("x-forwarded-for")
        .and_then(|h| h.to_str().ok())
        .and_then(|s| s.split(',').next())
        .and_then(|s| s.trim().parse().ok())
        .unwrap_or([127, 0, 0, 1].into());

    if limiter.check_key(&client_ip).is_err() {
        return Err(StatusCode::TOO_MANY_REQUESTS);
    }

    // 2. Extract tenant from subdomain
    let host = request
        .headers()
        .get(header::HOST)
        .and_then(|h| h.to_str().ok())
        .unwrap_or("");

    let tenant = host
        .split('.')
        .next()
        .filter(|s| !s.is_empty())
        .unwrap_or("default")
        .to_string();

    // 3. Parse feature flags from query params
    let features: Vec<String> = request
        .uri()
        .query()
        .and_then(|q| {
            q.split('&')
                .find(|p| p.starts_with("features="))
                .map(|p| p.trim_start_matches("features="))
        })
        .map(|f| f.split(',').map(String::from).collect())
        .unwrap_or_default();

    // Inject context for handlers
    let ctx = RequestContext { tenant, features };
    request.extensions_mut().insert(ctx);

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

// Usage with App:
// let limiter = create_rate_limiter();
// app.layer(Extension(limiter));
// app.layer(axum::middleware::from_fn(request_shaper));

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

    async fn handler() -> &'static str {
        "OK"
    }

    #[tokio::test]
    async fn test_auth_guard_valid_token() {
        let app = Router::new()
            .route("/", get(handler))
            .layer(axum::middleware::from_fn(auth_guard));

        let valid_token = create_test_jwt("user-123", vec!["admin".into()]);

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/")
                    .header("Authorization", format!("Bearer {}", valid_token))
                    .body(Body::empty())
                    .expect("failed to build request"),
            )
            .await
            .expect("request failed");

        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_auth_guard_missing_token() {
        let app = Router::new()
            .route("/", get(handler))
            .layer(axum::middleware::from_fn(auth_guard));

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/")
                    .body(Body::empty())
                    .expect("failed to build request"),
            )
            .await
            .expect("request failed");

        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn test_auth_guard_invalid_token() {
        let app = Router::new()
            .route("/", get(handler))
            .layer(axum::middleware::from_fn(auth_guard));

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/")
                    .header("Authorization", "Bearer invalid-token")
                    .body(Body::empty())
                    .expect("failed to build request"),
            )
            .await
            .expect("request failed");

        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
    }

    fn create_test_jwt(sub: &str, roles: Vec<String>) -> String {
        use jsonwebtoken::{encode, EncodingKey, Header};

        let claims = Claims {
            sub: sub.to_string(),
            roles,
            exp: (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp() as usize,
        };

        encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(b"your-secret-key"),
        )
        .expect("failed to encode JWT")
    }
}