oxidite_middleware/
request_id.rs1use 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#[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 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 let mut response = inner.call(req).await?;
49
50 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#[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}