Skip to main content

oxidite_middleware/
request_id.rs

1use oxidite_core::{OxiditeRequest, OxiditeResponse,  Error as CoreError};
2use tower::{Service, Layer};
3use std::task::{Context, Poll};
4use std::future::Future;
5use std::pin::Pin;
6use uuid::Uuid;
7
8const REQUEST_ID_HEADER: &str = "x-request-id";
9
10/// Request ID middleware
11#[derive(Clone)]
12pub struct RequestIdMiddleware<S> {
13    inner: S,
14}
15
16impl<S> RequestIdMiddleware<S> {
17    pub fn new(inner: S) -> Self {
18        Self { inner }
19    }
20}
21
22impl<S> Service<OxiditeRequest> for RequestIdMiddleware<S>
23where
24    S: Service<OxiditeRequest, Response = OxiditeResponse, Error = CoreError> + Clone + Send + 'static,
25    S::Future: Send + 'static,
26{
27    type Response = S::Response;
28    type Error = S::Error;
29    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
30
31    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
32        self.inner.poll_ready(cx)
33    }
34
35    fn call(&mut self, req: OxiditeRequest) -> Self::Future {
36        // Extract or generate request ID
37        let request_id = req
38            .headers()
39            .get(REQUEST_ID_HEADER)
40            .and_then(|h| h.to_str().ok())
41            .map(|s| s.to_string())
42            .unwrap_or_else(|| Uuid::new_v4().to_string());
43
44        let mut inner = self.inner.clone();
45
46        Box::pin(async move {
47            // TODO: Attach request_id to request extensions
48            let mut response = inner.call(req).await?;
49
50            // Add request ID to response headers
51            if let Ok(header_value) = request_id.parse() {
52                response.headers_mut().insert(REQUEST_ID_HEADER, header_value);
53            }
54
55            Ok(response)
56        })
57    }
58}
59
60/// Layer for Request ID middleware
61#[derive(Clone, Default)]
62pub struct RequestIdLayer;
63
64impl RequestIdLayer {
65    pub fn new() -> Self {
66        Self
67    }
68}
69
70impl<S> Layer<S> for RequestIdLayer {
71    type Service = RequestIdMiddleware<S>;
72
73    fn layer(&self, inner: S) -> Self::Service {
74        RequestIdMiddleware::new(inner)
75    }
76}