shared_framework/logging/

correlation.rs

//! Per-request correlation context.
//!
//! [`CorrelationContext`] carries request-scoped state — correlation and request IDs,
//! flow marker, user ID, raw body, headers, query parameters, decoded multipart
//! payload, arbitrary key/value data, and pagination state — through handlers.
//! It is cheaply clonable via `Arc` and can also be propagated via the
//! [`CORRELATION_CTX`] task-local.
//!
//! ```ignore
//! let ctx = CorrelationContext::new();
//! let page: Option<String> = ctx.query_param("page");
//! ```

use crate::ErrorResult;
use crate::doc::DocumentableDTO;
use crate::utils::request_parser::{MultipartBody, UploadedFile};
use crate::validation::Validate;
use serde_json::Value;
use std::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use uuid::Uuid;

/// Lifecycle marker for a correlated unit of work.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CorrelationFlow {
    /// A single standalone unit of work.
    Once,
    /// The start of a multi-step unit of work.
    Start,
    /// A middle step of a multi-step unit of work.
    Continue,
    /// The final step of a multi-step unit of work.
    End,
}

impl std::str::FromStr for CorrelationFlow {
    type Err = String;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.trim().to_uppercase().as_str() {
            "ONCE" => Ok(Self::Once),
            "START" => Ok(Self::Start),
            "CONTINUE" => Ok(Self::Continue),
            "END" => Ok(Self::End),
            other => Err(format!("unknown flow {other}")),
        }
    }
}

impl std::fmt::Display for CorrelationFlow {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = match self {
            Self::Once => "ONCE",
            Self::Start => "START",
            Self::Continue => "CONTINUE",
            Self::End => "END",
        };
        write!(f, "{s}")
    }
}

/// Per-request correlation context. Cheaply clonable via `Arc`.
/// Shares the underlying request state; mutating one clone is visible to the others.
#[derive(Debug, Clone)]
pub struct CorrelationContext {
    inner: Arc<Mutex<Inner>>,
}

#[derive(Debug)]
struct Inner {
    correlation_id: String,
    request_id: String,
    flow: CorrelationFlow,
    user_id: Option<String>,
    body: Option<Vec<u8>>,
    headers: http::HeaderMap,
    query_params: HashMap<String, String>,
    multipart: Option<MultipartBody>,
    data: HashMap<String, Box<dyn Any + Send + Sync>>,
    path_params: HashMap<String, String>,
    // pagination
    pagination_cursor: Option<String>,
    pagination_limit: usize,
}

impl CorrelationContext {
    /// Creates a context with fresh random correlation and request IDs,
    /// `Once` flow, no user, and a default pagination limit of 15.
    pub fn new() -> Self {
        let corr = Uuid::new_v4().to_string();
        let req = hex::encode(rand::random::<[u8; 8]>());
        Self {
            inner: Arc::new(Mutex::new(Inner {
                correlation_id: corr,
                request_id: req,
                flow: CorrelationFlow::Once,
                user_id: None,
                body: None,
                headers: http::HeaderMap::new(),
                query_params: HashMap::new(),
                path_params: HashMap::new(),
                multipart: None,
                data: HashMap::new(),
                pagination_cursor: None,
                pagination_limit: 15,
            })),
        }
    }

    /// Creates a context with the given correlation and request IDs.
    /// Other state matches [`CorrelationContext::new`].
    pub fn with_ids(correlation_id: &str, request_id: &str) -> Self {
        let ctx = Self::new();
        {
            let mut inner = ctx.inner.lock().unwrap();
            inner.correlation_id = correlation_id.to_string();
            inner.request_id = request_id.to_string();
        }
        ctx
    }

    pub(crate) fn set_body(&self, body: Vec<u8>) {
        self.inner.lock().unwrap().body = Some(body);
    }

    pub(crate) fn set_headers(&self, headers: http::HeaderMap) {
        self.inner.lock().unwrap().headers = headers;
    }

    pub(crate) fn set_params(&self, params: HashMap<String, String>) {
        self.inner.lock().unwrap().path_params = params;
    }

    pub(crate) fn set_query_params(&self, params: HashMap<String, String>) {
        self.inner.lock().unwrap().query_params = params;
    }

    pub(crate) fn set_multipart(&self, multipart: MultipartBody) {
        self.inner.lock().unwrap().multipart = Some(multipart);
    }

    /// Whether the request arrived as `multipart/form-data`.
    pub fn is_multipart(&self) -> bool {
        self.inner.lock().unwrap().multipart.is_some()
    }

    /// Decoded multipart payload, if the request was multipart.
    pub fn multipart(&self) -> Option<MultipartBody> {
        self.inner.lock().unwrap().multipart.clone()
    }

    /// First value of a multipart text field, if the request was multipart.
    pub fn form_field(&self, name: &str) -> Option<String> {
        self.inner
            .lock()
            .unwrap()
            .multipart
            .as_ref()
            .and_then(|mp| mp.field(name))
            .map(|s| s.to_string())
    }

    /// All values of a repeated multipart text field.
    pub fn form_field_all(&self, name: &str) -> Vec<String> {
        self.inner
            .lock()
            .unwrap()
            .multipart
            .as_ref()
            .and_then(|mp| mp.fields.get(name).cloned())
            .unwrap_or_default()
    }

    /// Uploaded files in part order. Returns an empty vector when not multipart.
    pub fn files(&self) -> Vec<UploadedFile> {
        self.inner
            .lock()
            .unwrap()
            .multipart
            .as_ref()
            .map(|mp| mp.files.clone())
            .unwrap_or_default()
    }

    /// Uploaded files for one form field.
    pub fn files_for(&self, field: &str) -> Vec<UploadedFile> {
        self.files()
            .into_iter()
            .filter(|f| f.field_name == field)
            .collect()
    }

    /// Raw request body bytes, if any.
    pub fn body_bytes(&self) -> Option<Vec<u8>> {
        self.inner.lock().unwrap().body.clone()
    }

    /// Raw request body as a string.
    /// For multipart requests the `body` form field wins when present;
    /// otherwise the raw payload is decoded as UTF-8.
    /// Returns a `400` error when there is no request body, or it is not valid UTF-8.
    /// No validation is applied to unstructured payloads.
    pub fn body_string(&self) -> Result<String, ErrorResult> {
        if let Some(payload) = self.form_field("body") {
            return Ok(payload);
        }
        let body_bytes_opt = self.inner.lock().unwrap().body.clone();
        let Some(body_bytes) = body_bytes_opt else {
            return Err(ErrorResult::bad_request("no body"));
        };
        String::from_utf8(body_bytes).map_err(|_| ErrorResult::bad_request("invalid body"))
    }

    /// Deserializes and validates the request body as `T`.
    ///
    /// Multipart requests decode the DTO from the `body` form field when present,
    /// otherwise from the merged text fields. Non-multipart bodies accept JSON
    /// with a form-urlencoded fallback. Returns a `400` error when the body is
    /// missing, cannot be deserialized, or fails [`Validate`](Validate).
    ///
    /// ```ignore
    /// let dto: MyDto = ctx.body::<MyDto>()?;
    /// ```
    pub fn body<T>(&self) -> Result<T, ErrorResult>
    where
        T: DocumentableDTO + Validate,
    {
        // Multipart requests carry the DTO separately from the files — decode it
        // in place instead of delegating parsing to the caller.
        if let Some(mp) = self.inner.lock().unwrap().multipart.clone() {
            return self.multipart_body(&mp);
        }

        // 1. Extract and clone the optional byte vector
        let body_bytes_opt = self.inner.lock().unwrap().body.clone();
        let Some(body_bytes) = body_bytes_opt else {
            return Err(ErrorResult::bad_request("no body"));
        };

        // 2. Deserialize from the bytes. JSON is tried first; form-urlencoded
        // is accepted as a fallback.
        let parsed_body: Option<T> = serde_json::from_slice(&body_bytes).ok().or_else(|| {
            serde_urlencoded::from_bytes::<Value>(&body_bytes)
                .ok()
                .and_then(|v| serde_json::from_value(v).ok())
        });
        let Some(parsed_body) = parsed_body else {
            return Err(ErrorResult::bad_request("invalid body"));
        };

        // 3. Validate the deserialized struct
        if let Err(validated) = parsed_body.validate() {
            return Err(ErrorResult::new(
                validated.message,
                Some(Value::String(validated.field)),
                400,
            ));
        }

        // 4. Return the validated body (requires T to implement Clone)
        Ok(parsed_body.clone())
    }

    /// Deserialize a DTO from a decoded multipart payload.
    fn multipart_body<T>(&self, mp: &MultipartBody) -> Result<T, ErrorResult>
    where
        T: DocumentableDTO + Validate,
    {
        // Multipart requests encode the actual JSON payload into a single
        // `body` form field.
        if let Some(payload) = mp.field("body") {
            let parsed: T = serde_json::from_str(payload)
                .map_err(|_| ErrorResult::bad_request("invalid body"))?;
            return self.validated(parsed);
        }
        // Otherwise merge the text fields into an object (single values as
        // strings, repeated names as arrays) and deserialize from that.
        let mut map = serde_json::Map::new();
        for (k, vs) in &mp.fields {
            let v = if vs.len() == 1 {
                serde_json::Value::String(vs[0].clone())
            } else {
                serde_json::Value::Array(
                    vs.iter().cloned().map(serde_json::Value::String).collect(),
                )
            };
            map.insert(k.clone(), v);
        }
        if map.is_empty() {
            return Err(ErrorResult::bad_request("no body"));
        }
        let parsed: T = serde_json::from_value(serde_json::Value::Object(map))
            .map_err(|_| ErrorResult::bad_request("invalid body"))?;
        self.validated(parsed)
    }

    fn validated<T>(&self, parsed: T) -> Result<T, ErrorResult>
    where
        T: DocumentableDTO + Validate,
    {
        if let Err(e) = parsed.validate() {
            return Err(ErrorResult::bad_request(e.message));
        }
        Ok(parsed)
    }

    /// All request headers attached at dispatch.
    pub fn headers(&self) -> http::HeaderMap {
        self.inner.lock().unwrap().headers.clone()
    }

    /// A single request header value, if present.
    pub fn header(&self, name: &str) -> Option<String> {
        self.inner
            .lock()
            .unwrap()
            .headers
            .get(name)
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string())
    }

    /// All query parameters attached at dispatch.
    pub fn query_params(&self) -> HashMap<String, String> {
        self.inner.lock().unwrap().query_params.clone()
    }

    /// A single query parameter value, if present.
    pub fn query_param(&self, name: &str) -> Option<String> {
        self.inner.lock().unwrap().query_params.get(name).cloned()
    }

    /// A query parameter value, or `default` when absent.
    pub fn query_param_or(&self, name: &str, default: &str) -> String {
        self.query_param(name)
            .unwrap_or_else(|| default.to_string())
    }

    /// All path parameters attached at dispatch.
    pub fn path_params(&self) -> HashMap<String, String> {
        self.inner.lock().unwrap().path_params.clone()
    }

    /// A single path parameter value, if present.
    pub fn path_param(&self, name: &str) -> Option<String> {
        self.inner.lock().unwrap().path_params.get(name).cloned()
    }

    /// A path parameter value, or `default` when absent.
    pub fn path_param_or(&self, name: &str, default: &str) -> String {
        self.query_param(name)
            .unwrap_or_else(|| default.to_string())
    }

    /// Returns the correlation ID shared across related requests.
    pub fn correlation_id(&self) -> String {
        self.inner.lock().unwrap().correlation_id.clone()
    }

    /// Returns the ID unique to this request.
    pub fn request_id(&self) -> String {
        self.inner.lock().unwrap().request_id.clone()
    }

    /// Replaces the request ID.
    pub fn set_request_id(&self, id: &str) {
        self.inner.lock().unwrap().request_id = id.to_string();
    }

    /// Replaces the correlation ID.
    pub fn set_correlation_id(&self, id: &str) {
        self.inner.lock().unwrap().correlation_id = id.to_string();
    }

    /// Returns the current lifecycle flow marker.
    pub fn flow(&self) -> CorrelationFlow {
        self.inner.lock().unwrap().flow
    }

    /// Replaces the lifecycle flow marker.
    pub fn set_flow(&self, flow: CorrelationFlow) {
        self.inner.lock().unwrap().flow = flow;
    }

    /// Sets the lifecycle flow marker and returns the context for chaining.
    pub fn with_flow(self, flow: CorrelationFlow) -> Self {
        self.set_flow(flow);
        self
    }

    /// Sets the correlation ID and returns the context for chaining.
    pub fn with_correlation_id(self, id: &str) -> Self {
        self.set_correlation_id(id);
        self
    }

    /// Returns the authenticated user ID, if one was attached.
    pub fn user_id(&self) -> Option<String> {
        self.inner.lock().unwrap().user_id.clone()
    }

    /// Sets or clears the authenticated user ID.
    pub fn set_user_id(&self, id: Option<String>) {
        self.inner.lock().unwrap().user_id = id;
    }

    /// Stores an arbitrary typed value under `key`.
    pub fn set<T>(&self, key: &str, value: T)
    where
        T: Any + Send + Sync,
    {
        self.inner
            .lock()
            .unwrap()
            .data
            .insert(key.to_string(), Box::new(value));
    }

    /// Returns a cloned value of type `T` stored under `key`, if the key exists
    /// and its value has that type.
    pub fn get<T>(&self, key: &str) -> Option<T>
    where
        T: Any + Clone,
    {
        self.inner
            .lock()
            .unwrap()
            .data
            .get(key)
            .and_then(|value| value.downcast_ref::<T>())
            .cloned()
    }

    /// Stores a string value under `key`.
    pub fn set_string(&self, key: &str, value: impl Into<String>) {
        self.set(key, value.into());
    }

    /// Returns a cloned string value stored under `key`, if any.
    pub fn get_string(&self, key: &str) -> Option<String> {
        self.get(key)
    }

    /// Stores a boolean value under `key`.
    pub fn set_bool(&self, key: &str, value: bool) {
        self.set(key, value);
    }

    /// Returns a boolean value stored under `key`, if any.
    pub fn get_bool(&self, key: &str) -> Option<bool> {
        self.get(key)
    }

    /// Stores a numeric `f64` value under `key`.
    pub fn set_number(&self, key: &str, value: f64) {
        self.set(key, value);
    }

    /// Returns a numeric `f64` value stored under `key`, if any.
    pub fn get_number(&self, key: &str) -> Option<f64> {
        self.get(key)
    }

    /// Returns the pagination cursor, if one was attached.
    pub fn pagination_cursor(&self) -> Option<String> {
        self.inner.lock().unwrap().pagination_cursor.clone()
    }

    /// Returns the pagination limit (defaults to 15).
    pub fn pagination_limit(&self) -> usize {
        self.inner.lock().unwrap().pagination_limit
    }

    /// Sets the pagination cursor and limit.
    pub fn set_pagination(&self, cursor: Option<String>, limit: usize) {
        let mut inner = self.inner.lock().unwrap();
        inner.pagination_cursor = cursor;
        inner.pagination_limit = limit;
    }

    /// Extracts the client IP, preferring `x-forwarded-for` (first entry),
    /// then `x-real-ip`, then the socket address, else `"unknown"`.
    pub fn client_ip(
        headers: &http::HeaderMap,
        remote_addr: Option<std::net::SocketAddr>,
    ) -> String {
        if let Some(v) = headers.get("x-forwarded-for").and_then(|h| h.to_str().ok()) {
            if let Some(first) = v.split(',').next() {
                let ip = first.trim();
                if !ip.is_empty() {
                    return ip.to_string();
                }
            }
        }
        if let Some(v) = headers.get("x-real-ip").and_then(|h| h.to_str().ok()) {
            return v.to_string();
        }
        remote_addr
            .map(|a| a.ip().to_string())
            .unwrap_or_else(|| "unknown".to_string())
    }
}

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

tokio::task_local! {
    /// Task-local [`CorrelationContext`] for the current async task, when set by the dispatcher.
    pub static CORRELATION_CTX: CorrelationContext;
}

#[cfg(test)]
mod tests {
    use super::CorrelationContext;

    #[test]
    fn stores_and_returns_owned_typed_values() {
        let context = CorrelationContext::new();
        context.set("count", 42_u32);

        assert_eq!(context.get::<u32>("count"), Some(42));
        assert_eq!(context.get::<String>("count"), None);
        assert_eq!(context.get::<u32>("missing"), None);
    }

    #[test]
    fn convenience_accessors_store_and_return_values() {
        let context = CorrelationContext::new();
        context.set_string("name", "Ada");
        context.set_bool("enabled", true);
        context.set_number("ratio", 1.5);

        assert_eq!(context.get_string("name"), Some("Ada".to_string()));
        assert_eq!(context.get_bool("enabled"), Some(true));
        assert_eq!(context.get_number("ratio"), Some(1.5));
    }
}