aro_web/

middleware.rs

//! Built-in middleware helpers for common cross-cutting concerns.
//!
//! Provides pre-configured CORS and request-ID layers, plus re-exports of
//! Axum's middleware combinators.

use axum::http::{HeaderName, HeaderValue, Method, Request, header};
use tower_http::cors::{AllowHeaders, AllowMethods, AllowOrigin, CorsLayer};
use tower_http::request_id::{
    MakeRequestId, PropagateRequestIdLayer, RequestId, SetRequestIdLayer,
};
use tower_http::set_header::SetResponseHeaderLayer;

// Re-export axum's middleware combinators so users can write
// `aro::middleware::from_fn` without depending on axum directly.
pub use axum::middleware::*;

/// Header name used for request IDs.
pub static X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");

// ---------------------------------------------------------------------------
// CORS
// ---------------------------------------------------------------------------

/// Configuration for the CORS middleware.
#[derive(Debug, Clone)]
pub struct CorsConfig {
    /// Origins allowed to access the API (see [`AllowOrigin`]).
    pub allow_origins: AllowOrigin,
    /// HTTP methods allowed for cross-origin requests.
    pub allow_methods: AllowMethods,
    /// Request headers browsers may send on cross-origin requests.
    pub allow_headers: AllowHeaders,
}

impl Default for CorsConfig {
    /// Returns a permissive configuration suitable for development.
    fn default() -> Self {
        Self {
            allow_origins: AllowOrigin::any(),
            allow_methods: AllowMethods::list([
                Method::GET,
                Method::POST,
                Method::PUT,
                Method::PATCH,
                Method::DELETE,
                Method::OPTIONS,
            ]),
            allow_headers: AllowHeaders::list([
                header::CONTENT_TYPE,
                header::AUTHORIZATION,
                header::ACCEPT,
                X_REQUEST_ID.clone(),
            ]),
        }
    }
}

/// Returns a CORS layer with permissive defaults suitable for development.
///
/// Allows any origin, common HTTP methods, and typical headers.
pub fn cors() -> CorsLayer {
    cors_with_config(CorsConfig::default())
}

/// Returns a CORS layer configured according to the given [`CorsConfig`].
pub fn cors_with_config(config: CorsConfig) -> CorsLayer {
    CorsLayer::new()
        .allow_origin(config.allow_origins)
        .allow_methods(config.allow_methods)
        .allow_headers(config.allow_headers)
}

// ---------------------------------------------------------------------------
// Request ID
// ---------------------------------------------------------------------------

/// UUID v4 request ID generator.
#[derive(Clone, Copy)]
pub struct UuidRequestId;

impl MakeRequestId for UuidRequestId {
    fn make_request_id<B>(&mut self, _request: &Request<B>) -> Option<RequestId> {
        let id = uuid::Uuid::new_v4().to_string();
        Some(RequestId::new(HeaderValue::from_str(&id).ok()?))
    }
}

/// Returns a tuple of layers that set and propagate an `x-request-id` header.
///
/// - `SetRequestIdLayer` assigns a UUID v4 request ID on incoming requests
///   that do not already carry one.
/// - `PropagateRequestIdLayer` copies the request ID to the response.
pub fn request_id() -> (SetRequestIdLayer<UuidRequestId>, PropagateRequestIdLayer) {
    (
        SetRequestIdLayer::new(X_REQUEST_ID.clone(), UuidRequestId),
        PropagateRequestIdLayer::new(X_REQUEST_ID.clone()),
    )
}

// ---------------------------------------------------------------------------
// Compression
// ---------------------------------------------------------------------------

/// Returns a compression layer that negotiates encoding based on the
/// `Accept-Encoding` request header.
///
/// Enabled algorithms depend on the active feature flags:
/// - `compression` enables gzip and brotli
/// - `compression-full` adds zstd and deflate
#[cfg(feature = "compression")]
pub fn compression() -> tower_http::compression::CompressionLayer {
    tower_http::compression::CompressionLayer::new()
}

// ---------------------------------------------------------------------------
// Timeout
// ---------------------------------------------------------------------------

/// Returns a layer that fails requests which are not completed within the
/// given `duration`.
///
/// Requests that exceed the timeout receive a `408 Request Timeout` response.
///
/// Use this when composing layers on a raw [`axum::Router`]. When building
/// via [`App`](crate::App), prefer [`App::timeout`](crate::App::timeout),
/// which applies the same layer to all routes including stateful DI routes.
pub fn timeout(duration: std::time::Duration) -> tower_http::timeout::TimeoutLayer {
    tower_http::timeout::TimeoutLayer::with_status_code(
        axum::http::StatusCode::REQUEST_TIMEOUT,
        duration,
    )
}

// ---------------------------------------------------------------------------
// Cache-Control
// ---------------------------------------------------------------------------

/// Returns a layer that sets the `Cache-Control` response header if not
/// already present.
///
/// The provided `directive` is used as the header value (e.g.
/// `"max-age=3600"` or `"no-cache, no-store, must-revalidate"`).
#[expect(
    clippy::expect_used,
    reason = "caller-provided directive is validated at the API boundary"
)]
pub fn cache_control(directive: &str) -> SetResponseHeaderLayer<HeaderValue> {
    let value = HeaderValue::from_str(directive).expect("valid Cache-Control header value");
    SetResponseHeaderLayer::if_not_present(header::CACHE_CONTROL, value)
}

/// Returns a layer that sets `Cache-Control: no-cache, no-store, must-revalidate`
/// if the response does not already include a `Cache-Control` header.
pub fn no_cache() -> SetResponseHeaderLayer<HeaderValue> {
    cache_control("no-cache, no-store, must-revalidate")
}

/// Returns a layer that sets `Cache-Control: no-store` if the response does
/// not already include a `Cache-Control` header.
pub fn no_store() -> SetResponseHeaderLayer<HeaderValue> {
    cache_control("no-store")
}

// ---------------------------------------------------------------------------
// Request Decompression
// ---------------------------------------------------------------------------

/// Returns a layer that transparently decompresses request bodies based on
/// the `Content-Encoding` header (gzip, br, and optionally zstd/deflate
/// when the `decompression-full` feature is enabled).
#[cfg(feature = "decompression")]
pub fn decompression() -> tower_http::decompression::RequestDecompressionLayer {
    tower_http::decompression::RequestDecompressionLayer::new()
}

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

    async fn ok_handler() -> &'static str {
        "ok"
    }

    #[tokio::test]
    async fn cors_allows_cross_origin_with_default_config() {
        let app = Router::new().route("/", get(ok_handler)).layer(cors());

        let request = http::Request::builder()
            .uri("/")
            .header(header::ORIGIN, "http://example.com")
            .body(Body::empty())
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(
            response
                .headers()
                .get(header::ACCESS_CONTROL_ALLOW_ORIGIN)
                .is_some()
        );
    }

    #[tokio::test]
    async fn cors_can_be_customized_with_specific_origins() {
        let config = CorsConfig {
            allow_origins: AllowOrigin::exact("http://allowed.com".parse().unwrap()),
            ..CorsConfig::default()
        };
        let app = Router::new()
            .route("/", get(ok_handler))
            .layer(cors_with_config(config));

        // Request from allowed origin
        let request = http::Request::builder()
            .uri("/")
            .header(header::ORIGIN, "http://allowed.com")
            .body(Body::empty())
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get(header::ACCESS_CONTROL_ALLOW_ORIGIN)
                .unwrap(),
            "http://allowed.com"
        );
    }

    #[tokio::test]
    async fn request_id_adds_header_to_response() {
        let (set_id, propagate_id) = request_id();
        let app = Router::new()
            .route("/", get(ok_handler))
            .layer(propagate_id)
            .layer(set_id);

        let request = http::Request::builder()
            .uri("/")
            .body(Body::empty())
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        let rid = response
            .headers()
            .get("x-request-id")
            .expect("x-request-id header should be present");
        // Should be a valid UUID v4
        let parsed = uuid::Uuid::parse_str(rid.to_str().unwrap());
        assert!(parsed.is_ok(), "x-request-id should be a valid UUID");
    }

    #[tokio::test]
    async fn request_id_preserves_existing_header() {
        let existing_id = "my-custom-request-id-123";
        let (set_id, propagate_id) = request_id();
        let app = Router::new()
            .route("/", get(ok_handler))
            .layer(propagate_id)
            .layer(set_id);

        let request = http::Request::builder()
            .uri("/")
            .header("x-request-id", existing_id)
            .body(Body::empty())
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get("x-request-id")
                .unwrap()
                .to_str()
                .unwrap(),
            existing_id
        );
    }

    #[tokio::test]
    async fn middleware_composes_with_app_layer() {
        let (set_id, propagate_id) = request_id();
        let app = crate::App::new()
            .routes(Router::new().route("/", get(ok_handler)))
            .layer(cors())
            .layer(propagate_id)
            .layer(set_id)
            .build();

        let request = http::Request::builder()
            .uri("/")
            .header(header::ORIGIN, "http://example.com")
            .body(Body::empty())
            .unwrap();

        let response = ServiceExt::<http::Request<Body>>::oneshot(app, request)
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        // CORS header present
        assert!(
            response
                .headers()
                .get(header::ACCESS_CONTROL_ALLOW_ORIGIN)
                .is_some()
        );
        // Request ID present
        assert!(response.headers().get("x-request-id").is_some());
    }

    #[tokio::test]
    async fn axum_from_fn_is_accessible() {
        // Verify the re-export works by creating a middleware from a function
        async fn noop_middleware(
            request: http::Request<Body>,
            next: axum::middleware::Next,
        ) -> axum::response::Response {
            next.run(request).await
        }

        let _layer = from_fn::<_, ()>(noop_middleware);
    }

    #[tokio::test]
    async fn cache_control_sets_header_on_response() {
        let app = Router::new()
            .route("/", get(ok_handler))
            .layer(cache_control("max-age=3600"));

        let request = http::Request::builder()
            .uri("/")
            .body(Body::empty())
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get(header::CACHE_CONTROL)
                .unwrap()
                .to_str()
                .unwrap(),
            "max-age=3600"
        );
    }

    #[tokio::test]
    async fn cache_control_does_not_overwrite_handler_set_header() {
        async fn handler_with_cache() -> impl axum::response::IntoResponse {
            ([(header::CACHE_CONTROL, "private, max-age=60")], "ok")
        }

        let app = Router::new()
            .route("/", get(handler_with_cache))
            .layer(no_cache());

        let request = http::Request::builder()
            .uri("/")
            .body(Body::empty())
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get(header::CACHE_CONTROL)
                .unwrap()
                .to_str()
                .unwrap(),
            "private, max-age=60"
        );
    }

    #[tokio::test]
    async fn no_cache_sets_full_directive() {
        let app = Router::new().route("/", get(ok_handler)).layer(no_cache());

        let request = http::Request::builder()
            .uri("/")
            .body(Body::empty())
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get(header::CACHE_CONTROL)
                .unwrap()
                .to_str()
                .unwrap(),
            "no-cache, no-store, must-revalidate"
        );
    }
}