b453d5d
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use axum::body::Body;
use axum::http::{HeaderName, HeaderValue, Request, Response, header};
use bytes::Bytes;
use http_body_util::BodyExt;
use serde_json::Value;
use tower::{Layer, Service};
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;
}
#[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()
}
/// 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 {
new_id()
}
#[derive(Debug, Clone, Copy, Default)]
pub struct AppTraceLayer;
impl AppTraceLayer {
pub fn new() -> Self {
Self
}
}
impl<S> Layer<S> for AppTraceLayer {
type Service = AppTraceService<S>;
fn layer(&self, inner: S) -> Self::Service {
AppTraceService { inner }
}
}
#[derive(Debug, Clone)]
pub struct AppTraceService<S> {
inner: S,
}
impl<S> Service<Request<Body>> for AppTraceService<S>
where
S: Service<Request<Body>, Response = Response<Body>> + Send + 'static,
S::Future: Send + 'static,
{
type Response = S::Response;
type Error = S::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut req: Request<Body>) -> Self::Future {
let request_id = inbound_id(req.headers().get(&REQUEST_ID_HEADER)).unwrap_or_else(new_id);
let log_id = new_id();
req.extensions_mut().insert(RequestId(request_id.clone()));
req.extensions_mut().insert(LogId(log_id.clone()));
let fut = REQUEST_ID.scope(request_id.clone(), self.inner.call(req));
Box::pin(async move {
let mut res = fut.await?;
if let Ok(value) = HeaderValue::try_from(request_id.as_str()) {
res.headers_mut().insert(REQUEST_ID_HEADER, value);
}
if let Ok(value) = HeaderValue::try_from(log_id.as_str()) {
res.headers_mut().insert(LOG_ID_HEADER, value);
}
inject_ids(&mut res, &request_id, &log_id).await;
Ok(res)
})
}
}
fn new_id() -> String {
Uuid::now_v7().to_string()
}
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())
}
async fn inject_ids(res: &mut Response<Body>, request_id: &str, log_id: &str) {
let is_json = res
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.starts_with("application/json"));
if !is_json {
return;
}
let (mut parts, body) = std::mem::take(res).into_parts();
let bytes = match body.collect().await {
Ok(collected) => collected.to_bytes(),
Err(_) => {
*res = Response::from_parts(parts, Body::empty());
return;
}
};
let output = match serde_json::from_slice::<Value>(&bytes) {
Ok(Value::Object(mut map))
if map.contains_key("request_id") && map.contains_key("log_id") =>
{
map.insert(
"request_id".to_string(),
Value::String(request_id.to_string()),
);
map.insert("log_id".to_string(), Value::String(log_id.to_string()));
match serde_json::to_vec(&map) {
Ok(rewritten) => Bytes::from(rewritten),
Err(_) => bytes,
}
}
_ => bytes,
};
if let Ok(length) = HeaderValue::try_from(output.len().to_string()) {
parts.headers.insert(header::CONTENT_LENGTH, length);
}
*res = Response::from_parts(parts, Body::from(output));
}