tork_core/middleware/
request_id.rs

1//! Request-identifier middleware.
2
3use http::{HeaderName, HeaderValue};
4
5use crate::error::Result;
6use crate::middleware::{DuplicatePolicy, Middleware, Next, Request};
7use crate::response::Response;
8use crate::router::BoxFuture;
9
10/// Default header carrying the request identifier.
11const DEFAULT_HEADER: &str = "x-request-id";
12/// Prefix applied to generated request identifiers.
13const REQUEST_ID_PREFIX: &str = "req-";
14
15/// Assigns a request identifier and echoes it on the response.
16///
17/// If the incoming request already carries the id header, its value is reused;
18/// otherwise a `req-<uuid>` value is generated. The id is written back onto the
19/// request (so downstream layers and handlers can read it) and onto the
20/// response.
21pub struct RequestId {
22    header: HeaderName,
23}
24
25impl RequestId {
26    /// Creates the middleware using the default `x-request-id` header.
27    pub fn new() -> Self {
28        Self {
29            header: HeaderName::from_static(DEFAULT_HEADER),
30        }
31    }
32
33    /// Sets the header name that carries the request id.
34    pub fn header_name(mut self, name: &'static str) -> Self {
35        self.header = HeaderName::from_static(name);
36        self
37    }
38}
39
40impl Default for RequestId {
41    fn default() -> Self {
42        Self::new()
43    }
44}
45
46impl Middleware for RequestId {
47    fn handle(&self, mut request: Request, next: Next) -> BoxFuture<'static, Result<Response>> {
48        let header = self.header.clone();
49
50        let id = request
51            .headers()
52            .get(&header)
53            .and_then(|value| value.to_str().ok())
54            .map(str::to_owned)
55            .unwrap_or_else(|| format!("{REQUEST_ID_PREFIX}{}", uuid::Uuid::new_v4()));
56
57        if let Ok(value) = HeaderValue::from_str(&id) {
58            request.headers_mut().insert(header.clone(), value);
59        }
60
61        Box::pin(async move {
62            let mut response = next.run(request).await?;
63            if let Ok(value) = HeaderValue::from_str(&id) {
64                response.headers_mut().insert(header, value);
65            }
66            Ok(response)
67        })
68    }
69
70    fn name(&self) -> &'static str {
71        "RequestId"
72    }
73
74    fn duplicate_policy(&self) -> DuplicatePolicy {
75        DuplicatePolicy::Reject
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn header_name_builder_replaces_default_header() {
85        let request_id = RequestId::new().header_name("x-correlation-id");
86        assert_eq!(
87            request_id.header,
88            HeaderName::from_static("x-correlation-id")
89        );
90    }
91}
tork_core/middleware/request_id.rs

tork_core/middleware/
request_id.rs
