at-jet 0.7.2

//! Request ID middleware
//!
//! Generates or propagates a unique request ID for every HTTP request.
//! The ID is:
//! 1. Read from the incoming `X-Request-Id` header (if present), or generated as a UUID v4
//! 2. Stored in request extensions (accessible to handlers via `Extension<RequestId>`)
//! 3. Recorded on the current tracing span (appears in all log lines for the request)
//! 4. Set on the response `X-Request-Id` header (returned to the client)
//!
//! # Example
//!
//! ```rust,ignore
//! use at_jet::middleware::RequestIdLayer;
//!
//! let app = Router::new()
//!     .route("/api/users", get(handler))
//!     .layer(RequestIdLayer::new());
//! ```
//!
//! Access in handlers:
//!
//! ```rust,ignore
//! use at_jet::middleware::RequestId;
//! use axum::Extension;
//!
//! async fn handler(Extension(req_id): Extension<RequestId>) -> String {
//!     format!("Your request ID: {}", req_id.as_str())
//! }
//! ```

use {axum::http::{HeaderName,
                  HeaderValue,
                  Request},
     std::{future::Future,
           pin::Pin,
           task::{Context,
                  Poll}},
     tower::{Layer,
             Service},
     tracing::Span,
     uuid::Uuid};

/// Header name for request ID propagation.
pub static X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");

/// A unique identifier for a single HTTP request.
#[derive(Clone, Debug)]
pub struct RequestId(String);

impl RequestId {
  /// Get the request ID as a string slice.
  pub fn as_str(&self) -> &str {
    &self.0
  }
}

impl std::fmt::Display for RequestId {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.write_str(&self.0)
  }
}

/// Layer that adds request ID generation and propagation.
#[derive(Clone, Default)]
pub struct RequestIdLayer;

impl RequestIdLayer {
  /// Create a new `RequestIdLayer`.
  pub fn new() -> Self {
    Self
  }
}

impl<S> Layer<S> for RequestIdLayer {
  type Service = RequestIdMiddleware<S>;

  fn layer(&self, inner: S) -> Self::Service {
    RequestIdMiddleware { inner }
  }
}

/// Middleware service that generates or propagates request IDs.
#[derive(Clone)]
pub struct RequestIdMiddleware<S> {
  inner: S,
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestIdMiddleware<S>
where
  S: Service<Request<ReqBody>, Response = axum::response::Response<ResBody>> + Clone + Send + 'static,
  S::Future: Send,
  ReqBody: Send + 'static,
  ResBody: Send + 'static,
{
  type Error = S::Error;
  type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
  type Response = S::Response;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
    self.inner.poll_ready(cx)
  }

  fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
    // 1. Read from header or generate new UUID v4
    let id = req
      .headers()
      .get(&X_REQUEST_ID)
      .and_then(|v| v.to_str().ok())
      .filter(|v| !v.is_empty())
      .map(String::from)
      .unwrap_or_else(|| Uuid::new_v4().to_string());

    // 2. Record on the current tracing span
    Span::current().record("request_id", id.as_str());

    // 3. Build response header value eagerly (avoids cloning the id String)
    let header_val = HeaderValue::from_str(&id).ok();

    // 4. Store in request extensions (moves id, no clone needed)
    req.extensions_mut().insert(RequestId(id));

    let mut inner = self.inner.clone();

    Box::pin(async move {
      let mut response = inner.call(req).await?;

      // 5. Set on response header
      if let Some(val) = header_val {
        response.headers_mut().insert(&X_REQUEST_ID, val);
      }

      Ok(response)
    })
  }
}