nidus-http 1.0.15

//! Request context primitives shared by middleware, handlers, and observers.

use std::{
    future::Future,
    net::{IpAddr, SocketAddr},
    sync::Arc,
};

use axum::extract::FromRequestParts;
use http::{HeaderMap, Method, request::Parts};
use serde::Serialize;

/// Client classification inferred from request boundary headers.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ClientKind {
    /// Request uses an API key style credential.
    ApiKey,
    /// Request carries a bearer token or other authorization header.
    Authenticated,
    /// Request has no recognized application credential.
    Anonymous,
}

impl ClientKind {
    /// Returns the stable string label for this client kind.
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::ApiKey => "api_key",
            Self::Authenticated => "authenticated",
            Self::Anonymous => "anonymous",
        }
    }
}

/// Request/correlation context attached to request extensions.
///
/// `RequestContext` is available to handlers when the router uses
/// [`crate::middleware::validated_request_id_layer`] plus
/// [`crate::middleware::request_context_layer`], or when it is wrapped by
/// [`crate::middleware::ApiDefaults::production`]. Extracting it without those
/// extensions rejects the request with `500 Internal Server Error`.
///
/// Fields are inferred from request headers and Axum extensions:
/// - `request_id`: the final validated/generated `x-request-id`
/// - `correlation_id`: `x-correlation-id`, falling back to the request ID
/// - `trace_id` / `span_id`: validated IDs from `traceparent`
/// - `client_kind`: `x-api-key` means API key, otherwise `Authorization` means
///   authenticated, otherwise anonymous
/// - `route`: Axum's [`axum::extract::MatchedPath`] when it is available at the
///   point the context layer runs
///
/// Clones share the immutable string metadata. Consuming enrichment methods
/// such as [`Self::with_user_id`] use copy-on-write, so changing a clone does
/// not change the original context.
///
/// ```
/// use nidus_http::{Json, context::RequestContext};
///
/// async fn handler(context: RequestContext) -> Json<serde_json::Value> {
///     Json(serde_json::json!({
///         "requestId": context.request_id(),
///         "correlationId": context.correlation_id(),
///     }))
/// }
/// ```
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RequestContext {
    inner: Arc<RequestContextInner>,
    method: Method,
    client_kind: ClientKind,
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct RequestContextInner {
    request_id: String,
    correlation_id: CorrelationId,
    route: Option<String>,
    path: String,
    trace_id: Option<String>,
    span_id: Option<String>,
    user_id: Option<String>,
    tenant_id: Option<String>,
    session_id: Option<String>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum CorrelationId {
    None,
    RequestId,
    Explicit(String),
}

impl CorrelationId {
    fn as_str<'a>(&'a self, request_id: &'a str) -> Option<&'a str> {
        match self {
            Self::None => None,
            Self::RequestId => Some(request_id),
            Self::Explicit(value) => Some(value),
        }
    }
}

impl RequestContext {
    /// Creates a context for the current request boundary.
    ///
    /// This constructor is useful in tests or custom middleware. It does not
    /// inspect headers, so optional correlation, trace, route, and client fields
    /// remain empty/default until set explicitly or built via [`Self::from_parts`].
    pub fn new(request_id: impl Into<String>, method: Method, path: impl Into<String>) -> Self {
        Self {
            inner: Arc::new(RequestContextInner {
                request_id: request_id.into(),
                correlation_id: CorrelationId::None,
                route: None,
                path: path.into(),
                trace_id: None,
                span_id: None,
                user_id: None,
                tenant_id: None,
                session_id: None,
            }),
            method,
            client_kind: ClientKind::Anonymous,
        }
    }

    /// Creates a context from request parts.
    ///
    /// This reads `x-correlation-id`, `traceparent`, `x-api-key`,
    /// `Authorization`, and [`axum::extract::MatchedPath`] from the request
    /// boundary. The supplied `request_id` is expected to be the final ID chosen
    /// by request ID middleware.
    pub fn from_parts(parts: &Parts, request_id: impl Into<String>) -> Self {
        let request_id = request_id.into();
        let correlation_id = match header_to_string(&parts.headers, "x-correlation-id") {
            Some(value) => CorrelationId::Explicit(value),
            None if request_id.is_empty() => CorrelationId::None,
            None => CorrelationId::RequestId,
        };
        let mut context = Self::new(request_id, parts.method.clone(), parts.uri.path());
        let inner = Arc::make_mut(&mut context.inner);
        inner.correlation_id = correlation_id;
        inner.route = parts
            .extensions
            .get::<axum::extract::MatchedPath>()
            .map(|path| path.as_str().to_owned());
        context.client_kind = infer_client_kind(&parts.headers);
        if let Some(trace_context) = parts
            .headers
            .get("traceparent")
            .and_then(|value| value.to_str().ok())
            .and_then(parse_traceparent)
        {
            inner.trace_id = Some(trace_context.trace_id.to_owned());
            inner.span_id = Some(trace_context.parent_id.to_owned());
        }
        context
    }

    /// Returns the final request id.
    ///
    /// With [`crate::middleware::validated_request_id_layer`], this is either a
    /// valid inbound UUID v4 or a generated ID.
    pub fn request_id(&self) -> &str {
        &self.inner.request_id
    }

    pub(crate) fn into_request_id(self) -> String {
        match Arc::try_unwrap(self.inner) {
            Ok(inner) => inner.request_id,
            Err(inner) => inner.request_id.clone(),
        }
    }

    /// Returns the correlation id when available.
    ///
    /// [`Self::from_parts`] prefers `x-correlation-id` and falls back to the
    /// request ID when no correlation header is present.
    pub fn correlation_id(&self) -> Option<&str> {
        self.inner.correlation_id.as_str(&self.inner.request_id)
    }

    /// Returns the request method.
    pub const fn method(&self) -> &Method {
        &self.method
    }

    /// Returns the stable matched route pattern when available.
    ///
    /// This depends on Axum's [`axum::extract::MatchedPath`] extension being
    /// present before the context is built. Layer placement can affect whether
    /// this is available for a given router shape.
    pub fn route(&self) -> Option<&str> {
        self.inner.route.as_deref()
    }

    /// Returns the raw request path.
    pub fn path(&self) -> &str {
        &self.inner.path
    }

    /// Returns the trace id when available.
    ///
    /// The value is extracted only when the W3C `traceparent` header is valid.
    pub fn trace_id(&self) -> Option<&str> {
        self.inner.trace_id.as_deref()
    }

    /// Returns the validated parent span id from `traceparent` when available.
    pub fn span_id(&self) -> Option<&str> {
        self.inner.span_id.as_deref()
    }

    /// Returns the inferred client kind.
    ///
    /// `x-api-key` takes precedence over `Authorization`; otherwise the request
    /// is classified as anonymous.
    pub const fn client_kind(&self) -> ClientKind {
        self.client_kind
    }

    /// Returns the optional application user id.
    pub fn user_id(&self) -> Option<&str> {
        self.inner.user_id.as_deref()
    }

    /// Returns the optional application tenant id.
    pub fn tenant_id(&self) -> Option<&str> {
        self.inner.tenant_id.as_deref()
    }

    /// Returns the optional application session id.
    pub fn session_id(&self) -> Option<&str> {
        self.inner.session_id.as_deref()
    }

    /// Sets the stable matched route pattern.
    pub fn with_route(mut self, route: impl Into<String>) -> Self {
        Arc::make_mut(&mut self.inner).route = Some(route.into());
        self
    }

    /// Sets an application user id.
    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
        Arc::make_mut(&mut self.inner).user_id = Some(user_id.into());
        self
    }

    /// Sets an application tenant id.
    pub fn with_tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
        Arc::make_mut(&mut self.inner).tenant_id = Some(tenant_id.into());
        self
    }

    /// Sets an application session id.
    pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
        Arc::make_mut(&mut self.inner).session_id = Some(session_id.into());
        self
    }
}

impl<S> FromRequestParts<S> for RequestContext
where
    S: Send + Sync,
{
    type Rejection = axum::http::StatusCode;

    fn from_request_parts(
        parts: &mut Parts,
        _state: &S,
    ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
        let context = parts.extensions.get::<Self>().cloned();
        async move { context.ok_or(axum::http::StatusCode::INTERNAL_SERVER_ERROR) }
    }
}

/// Request identity used by rate limiters and observers.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct RequestIdentity(String);

impl RequestIdentity {
    /// Creates a request identity from a stable label.
    pub fn new(value: impl Into<String>) -> Self {
        Self(value.into())
    }

    /// Returns the identity label.
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// Extracts a rate-limit identity from request parts.
pub trait IdentityExtractor: Clone + Send + Sync + 'static {
    /// Returns the identity for this request.
    fn extract(&self, parts: &Parts) -> Option<RequestIdentity>;
}

impl<F> IdentityExtractor for F
where
    F: Fn(&Parts) -> Option<RequestIdentity> + Clone + Send + Sync + 'static,
{
    fn extract(&self, parts: &Parts) -> Option<RequestIdentity> {
        self(parts)
    }
}

/// Builds an identity extractor that prefers user/tenant/API key context fields.
pub fn context_identity() -> impl IdentityExtractor {
    |parts: &Parts| {
        if let Some(context) = parts.extensions.get::<RequestContext>()
            && let Some(value) = context.user_id().or_else(|| context.tenant_id())
        {
            return Some(RequestIdentity::new(value.to_owned()));
        }
        header_to_string(&parts.headers, "x-api-key").map(RequestIdentity::new)
    }
}

/// Builds an identity extractor from API key headers.
pub fn api_key_identity() -> impl IdentityExtractor {
    |parts: &Parts| header_to_string(&parts.headers, "x-api-key").map(RequestIdentity::new)
}

/// Builds an identity extractor from the connected client IP address.
///
/// This extractor uses Axum's [`axum::extract::ConnectInfo<SocketAddr>`]
/// extension and ignores `X-Forwarded-For`. Nidus serving helpers populate
/// `ConnectInfo` on the normal `listen`/`serve` path. If a router is exercised
/// without peer information, the identity falls back to `"anonymous"`.
pub fn client_ip_identity() -> impl IdentityExtractor {
    |parts: &Parts| {
        peer_ip(parts)
            .map(|ip| RequestIdentity::new(ip.to_string()))
            .or_else(|| Some(RequestIdentity::new("anonymous")))
    }
}

/// Builds an identity extractor that trusts `X-Forwarded-For` only from known proxies.
///
/// Use this when Nidus runs behind a reverse proxy that rewrites or appends
/// `X-Forwarded-For` and the direct peer address is one of the configured
/// trusted proxy IPs. Requests from untrusted peers ignore `X-Forwarded-For`
/// and use the direct peer IP. Requests without peer information fall back to
/// `"anonymous"`.
///
/// Forwarded chains are evaluated from right to left across every header value.
/// Traversal continues only while the current hop is trusted and stops at the
/// first untrusted or malformed address. Configure every proxy that is allowed
/// to extend the chain; untrusted prefixes cannot select the returned identity.
pub fn trusted_proxy_client_ip_identity(
    trusted_proxies: impl IntoIterator<Item = IpAddr>,
) -> impl IdentityExtractor {
    let trusted_proxies: Arc<[IpAddr]> = trusted_proxies.into_iter().collect::<Vec<_>>().into();
    move |parts: &Parts| {
        peer_ip(parts)
            .map(|peer| {
                let client_ip =
                    trusted_forwarded_client_ip(&parts.headers, peer, trusted_proxies.as_ref());
                RequestIdentity::new(client_ip.to_string())
            })
            .or_else(|| Some(RequestIdentity::new("anonymous")))
    }
}

fn peer_ip(parts: &Parts) -> Option<IpAddr> {
    parts
        .extensions
        .get::<axum::extract::ConnectInfo<SocketAddr>>()
        .map(|connect| connect.0.ip())
}

fn trusted_forwarded_client_ip(
    headers: &HeaderMap,
    peer: IpAddr,
    trusted_proxies: &[IpAddr],
) -> IpAddr {
    let mut client_ip = peer;
    'chain: for value in headers.get_all("x-forwarded-for").iter().rev() {
        let Ok(value) = value.to_str() else {
            break;
        };
        for hop in value.rsplit(',') {
            if !trusted_proxies.contains(&client_ip) {
                break 'chain;
            }
            let Ok(hop) = hop.trim().parse() else {
                break 'chain;
            };
            client_ip = hop;
        }
    }
    client_ip
}

pub(crate) fn header_to_string(headers: &HeaderMap, name: &'static str) -> Option<String> {
    header_to_str(headers, name).map(str::to_owned)
}

pub(crate) fn header_to_str<'a>(headers: &'a HeaderMap, name: &'static str) -> Option<&'a str> {
    headers
        .get(name)
        .and_then(|value| value.to_str().ok())
        .filter(|value| !value.is_empty())
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct ParsedTraceParent<'a> {
    pub(crate) trace_id: &'a str,
    pub(crate) parent_id: &'a str,
    pub(crate) flags: u8,
}

pub(crate) fn parse_traceparent(value: &str) -> Option<ParsedTraceParent<'_>> {
    let bytes = value.as_bytes();
    if bytes.len() < 55 || bytes[2] != b'-' || bytes[35] != b'-' || bytes[52] != b'-' {
        return None;
    }

    let version = &bytes[..2];
    let trace_id = &bytes[3..35];
    let parent_id = &bytes[36..52];
    let flags = &bytes[53..55];
    if !is_lower_hex(version)
        || version == b"ff"
        || !is_valid_id(trace_id)
        || !is_valid_id(parent_id)
        || !is_lower_hex(flags)
        || (version == b"00" && bytes.len() != 55)
        || (bytes.len() > 55 && bytes[55] != b'-')
    {
        return None;
    }

    Some(ParsedTraceParent {
        trace_id: &value[3..35],
        parent_id: &value[36..52],
        flags: u8::from_str_radix(&value[53..55], 16).ok()?,
    })
}

fn is_valid_id(value: &[u8]) -> bool {
    is_lower_hex(value) && value.iter().any(|byte| *byte != b'0')
}

fn is_lower_hex(value: &[u8]) -> bool {
    value
        .iter()
        .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(byte))
}

fn infer_client_kind(headers: &HeaderMap) -> ClientKind {
    if headers.contains_key("x-api-key") {
        ClientKind::ApiKey
    } else if headers.contains_key(http::header::AUTHORIZATION) {
        ClientKind::Authenticated
    } else {
        ClientKind::Anonymous
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use http::Request;

    #[test]
    fn request_context_can_consume_request_id() {
        let context = RequestContext::new("req-123", Method::GET, "/users");

        assert_eq!(context.into_request_id(), "req-123");
    }

    #[test]
    fn request_context_new_does_not_implicitly_set_a_correlation_id() {
        let context = RequestContext::new("req-123", Method::GET, "/users");

        assert_eq!(context.correlation_id(), None);
        assert_eq!(context.inner.correlation_id, CorrelationId::None);
    }

    #[test]
    fn request_context_from_parts_reuses_request_id_as_correlation_fallback() {
        let (parts, ()) = Request::builder()
            .uri("/users")
            .body(())
            .unwrap()
            .into_parts();

        let context = RequestContext::from_parts(&parts, "req-123");

        assert_eq!(context.correlation_id(), Some("req-123"));
        assert_eq!(context.inner.correlation_id, CorrelationId::RequestId);
    }

    #[test]
    fn request_context_from_parts_preserves_explicit_correlation_id() {
        let (parts, ()) = Request::builder()
            .uri("/users")
            .header("x-correlation-id", "corr-456")
            .body(())
            .unwrap()
            .into_parts();

        let context = RequestContext::from_parts(&parts, "req-123");

        assert_eq!(context.correlation_id(), Some("corr-456"));
        assert_eq!(
            context.inner.correlation_id,
            CorrelationId::Explicit("corr-456".to_owned())
        );
    }

    #[test]
    fn request_context_from_parts_keeps_empty_ids_absent() {
        let (parts, ()) = Request::builder()
            .uri("/users")
            .body(())
            .unwrap()
            .into_parts();

        let context = RequestContext::from_parts(&parts, "");

        assert_eq!(context.correlation_id(), None);
        assert_eq!(context.inner.correlation_id, CorrelationId::None);
    }

    #[test]
    fn cloned_request_context_uses_copy_on_write_for_enrichment() {
        let original = RequestContext::new("req-123", Method::GET, "/users");
        let shared = original.clone();
        assert!(Arc::ptr_eq(&original.inner, &shared.inner));

        let enriched = shared
            .with_route("/users/{id}")
            .with_user_id("user-42")
            .with_tenant_id("tenant-7")
            .with_session_id("session-9");

        assert!(!Arc::ptr_eq(&original.inner, &enriched.inner));
        assert_eq!(original.route(), None);
        assert_eq!(original.user_id(), None);
        assert_eq!(enriched.route(), Some("/users/{id}"));
        assert_eq!(enriched.user_id(), Some("user-42"));
        assert_eq!(enriched.tenant_id(), Some("tenant-7"));
        assert_eq!(enriched.session_id(), Some("session-9"));
    }

    #[test]
    fn request_context_extracts_valid_trace_and_parent_span_ids() {
        let (parts, ()) = Request::builder()
            .uri("/users/42")
            .header(
                "traceparent",
                "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
            )
            .body(())
            .unwrap()
            .into_parts();

        let context = RequestContext::from_parts(&parts, "request-1");

        assert_eq!(context.trace_id(), Some("4bf92f3577b34da6a3ce929d0e0e4736"));
        assert_eq!(context.span_id(), Some("00f067aa0ba902b7"));
    }

    #[test]
    fn request_context_ignores_invalid_traceparent_values() {
        for traceparent in [
            "not-a-traceparent",
            "ff-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
            "00-00000000000000000000000000000000-00f067aa0ba902b7-01",
            "00-4BF92F3577B34DA6A3CE929D0E0E4736-00f067aa0ba902b7-01",
            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01-extra",
        ] {
            let (parts, ()) = Request::builder()
                .uri("/")
                .header("traceparent", traceparent)
                .body(())
                .unwrap()
                .into_parts();
            let context = RequestContext::from_parts(&parts, "request-1");

            assert_eq!(context.trace_id(), None, "{traceparent}");
            assert_eq!(context.span_id(), None, "{traceparent}");
        }
    }

    #[test]
    fn client_ip_identity_ignores_forwarded_headers_without_peer_info() {
        let parts = request_parts(None, Some("203.0.113.10"));
        let identity = client_ip_identity().extract(&parts).unwrap();

        assert_eq!(identity.as_str(), "anonymous");
    }

    #[test]
    fn trusted_proxy_client_ip_identity_uses_forwarded_header_from_trusted_peer() {
        let parts = request_parts(Some("127.0.0.1:5000"), Some("203.0.113.10, 10.0.0.5"));
        let trusted_proxy = "127.0.0.1".parse::<IpAddr>().unwrap();
        let intermediate_proxy = "10.0.0.5".parse::<IpAddr>().unwrap();
        let identity = trusted_proxy_client_ip_identity([trusted_proxy, intermediate_proxy])
            .extract(&parts)
            .unwrap();

        assert_eq!(identity.as_str(), "203.0.113.10");
    }

    #[test]
    fn trusted_proxy_client_ip_identity_rejects_untrusted_forwarded_prefixes() {
        let parts = request_parts(Some("127.0.0.1:5000"), Some("198.51.100.200, 203.0.113.10"));
        let trusted_proxy = "127.0.0.1".parse::<IpAddr>().unwrap();
        let identity = trusted_proxy_client_ip_identity([trusted_proxy])
            .extract(&parts)
            .unwrap();

        assert_eq!(identity.as_str(), "203.0.113.10");
    }

    #[test]
    fn trusted_proxy_client_ip_identity_scans_all_forwarded_header_values() {
        let mut parts = request_parts(Some("127.0.0.1:5000"), Some("198.51.100.200"));
        parts
            .headers
            .append("x-forwarded-for", "203.0.113.10".parse().unwrap());
        let trusted_proxy = "127.0.0.1".parse::<IpAddr>().unwrap();
        let identity = trusted_proxy_client_ip_identity([trusted_proxy])
            .extract(&parts)
            .unwrap();

        assert_eq!(identity.as_str(), "203.0.113.10");
    }

    #[test]
    fn trusted_proxy_client_ip_identity_stops_at_malformed_forwarded_hop() {
        let parts = request_parts(Some("127.0.0.1:5000"), Some("203.0.113.10, not-an-ip"));
        let trusted_proxy = "127.0.0.1".parse::<IpAddr>().unwrap();
        let identity = trusted_proxy_client_ip_identity([trusted_proxy])
            .extract(&parts)
            .unwrap();

        assert_eq!(identity.as_str(), "127.0.0.1");
    }

    #[test]
    fn trusted_proxy_client_ip_identity_ignores_forwarded_header_from_untrusted_peer() {
        let parts = request_parts(Some("127.0.0.1:5000"), Some("203.0.113.10"));
        let trusted_proxy = "10.0.0.1".parse::<IpAddr>().unwrap();
        let identity = trusted_proxy_client_ip_identity([trusted_proxy])
            .extract(&parts)
            .unwrap();

        assert_eq!(identity.as_str(), "127.0.0.1");
    }

    fn request_parts(peer: Option<&str>, forwarded_for: Option<&str>) -> Parts {
        let mut builder = Request::builder().uri("/");
        if let Some(forwarded_for) = forwarded_for {
            builder = builder.header("x-forwarded-for", forwarded_for);
        }
        let (mut parts, ()) = builder.body(()).unwrap().into_parts();
        if let Some(peer) = peer {
            parts.extensions.insert(axum::extract::ConnectInfo(
                peer.parse::<SocketAddr>().unwrap(),
            ));
        }
        parts
    }
}