salvo_extra 0.95.0

//! Request id middleware.
//!
//! # Example
//!
//! ```no_run
//! use salvo_core::prelude::*;
//! use salvo_extra::request_id::RequestId;
//!
//! #[handler]
//! async fn hello(req: &mut Request) -> String {
//!     format!("Request id: {:?}", req.header::<String>("x-request-id"))
//! }
//!
//! #[tokio::main]
//! async fn main() {
//!     let acceptor = TcpListener::new("0.0.0.0:8698").bind().await;
//!     let router = Router::new().hoop(RequestId::new()).get(hello);
//!     Server::new(acceptor).serve(router).await;
//! }
//! ```
use std::fmt::{self, Debug, Formatter};
use tracing::Instrument;
use ulid::Ulid;

use salvo_core::http::{HeaderValue, Request, Response, header::HeaderName};
use salvo_core::{Depot, FlowCtrl, Handler, async_trait};

/// Key used to store the request id in the depot.
pub const REQUEST_ID_KEY: &str = "::salvo::request_id";

/// Extension trait that exposes the current request id stored on the [`Depot`].
pub trait RequestIdDepotExt {
    /// Get a reference to the request id from the depot, if one has been set.
    fn request_id(&self) -> Option<&str>;
}

impl RequestIdDepotExt for Depot {
    #[inline]
    fn request_id(&self) -> Option<&str> {
        self.get::<String>(REQUEST_ID_KEY).map(|v| &**v).ok()
    }
}

/// Middleware that assigns a request id to every incoming request.
#[non_exhaustive]
pub struct RequestId {
    /// The header name used to carry the request id.
    pub header_name: HeaderName,
    /// Whether to overwrite an existing request id. Default is `true`.
    ///
    /// When set to `false`, a client-supplied id in the request header is trusted
    /// and propagated to logs, the response and the depot. Only disable this
    /// behind a trusted proxy that sets/sanitizes the header, otherwise a client
    /// can inject arbitrary ids (log forging / trace-correlation confusion).
    pub overwrite: bool,
    /// The generator used to produce request ids.
    pub generator: Box<dyn IdGenerator + Send + Sync>,
}

impl Debug for RequestId {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("RequestId")
            .field("header_name", &self.header_name)
            .field("overwrite", &self.overwrite)
            .finish()
    }
}

impl RequestId {
    /// Create a new `RequestId` middleware.
    #[must_use]
    pub fn new() -> Self {
        Self {
            header_name: HeaderName::from_static("x-request-id"),
            overwrite: true,
            generator: Box::new(UlidGenerator::new()),
        }
    }

    /// Set the header name used to carry the request id.
    #[must_use]
    pub fn header_name(mut self, name: HeaderName) -> Self {
        self.header_name = name;
        self
    }

    /// Set whether to overwrite an existing request id. Default is `true`.
    ///
    /// # Security
    ///
    /// With `false`, a client-supplied id is trusted and forwarded to logs, the
    /// response and the depot. Only disable overwriting when a trusted proxy
    /// sets/sanitizes the header; otherwise clients can inject arbitrary ids.
    #[must_use]
    pub fn overwrite(mut self, overwrite: bool) -> Self {
        self.overwrite = overwrite;
        self
    }

    /// Set the generator used to produce request ids.
    #[must_use]
    pub fn generator(mut self, generator: impl IdGenerator + Send + Sync + 'static) -> Self {
        self.generator = Box::new(generator);
        self
    }

    fn generate_id(&self, req: &mut Request, depot: &mut Depot) -> HeaderValue {
        let id = self.generator.generate(req, depot);
        match HeaderValue::from_str(&id) {
            Ok(header_value) => header_value,
            Err(error) => {
                tracing::warn!(
                    error = ?error,
                    generated_id = %id,
                    "request id generator returned an invalid header value; falling back to ULID"
                );
                HeaderValue::from_str(&Ulid::r#gen().to_string())
                    .expect("ULID should always be a valid header value")
            }
        }
    }
}

impl Default for RequestId {
    fn default() -> Self {
        Self::new()
    }
}

/// Trait for types that can generate a new request id.
pub trait IdGenerator {
    /// Generate a new request id.
    fn generate(&self, req: &mut Request, depot: &mut Depot) -> String;
}

impl<F> IdGenerator for F
where
    F: Fn() -> String + Send + Sync,
{
    fn generate(&self, _req: &mut Request, _depot: &mut Depot) -> String {
        self()
    }
}

/// An [`IdGenerator`] that produces request ids using ULIDs.
#[derive(Default, Debug)]
pub struct UlidGenerator {}
impl UlidGenerator {
    /// Create a new `UlidGenerator`.
    #[must_use]
    pub fn new() -> Self {
        Self {}
    }
}
impl IdGenerator for UlidGenerator {
    fn generate(&self, _req: &mut Request, _depot: &mut Depot) -> String {
        Ulid::r#gen().to_string()
    }
}

#[async_trait]
impl Handler for RequestId {
    async fn handle(
        &self,
        req: &mut Request,
        depot: &mut Depot,
        res: &mut Response,
        ctrl: &mut FlowCtrl,
    ) {
        let request_id = match req.headers().get(&self.header_name) {
            None => self.generate_id(req, depot),
            Some(value) => {
                if self.overwrite {
                    self.generate_id(req, depot)
                } else {
                    value.clone()
                }
            }
        };

        // Overwrite (not append) so the request carries exactly the resolved id:
        // appending left a client-supplied id in place when `overwrite` generated a
        // new one (downstream `req.header(..)` then read the stale value), and
        // duplicated the header in the pass-through case.
        let _ = req.add_header(self.header_name.clone(), &request_id, true);

        let span = tracing::info_span!("request", ?request_id);
        res.headers_mut()
            .insert(self.header_name.clone(), request_id.clone());
        // Store the id as a `String` so `RequestIdDepotExt::request_id()`, which
        // looks it up by the `String` type, can actually retrieve it. Inserting the
        // raw `HeaderValue` made that accessor always return `None`.
        if let Ok(id) = request_id.to_str() {
            depot.insert(REQUEST_ID_KEY, id.to_owned());
        }

        async move {
            ctrl.call_next(req, depot, res).await;
        }
        .instrument(span)
        .await;
    }
}
#[cfg(test)]
mod tests {
    use salvo_core::prelude::*;
    use salvo_core::test::{ResponseExt, TestClient};

    use super::*;

    #[tokio::test]
    async fn test_request_id_added() {
        let handler = RequestId::new();
        let router = Router::new().hoop(handler).get(endpoint);
        let service = Service::new(router);

        let response = TestClient::get("http://127.0.0.1:8698/")
            .send(&service)
            .await;
        assert_eq!(response.status_code, Some(StatusCode::OK));
        assert!(response.headers.contains_key("x-request-id"));
    }

    #[tokio::test]
    async fn test_request_id_overwrite() {
        let handler = RequestId::new().overwrite(true);
        let router = Router::new().hoop(handler).get(endpoint);
        let service = Service::new(router);

        let response = TestClient::get("http://127.0.0.1:8698/")
            .add_header("x-request-id", "existing-id", true)
            .send(&service)
            .await;
        assert_eq!(response.status_code, Some(StatusCode::OK));
        assert_ne!(response.headers.get("x-request-id").unwrap(), "existing-id");
    }

    #[tokio::test]
    async fn test_request_id_no_overwrite() {
        let handler = RequestId::new().overwrite(false);
        let router = Router::new().hoop(handler).get(endpoint);
        let service = Service::new(router);

        let response = TestClient::get("http://127.0.0.1:8698/")
            .add_header("x-request-id", "existing-id", true)
            .send(&service)
            .await;
        assert_eq!(response.status_code, Some(StatusCode::OK));
        assert_eq!(response.headers.get("x-request-id").unwrap(), "existing-id");
    }

    #[tokio::test]
    async fn test_custom_generator() {
        let handler = RequestId::new().generator(|| "custom-id".to_owned());
        let router = Router::new().hoop(handler).get(endpoint);
        let service = Service::new(router);

        let response = TestClient::get("http://127.0.0.1:8698/")
            .send(&service)
            .await;
        assert_eq!(response.status_code, Some(StatusCode::OK));
        assert_eq!(response.headers.get("x-request-id").unwrap(), "custom-id");
    }

    #[tokio::test]
    async fn test_invalid_custom_generator_falls_back() {
        let handler = RequestId::new().generator(|| "bad\r\nvalue".to_owned());
        let router = Router::new().hoop(handler).get(endpoint);
        let service = Service::new(router);

        let response = TestClient::get("http://127.0.0.1:8698/")
            .send(&service)
            .await;
        assert_eq!(response.status_code, Some(StatusCode::OK));
        let request_id = response
            .headers
            .get("x-request-id")
            .unwrap()
            .to_str()
            .unwrap();
        assert_ne!(request_id, "bad\r\nvalue");
        assert_eq!(request_id.len(), 26);
    }

    #[tokio::test]
    async fn test_depot_storage() {
        let handler = RequestId::new();
        #[handler]
        async fn depot_checker(depot: &mut Depot, res: &mut Response) {
            // Exercise the public accessor, which reads the id back as a `String`.
            let id = depot
                .request_id()
                .expect("request id should be retrievable via RequestIdDepotExt")
                .to_owned();
            res.render(Text::Plain(id));
        }
        let router = Router::new().hoop(handler).get(depot_checker);
        let service = Service::new(router);

        let mut response = TestClient::get("http://127.0.0.1:8698/")
            .send(&service)
            .await;
        assert_eq!(response.status_code, Some(StatusCode::OK));
        let header_id = response
            .headers
            .get("x-request-id")
            .unwrap()
            .to_str()
            .unwrap().to_owned();
        let body = response.take_string().await.unwrap();
        assert_eq!(header_id, body);
    }

    #[handler]
    async fn endpoint() {}
}