Skip to main content

nidus_http/middleware/
request_id.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::Arc,
5    task::{Context, Poll},
6};
7
8use axum::{body::Body, response::IntoResponse};
9use http::{HeaderValue, Request, Response, StatusCode, header::HeaderName};
10use serde::Serialize;
11use tower::{Layer, Service};
12use uuid::{Uuid, Version};
13
14use crate::context::RequestContext;
15
16/// Legacy Tower layer that adds an `x-request-id` response header when absent.
17///
18/// Incoming request IDs are propagated to the response unless the inner service
19/// already set a response ID. Requests without an ID receive a generated
20/// UUID v4 value.
21///
22/// This layer does not validate inbound IDs and does not populate
23/// [`RequestContext`]. Prefer [`validated_request_id_layer`] for production API
24/// defaults, UUID v4 generation, strict/permissive validation, request
25/// extension insertion, and consistent error responses.
26#[derive(Clone, Copy, Debug, Default)]
27pub struct RequestIdLayer;
28
29impl<S> Layer<S> for RequestIdLayer {
30    type Service = RequestIdService<S>;
31
32    fn layer(&self, inner: S) -> Self::Service {
33        RequestIdService { inner }
34    }
35}
36
37/// Service produced by [`RequestIdLayer`].
38#[derive(Clone, Debug)]
39pub struct RequestIdService<S> {
40    inner: S,
41}
42
43impl<S, RequestBody, ResponseBody> Service<Request<RequestBody>> for RequestIdService<S>
44where
45    S: Service<Request<RequestBody>, Response = Response<ResponseBody>> + Send + 'static,
46    S::Future: Send + 'static,
47    S::Error: Send + 'static,
48    RequestBody: Send + 'static,
49    ResponseBody: Send + 'static,
50{
51    type Response = Response<ResponseBody>;
52    type Error = S::Error;
53    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
54
55    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
56        self.inner.poll_ready(cx)
57    }
58
59    fn call(&mut self, request: Request<RequestBody>) -> Self::Future {
60        let request_id = request.headers().get(&REQUEST_ID_HEADER).cloned();
61        let future = self.inner.call(request);
62        Box::pin(async move {
63            let mut response = future.await?;
64            response
65                .headers_mut()
66                .entry(&REQUEST_ID_HEADER)
67                .or_insert_with(|| request_id.unwrap_or_else(new_request_id));
68            Ok(response)
69        })
70    }
71}
72
73static REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
74
75fn new_request_id() -> HeaderValue {
76    HeaderValue::from_str(&Uuid::new_v4().to_string())
77        .expect("generated request id contains only valid header characters")
78}
79
80/// Request ID validation behavior for inbound `x-request-id` values.
81///
82/// Valid inbound IDs must parse as UUID v4 values. Invalid header syntax,
83/// non-UUID strings, and UUIDs from other versions are treated as malformed.
84/// Missing IDs are never rejected; the configured generator is used instead.
85#[derive(Clone, Copy, Debug, Eq, PartialEq)]
86pub enum RequestIdMode {
87    /// Propagate valid UUID v4 IDs and replace malformed incoming IDs.
88    ///
89    /// This mode is useful in development and at integration boundaries where
90    /// clients may still send legacy IDs. A malformed inbound value is not
91    /// exposed to handlers; it is replaced with a generated ID before the
92    /// request reaches the inner service.
93    Permissive,
94    /// Propagate valid UUID v4 IDs and reject malformed incoming IDs.
95    ///
96    /// A malformed inbound value returns `400 Bad Request` with an
97    /// `invalid_request_id` JSON error body. The rejection response still
98    /// receives a generated request ID in the configured response header.
99    Strict,
100}
101
102/// Compatibility alias for naming request ID validation policy.
103pub type RequestIdPolicy = RequestIdMode;
104
105/// Typed configuration for validated request ID propagation.
106///
107/// The default production config uses the `x-request-id` header, strict inbound
108/// validation, and UUID v4 generation. Custom generators are accepted, but their
109/// output must be a valid HTTP header value because generated IDs are inserted
110/// into request headers, request extensions, and response headers. If a custom
111/// generator returns an invalid header value, the middleware returns a stable
112/// `500 Internal Server Error` response with code
113/// `invalid_generated_request_id` instead of panicking or calling the inner
114/// service.
115#[derive(Clone)]
116pub struct RequestIdConfig {
117    header_name: HeaderName,
118    mode: RequestIdMode,
119    generator: Arc<dyn Fn() -> String + Send + Sync>,
120}
121
122impl std::fmt::Debug for RequestIdConfig {
123    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        formatter
125            .debug_struct("RequestIdConfig")
126            .field("header_name", &self.header_name)
127            .field("mode", &self.mode)
128            .finish_non_exhaustive()
129    }
130}
131
132impl RequestIdConfig {
133    /// Creates a production request ID policy.
134    ///
135    /// Defaults:
136    /// - header: `x-request-id`
137    /// - inbound validation: [`RequestIdMode::Strict`]
138    /// - generated IDs: UUID v4 strings
139    ///
140    /// In strict mode, a present but malformed inbound ID is rejected with
141    /// `400 Bad Request`. Missing IDs are generated and accepted.
142    pub fn production() -> Self {
143        Self {
144            header_name: REQUEST_ID_HEADER.clone(),
145            mode: RequestIdMode::Strict,
146            generator: Arc::new(|| Uuid::new_v4().to_string()),
147        }
148    }
149
150    /// Creates a development request ID policy.
151    ///
152    /// Development uses the same `x-request-id` header and UUID v4 generator as
153    /// production, but switches to [`RequestIdMode::Permissive`]. Valid inbound
154    /// UUID v4 IDs are propagated; malformed inbound IDs are replaced with
155    /// generated UUID v4 IDs instead of returning `400`.
156    pub fn development() -> Self {
157        Self::production().mode(RequestIdMode::Permissive)
158    }
159
160    /// Sets the request ID header name.
161    pub fn header_name(mut self, header_name: HeaderName) -> Self {
162        self.header_name = header_name;
163        self
164    }
165
166    /// Sets request ID validation behavior for present inbound IDs.
167    pub fn mode(mut self, mode: RequestIdMode) -> Self {
168        self.mode = mode;
169        self
170    }
171
172    /// Replaces the request ID generator.
173    ///
174    /// [`RequestIdConfig::production`] and [`RequestIdConfig::development`] use
175    /// UUID v4 strings. If you provide a custom generator, keep it deterministic
176    /// enough for your tests and ensure it returns values that can be stored in
177    /// an HTTP header. Invalid generated header values return a structured
178    /// framework error response before the request reaches the inner service.
179    pub fn generator(mut self, generator: impl Fn() -> String + Send + Sync + 'static) -> Self {
180        self.generator = Arc::new(generator);
181        self
182    }
183
184    /// Returns the configured request ID header name.
185    pub fn header(&self) -> &HeaderName {
186        &self.header_name
187    }
188
189    /// Returns the configured validation mode.
190    pub const fn validation_mode(&self) -> RequestIdMode {
191        self.mode
192    }
193
194    fn generate(&self) -> String {
195        (self.generator)()
196    }
197}
198
199impl Default for RequestIdConfig {
200    fn default() -> Self {
201        Self::production()
202    }
203}
204
205/// Creates a validated request ID layer.
206///
207/// The layer validates or generates a request ID before the inner service runs,
208/// inserts that value into the configured request header, stores a
209/// [`RequestContext`] in request extensions, and mirrors the same header onto
210/// the response when the inner service has not already set it.
211///
212/// ```
213/// use axum::{Router, routing::get};
214/// use nidus_http::middleware::{
215///     RequestIdConfig, RequestIdMode, validated_request_id_layer,
216/// };
217/// # async fn handler() -> &'static str { "user" }
218///
219/// let app = Router::new()
220///     .route("/users/{id}", get(handler))
221///     .layer(validated_request_id_layer(
222///         RequestIdConfig::production().mode(RequestIdMode::Strict),
223///     ));
224/// # let _: Router = app;
225/// ```
226pub fn validated_request_id_layer(config: RequestIdConfig) -> ValidatedRequestIdLayer {
227    ValidatedRequestIdLayer::new(config)
228}
229
230/// Tower layer that validates, generates, stores, and propagates request IDs.
231///
232/// Valid inbound request IDs are UUID v4 strings. With
233/// [`RequestIdMode::Strict`], malformed inbound IDs receive `400 Bad Request`.
234/// With [`RequestIdMode::Permissive`], malformed inbound IDs are replaced with a
235/// generated ID. Generated IDs are UUID v4 by default.
236///
237/// On accepted requests, the final ID is inserted into the configured request
238/// header, added to request extensions through [`RequestContext`], and copied to
239/// the response header if the inner service did not set one. Use
240/// [`crate::middleware::request_context_layer`] after this layer when you want
241/// the context enriched with route and correlation fields before handlers run.
242#[derive(Clone, Debug)]
243pub struct ValidatedRequestIdLayer {
244    config: RequestIdConfig,
245}
246
247impl ValidatedRequestIdLayer {
248    /// Creates a validated request ID layer from typed config.
249    pub fn new(config: RequestIdConfig) -> Self {
250        Self { config }
251    }
252}
253
254impl<S> Layer<S> for ValidatedRequestIdLayer {
255    type Service = ValidatedRequestIdService<S>;
256
257    fn layer(&self, inner: S) -> Self::Service {
258        ValidatedRequestIdService {
259            inner,
260            config: self.config.clone(),
261        }
262    }
263}
264
265/// Service produced by [`ValidatedRequestIdLayer`].
266#[derive(Clone, Debug)]
267pub struct ValidatedRequestIdService<S> {
268    inner: S,
269    config: RequestIdConfig,
270}
271
272impl<S> Service<Request<Body>> for ValidatedRequestIdService<S>
273where
274    S: Service<Request<Body>, Response = Response<Body>> + Send + 'static,
275    S::Future: Send + 'static,
276    S::Error: Send + 'static,
277{
278    type Response = Response<Body>;
279    type Error = S::Error;
280    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
281
282    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
283        self.inner.poll_ready(cx)
284    }
285
286    fn call(&mut self, request: Request<Body>) -> Self::Future {
287        let config = &self.config;
288        let (mut parts, body) = request.into_parts();
289
290        enum Inbound {
291            Valid(String, HeaderValue),
292            Invalid,
293            Missing,
294        }
295
296        let inbound = match parts.headers.get(config.header()) {
297            None => Inbound::Missing,
298            Some(value) => match value.to_str() {
299                // Reuse the inbound header value instead of re-parsing the
300                // validated string into a fresh HeaderValue.
301                Ok(text) if is_valid_request_id(text) => {
302                    Inbound::Valid(text.to_owned(), value.clone())
303                }
304                _ => Inbound::Invalid,
305            },
306        };
307
308        let (request_id, header_value) = match inbound {
309            Inbound::Valid(request_id, header_value) => (request_id, header_value),
310            Inbound::Invalid if config.validation_mode() == RequestIdMode::Strict => {
311                let (request_id, header_value) = match generated_request_id_header(config) {
312                    Some(generated) => generated,
313                    None => {
314                        return Box::pin(async move {
315                            Ok(invalid_generated_request_id_response(parts.uri.path()))
316                        });
317                    }
318                };
319                let mut response = invalid_request_id_response(&request_id, parts.uri.path());
320                response
321                    .headers_mut()
322                    .insert(config.header().clone(), header_value);
323                return Box::pin(async move { Ok(response) });
324            }
325            Inbound::Invalid | Inbound::Missing => match generated_request_id_header(config) {
326                Some(generated) => generated,
327                None => {
328                    return Box::pin(async move {
329                        Ok(invalid_generated_request_id_response(parts.uri.path()))
330                    });
331                }
332            },
333        };
334
335        let response_header = config.header().clone();
336        parts
337            .headers
338            .insert(response_header.clone(), header_value.clone());
339        let context = RequestContext::from_parts(&parts, request_id);
340        parts.extensions.insert(context);
341        let future = self.inner.call(Request::from_parts(parts, body));
342
343        Box::pin(async move {
344            let mut response = future.await?;
345            response
346                .headers_mut()
347                .entry(response_header)
348                .or_insert(header_value);
349            Ok(response)
350        })
351    }
352}
353
354fn generated_request_id_header(config: &RequestIdConfig) -> Option<(String, HeaderValue)> {
355    let request_id = config.generate();
356    let header_value = HeaderValue::from_str(&request_id).ok()?;
357    Some((request_id, header_value))
358}
359
360fn is_valid_request_id(value: &str) -> bool {
361    Uuid::parse_str(value)
362        .ok()
363        .and_then(|uuid| uuid.get_version())
364        == Some(Version::Random)
365}
366
367fn invalid_request_id_response(request_id: &str, path: &str) -> Response<Body> {
368    let timestamp = crate::error::timestamp_now();
369    (
370        StatusCode::BAD_REQUEST,
371        axum::Json(RequestIdErrorBody {
372            error: RequestIdErrorDetails {
373                status_code: StatusCode::BAD_REQUEST.as_u16(),
374                code: "invalid_request_id",
375                message: "invalid request id",
376                details: serde_json::Value::Null,
377                timestamp,
378                path: path.to_owned(),
379                request_id: Some(request_id.to_owned()),
380            },
381        }),
382    )
383        .into_response()
384}
385
386fn invalid_generated_request_id_response(path: &str) -> Response<Body> {
387    let timestamp = crate::error::timestamp_now();
388    (
389        StatusCode::INTERNAL_SERVER_ERROR,
390        axum::Json(RequestIdErrorBody {
391            error: RequestIdErrorDetails {
392                status_code: StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
393                code: "invalid_generated_request_id",
394                message: "generated request id was not a valid HTTP header value",
395                details: serde_json::Value::Null,
396                timestamp,
397                path: path.to_owned(),
398                request_id: None,
399            },
400        }),
401    )
402        .into_response()
403}
404
405#[derive(Debug, Serialize)]
406struct RequestIdErrorBody {
407    error: RequestIdErrorDetails,
408}
409
410#[derive(Debug, Serialize)]
411#[serde(rename_all = "camelCase")]
412struct RequestIdErrorDetails {
413    status_code: u16,
414    code: &'static str,
415    message: &'static str,
416    details: serde_json::Value,
417    timestamp: String,
418    path: String,
419    #[serde(skip_serializing_if = "Option::is_none")]
420    request_id: Option<String>,
421}
422
423#[cfg(test)]
424mod tests {
425    use std::{convert::Infallible, sync::Arc};
426
427    use tower::{Service, service_fn};
428
429    use super::*;
430
431    #[tokio::test]
432    async fn service_does_not_retain_an_extra_generator_owner_across_the_inner_future() {
433        let config = RequestIdConfig::production();
434        let generator = Arc::clone(&config.generator);
435        let mut service = ValidatedRequestIdService {
436            inner: service_fn(|_request: Request<Body>| async {
437                Ok::<_, Infallible>(Response::new(Body::empty()))
438            }),
439            config,
440        };
441        assert_eq!(Arc::strong_count(&generator), 2);
442
443        let future = service.call(Request::new(Body::empty()));
444        assert_eq!(Arc::strong_count(&generator), 2);
445        let response = future.await.unwrap();
446
447        assert_eq!(response.status(), StatusCode::OK);
448        assert_eq!(Arc::strong_count(&generator), 2);
449    }
450}