tako_rs_plugins/middleware/
request_id.rs

1//! Request ID middleware for tracing and correlation.
2//!
3//! Generates or propagates a unique request identifier via the `X-Request-ID` header.
4//! If the incoming request already has the header, it is preserved; otherwise a new
5//! UUID v4 is generated. The ID is injected into both request extensions and
6//! the response header.
7
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12use http::HeaderName;
13use http::HeaderValue;
14use tako_rs_core::middleware::IntoMiddleware;
15use tako_rs_core::middleware::Next;
16use tako_rs_core::types::Request;
17use tako_rs_core::types::Response;
18
19/// A request ID value that can be extracted from request extensions.
20#[derive(Debug, Clone)]
21pub struct RequestIdValue(pub String);
22
23/// Request ID middleware configuration.
24///
25/// # Examples
26///
27/// ```rust
28/// use tako::middleware::request_id::RequestId;
29/// use tako::middleware::IntoMiddleware;
30///
31/// // Default: uses X-Request-ID header with UUID v4
32/// let mw = RequestId::new().into_middleware();
33///
34/// // Custom header name
35/// let mw = RequestId::new().header_name("X-Correlation-ID").into_middleware();
36/// ```
37pub struct RequestId {
38  header: HeaderName,
39  generator: Arc<dyn Fn() -> String + Send + Sync + 'static>,
40}
41
42impl Default for RequestId {
43  fn default() -> Self {
44    Self::new()
45  }
46}
47
48impl RequestId {
49  /// Creates a new `RequestId` middleware with default settings (X-Request-ID, UUID v4).
50  pub fn new() -> Self {
51    Self {
52      header: HeaderName::from_static("x-request-id"),
53      generator: Arc::new(|| uuid::Uuid::new_v4().to_string()),
54    }
55  }
56
57  /// Sets a custom header name for the request ID.
58  pub fn header_name(mut self, name: &'static str) -> Self {
59    self.header = HeaderName::from_static(name);
60    self
61  }
62
63  /// Sets a custom ID generator function.
64  pub fn generator(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self {
65    self.generator = Arc::new(f);
66    self
67  }
68}
69
70impl IntoMiddleware for RequestId {
71  fn into_middleware(
72    self,
73  ) -> impl Fn(Request, Next) -> Pin<Box<dyn Future<Output = Response> + Send + 'static>>
74  + Clone
75  + Send
76  + Sync
77  + 'static {
78    let header = self.header;
79    let generator = self.generator;
80
81    move |mut req: Request, next: Next| {
82      let header = header.clone();
83      let generator = generator.clone();
84
85      Box::pin(async move {
86        // Use existing request ID or generate a new one. Cap the inbound
87        // header to 256 bytes — without a limit a peer can pin an
88        // arbitrarily large `String` into the request extensions and the
89        // response headers, which both flows downstream (log lines, tracing
90        // spans, response header allocator). 256 matches the X-Request-ID
91        // hygiene most CDNs already enforce and is plenty for ULIDs, UUIDs,
92        // and traceparent fragments.
93        const MAX_INBOUND_LEN: usize = 256;
94        let id = req
95          .headers()
96          .get(&header)
97          .and_then(|v| v.to_str().ok())
98          .filter(|s| !s.is_empty() && s.len() <= MAX_INBOUND_LEN)
99          .map_or_else(|| generator(), std::string::ToString::to_string);
100
101        // Inject into request extensions for handler access
102        req.extensions_mut().insert(RequestIdValue(id.clone()));
103
104        let mut resp = next.run(req).await;
105
106        // Add to response headers
107        if let Ok(val) = HeaderValue::from_str(&id) {
108          resp.headers_mut().insert(header, val);
109        }
110
111        resp
112      })
113    }
114  }
115}
tako_rs_plugins/middleware/request_id.rs

tako_rs_plugins/middleware/
request_id.rs
