use crate::{
core::router::Handler,
core::{Request, Response},
middleware::Middleware,
};
use std::sync::Arc;
#[derive(Clone)]
pub struct RequestId {
header: &'static str,
}
impl RequestId {
pub fn new() -> Self {
Self {
header: "x-request-id",
}
}
}
impl Default for RequestId {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl Middleware for RequestId {
async fn handle(&self, mut req: Request, next: Arc<dyn Handler>) -> Response {
let request_id = req
.headers()
.get(self.header)
.and_then(|v| v.to_str().ok())
.filter(|s| !s.is_empty())
.map(|s| s.to_string())
.unwrap_or_else(crate::utils::request_id::generate);
req.headers_mut()
.insert(self.header, request_id.clone().try_into().unwrap());
let mut res = next.handle(req).await;
if !res.headers.contains_key(self.header) {
let _ = res.headers.insert(
self.header,
http::HeaderValue::from_str(&request_id).unwrap(),
);
}
res
}
}