b7cd2b1
use std::future::Future;
use actix_web::body::{BoxBody, MessageBody};
use actix_web::dev::{ServiceRequest, ServiceResponse};
use actix_web::http::header::{HeaderName, HeaderValue};
use actix_web::middleware::Next;
use actix_web::{Error, HttpMessage};
use uuid::Uuid;
pub const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
pub const LOG_ID_HEADER: HeaderName = HeaderName::from_static("x-log-id");
const MAX_INBOUND_ID_LEN: usize = 128;
tokio::task_local! {
static REQUEST_ID: String;
static LOG_ID: String;
}
#[derive(Debug, Clone)]
pub struct RequestId(pub String);
#[derive(Debug, Clone)]
pub struct LogId(pub String);
/// The request id of the request currently being handled, when one is set.
pub fn request_id() -> Option<String> {
REQUEST_ID.try_with(|id| id.clone()).ok()
}
/// The log id of the request currently being handled, when one is set.
pub fn log_id() -> Option<String> {
LOG_ID.try_with(|id| id.clone()).ok()
}
/// Runs `future` with `request_id` available to nested code, so queue
/// consumers keep the id of the request that produced the message.
pub async fn with_request_id<F: Future>(request_id: Option<&str>, future: F) -> F::Output {
match request_id {
Some(request_id) => REQUEST_ID.scope(request_id.to_string(), future).await,
None => future.await,
}
}
pub fn new_request_id() -> String {
Uuid::now_v7().to_string()
}
pub async fn trace_request<B>(
request: ServiceRequest,
next: Next<B>,
) -> Result<ServiceResponse<BoxBody>, Error>
where
B: MessageBody + 'static,
{
let request_id =
inbound_id(request.headers().get(&REQUEST_ID_HEADER)).unwrap_or_else(new_request_id);
let log_id = new_request_id();
request
.extensions_mut()
.insert(RequestId(request_id.clone()));
request.extensions_mut().insert(LogId(log_id.clone()));
let response = REQUEST_ID
.scope(
request_id.clone(),
LOG_ID.scope(log_id.clone(), next.call(request)),
)
.await?;
let mut response = response.map_into_boxed_body();
if let Ok(value) = HeaderValue::try_from(request_id.as_str()) {
response.headers_mut().insert(REQUEST_ID_HEADER, value);
}
if let Ok(value) = HeaderValue::try_from(log_id.as_str()) {
response.headers_mut().insert(LOG_ID_HEADER, value);
}
Ok(response)
}
fn inbound_id(value: Option<&HeaderValue>) -> Option<String> {
let value = value?.to_str().ok()?.trim();
if value.is_empty() || value.len() > MAX_INBOUND_ID_LEN {
return None;
}
Some(value.to_string())
}