aster_forge_middleware/axum/
request_id.rs

1//! Axum request-id adapter.
2use axum::{extract::Request, middleware::Next, response::Response};
3use http::header::{HeaderName, HeaderValue};
4use tracing::Instrument;
5
6#[derive(Clone, Debug)]
7/// Server-generated UUID stored in request extensions.
8pub struct RequestId(pub String);
9
10/// Generates a new UUID v4 per request, ignoring inbound IDs, and returns it in the response.
11pub async fn request_id(request: Request, next: Next) -> Response {
12    let request_id = uuid::Uuid::new_v4().to_string();
13    let method = request.method().to_string();
14    let path = request.uri().path().to_string();
15    let mut request = request;
16    request
17        .extensions_mut()
18        .insert(RequestId(request_id.clone()));
19    let span = tracing::info_span!("request", request_id = %request_id, method = %method, path = %path, user_id = tracing::field::Empty);
20    let mut response = next.run(request).instrument(span).await;
21    if let Ok(value) = HeaderValue::from_str(&request_id) {
22        response
23            .headers_mut()
24            .insert(HeaderName::from_static("x-request-id"), value);
25    }
26    response
27}
28
29#[cfg(test)]
30mod tests {
31    use super::*;
32    use axum::{
33        Extension, Router,
34        body::{Body, to_bytes},
35        middleware,
36        response::IntoResponse,
37        routing::get,
38    };
39    use http::{Request as HttpRequest, StatusCode};
40    use tower::ServiceExt;
41
42    #[tokio::test]
43    async fn generates_stores_and_returns_unique_ids_for_success_and_failure() {
44        let router = Router::new()
45            .route(
46                "/",
47                get(|Extension(id): Extension<RequestId>| async move {
48                    let mut response = (StatusCode::BAD_REQUEST, id.0).into_response();
49                    response
50                        .headers_mut()
51                        .insert("x-request-id", HeaderValue::from_static("downstream-id"));
52                    response
53                }),
54            )
55            .layer(middleware::from_fn(request_id));
56        let mut ids = Vec::new();
57        for inbound in [
58            None,
59            Some("untrusted-id"),
60            Some("11111111-1111-4111-8111-111111111111"),
61        ] {
62            let mut request = HttpRequest::builder();
63            if let Some(inbound) = inbound {
64                request = request.header("x-request-id", inbound);
65            }
66            let response = router
67                .clone()
68                .oneshot(request.body(Body::empty()).unwrap())
69                .await
70                .unwrap();
71            assert_eq!(response.status(), StatusCode::BAD_REQUEST);
72            let id = response.headers()["x-request-id"]
73                .to_str()
74                .unwrap()
75                .to_string();
76            assert_eq!(uuid::Uuid::parse_str(&id).unwrap().get_version_num(), 4);
77            assert_ne!(Some(id.as_str()), inbound);
78            assert_eq!(
79                to_bytes(response.into_body(), 4096).await.unwrap().as_ref(),
80                id.as_bytes()
81            );
82            ids.push(id);
83        }
84        assert_ne!(ids[0], ids[1]);
85        assert_ne!(ids[1], ids[2]);
86    }
87}
aster_forge_middleware/axum/request_id.rs

aster_forge_middleware/axum/
request_id.rs
