use axum::{extract::Request, middleware::Next, response::Response};
fn sanitize_request_header(value: &str) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
if trimmed.len() > 128 {
return None;
}
Some(trimmed.to_string())
}
pub async fn request_context_middleware(mut request: Request, next: Next) -> Response {
let mut correlation_id = request
.headers()
.get("x-correlation-id")
.and_then(|v| v.to_str().ok())
.and_then(sanitize_request_header)
.or_else(|| {
request
.headers()
.get("x-trace-id")
.and_then(|v| v.to_str().ok())
.and_then(sanitize_request_header)
});
let request_id = request
.headers()
.get("x-request-id")
.and_then(|v| v.to_str().ok())
.and_then(sanitize_request_header)
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
if correlation_id.is_none() {
correlation_id = Some(request_id.clone());
}
if request.headers().get("x-request-id").is_none() {
if let Ok(value) = axum::http::HeaderValue::from_str(&request_id) {
request.headers_mut().insert("x-request-id", value);
}
}
if request.headers().get("x-correlation-id").is_none() {
if let Some(correlation_id) = correlation_id.as_deref() {
if let Ok(value) = axum::http::HeaderValue::from_str(correlation_id) {
request.headers_mut().insert("x-correlation-id", value);
}
}
}
let ctx = crate::modules::system::request_context::RequestContext {
correlation_id: correlation_id.clone(),
request_id: Some(request_id.clone()),
};
let mut response =
crate::modules::system::request_context::with_request_context(ctx, async move {
next.run(request).await
})
.await;
if let Ok(value) = axum::http::HeaderValue::from_str(&request_id) {
response.headers_mut().insert("x-request-id", value);
}
if let Some(correlation_id) = correlation_id.as_deref() {
if let Ok(value) = axum::http::HeaderValue::from_str(correlation_id) {
response.headers_mut().insert("x-correlation-id", value);
}
}
response
}