Skip to main content

typeway_server/
request_id.rs

1//! Request ID middleware — assigns a unique ID to each request.
2//!
3//! The ID is available via the [`RequestId`] extractor and is also
4//! set as the `x-request-id` response header.
5//!
6//! # Example
7//!
8//! ```ignore
9//! use typeway_server::request_id::{RequestId, RequestIdLayer};
10//!
11//! Server::<API>::new(handlers)
12//!     .layer(RequestIdLayer::new())
13//!     .serve(addr)
14//!     .await?;
15//!
16//! // In handlers:
17//! async fn handler(Extension(id): Extension<RequestId>) -> String {
18//!     format!("request: {}", id.0)
19//! }
20//! ```
21
22use std::convert::Infallible;
23use std::future::Future;
24use std::pin::Pin;
25use std::task::{Context, Poll};
26
27use crate::body::BoxBody;
28
29/// A unique request identifier.
30#[derive(Debug, Clone)]
31pub struct RequestId(pub String);
32
33/// Tower layer that assigns a unique `x-request-id` to each request.
34///
35/// - Generates a UUID v4 if no `x-request-id` header is present
36/// - Preserves an existing `x-request-id` header if present
37/// - Injects [`RequestId`] into request extensions (accessible via `Extension<RequestId>`)
38/// - Copies the ID to the response `x-request-id` header
39#[derive(Clone)]
40pub struct RequestIdLayer;
41
42impl RequestIdLayer {
43    pub fn new() -> Self {
44        RequestIdLayer
45    }
46}
47
48impl Default for RequestIdLayer {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl<S> tower_layer::Layer<S> for RequestIdLayer {
55    type Service = RequestIdService<S>;
56
57    fn layer(&self, inner: S) -> Self::Service {
58        RequestIdService { inner }
59    }
60}
61
62/// The service produced by [`RequestIdLayer`].
63#[derive(Clone)]
64pub struct RequestIdService<S> {
65    inner: S,
66}
67
68impl<S, B> tower_service::Service<http::Request<B>> for RequestIdService<S>
69where
70    S: tower_service::Service<
71            http::Request<B>,
72            Response = http::Response<BoxBody>,
73            Error = Infallible,
74        > + Clone
75        + Send
76        + 'static,
77    S::Future: Send + 'static,
78    B: Send + 'static,
79{
80    type Response = http::Response<BoxBody>;
81    type Error = Infallible;
82    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
83
84    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
85        self.inner.poll_ready(cx)
86    }
87
88    fn call(&mut self, mut req: http::Request<B>) -> Self::Future {
89        // Extract or generate request ID.
90        let id = req
91            .headers()
92            .get("x-request-id")
93            .and_then(|v| v.to_str().ok())
94            .map(|s| s.to_string())
95            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
96
97        // Inject into extensions for handler access.
98        req.extensions_mut().insert(RequestId(id.clone()));
99
100        let mut inner = self.inner.clone();
101        Box::pin(async move {
102            let mut resp = inner.call(req).await?;
103            // Set response header.
104            if let Ok(val) = http::HeaderValue::from_str(&id) {
105                resp.headers_mut().insert("x-request-id", val);
106            }
107            Ok(resp)
108        })
109    }
110}