diff --git a/src/api/middleware.rs b/src/api/middleware.rs index 9b90da2..763f407 100644 --- a/src/api/middleware.rs +++ b/src/api/middleware.rs @@ -12,7 +12,7 @@ use std::net::SocketAddr; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use tracing::warn; +use tracing::{warn, Instrument}; lazy_static::lazy_static! { static ref START_INSTANT: Instant = Instant::now(); @@ -398,6 +398,23 @@ pub async fn rate_limit_layer( next.run(req).await } +pub async fn correlation_id_layer(mut req: Request, next: Next) -> Response { + let correlation_id = crate::tracing::correlation::correlation_id_from_headers(req.headers()); + req.extensions_mut().insert(correlation_id.clone()); + + let span = tracing::info_span!( + "http.correlation", + correlation_id = %correlation_id.as_str() + ); + + let mut response = crate::tracing::correlation::scope(correlation_id.clone(), async move { + next.run(req).instrument(span).await + }) + .await; + crate::tracing::correlation::insert_header(response.headers_mut(), &correlation_id); + response +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/api/router.rs b/src/api/router.rs index fa95ab7..d7e89b0 100644 --- a/src/api/router.rs +++ b/src/api/router.rs @@ -109,6 +109,9 @@ pub async fn build_router(state: AppState) -> anyhow::Result { .layer(axum_mw::from_fn( crate::api::middleware::slo_monitoring_layer, )) + .layer(axum_mw::from_fn( + crate::api::middleware::correlation_id_layer, + )) .layer(cors) .with_state(state); diff --git a/src/service_mesh/client.rs b/src/service_mesh/client.rs index 5436ffa..01cc0fa 100644 --- a/src/service_mesh/client.rs +++ b/src/service_mesh/client.rs @@ -144,9 +144,12 @@ impl ServiceMeshClient { req = req.header("content-type", "application/json").body(b); } - req.header("x-service-name", &self.identity.service_account) + req = req + .header("x-service-name", &self.identity.service_account) .header("x-trust-domain", &self.identity.trust_domain) - .timeout(std::time::Duration::from_secs(endpoint.timeout_secs)) + .timeout(std::time::Duration::from_secs(endpoint.timeout_secs)); + + crate::tracing::correlation::propagate_reqwest(req) .send() .await .map_err(|e| ServiceMeshClientError::HttpError(e.to_string())) @@ -164,10 +167,12 @@ impl ServiceMeshClient { .join(endpoint.health_check_path.trim_start_matches('/')) .map_err(|e| ServiceMeshClientError::HttpError(e.to_string()))?; - match self + let req = self .http_client .get(url) - .timeout(std::time::Duration::from_secs(5)) + .timeout(std::time::Duration::from_secs(5)); + + match crate::tracing::correlation::propagate_reqwest(req) .send() .await { diff --git a/src/tracing/correlation.rs b/src/tracing/correlation.rs new file mode 100644 index 0000000..da554b5 --- /dev/null +++ b/src/tracing/correlation.rs @@ -0,0 +1,143 @@ +use std::collections::HashMap; +use std::future::Future; + +use axum::http::{HeaderMap, HeaderValue}; +use uuid::Uuid; + +pub const CORRELATION_ID_HEADER: &str = "x-correlation-id"; +const MAX_CORRELATION_ID_LEN: usize = 128; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct CorrelationId(String); + +impl CorrelationId { + pub fn generate() -> Self { + Self(Uuid::new_v4().to_string()) + } + + pub fn parse(value: &str) -> Option { + let trimmed = value.trim(); + if trimmed.is_empty() || trimmed.len() > MAX_CORRELATION_ID_LEN { + return None; + } + + if trimmed + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.')) + { + Some(Self(trimmed.to_string())) + } else { + None + } + } + + pub fn as_str(&self) -> &str { + &self.0 + } + + fn to_header_value(&self) -> HeaderValue { + HeaderValue::from_str(self.as_str()).expect("correlation ids are sanitized") + } +} + +tokio::task_local! { + static CURRENT_CORRELATION_ID: String; +} + +pub async fn scope(id: CorrelationId, future: F) -> T +where + F: Future, +{ + CURRENT_CORRELATION_ID.scope(id.0, future).await +} + +pub fn current_correlation_id() -> Option { + CURRENT_CORRELATION_ID.try_with(Clone::clone).ok() +} + +pub fn correlation_id_from_headers(headers: &HeaderMap) -> CorrelationId { + headers + .get(CORRELATION_ID_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(CorrelationId::parse) + .unwrap_or_else(CorrelationId::generate) +} + +pub fn insert_header(headers: &mut HeaderMap, id: &CorrelationId) { + headers.insert(CORRELATION_ID_HEADER, id.to_header_value()); +} + +pub fn inject_current_into_message_headers(headers: &mut HashMap) { + if let Some(id) = current_correlation_id() { + headers.insert(CORRELATION_ID_HEADER.to_string(), id); + } +} + +pub fn inject_into_message_headers(headers: &mut HashMap, id: &CorrelationId) { + headers.insert(CORRELATION_ID_HEADER.to_string(), id.as_str().to_string()); +} + +pub fn extract_from_message_headers(headers: &HashMap) -> Option { + headers + .get(CORRELATION_ID_HEADER) + .or_else(|| headers.get("X-Correlation-ID")) + .and_then(|value| CorrelationId::parse(value)) +} + +pub fn propagate_reqwest(builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder { + match current_correlation_id() { + Some(id) => builder.header(CORRELATION_ID_HEADER, id), + None => builder, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_only_safe_correlation_ids() { + assert_eq!( + CorrelationId::parse("abc-123_DEF.456").unwrap().as_str(), + "abc-123_DEF.456" + ); + assert!(CorrelationId::parse("").is_none()); + assert!(CorrelationId::parse("has spaces").is_none()); + assert!(CorrelationId::parse("secret\nnext").is_none()); + assert!(CorrelationId::parse(&"a".repeat(MAX_CORRELATION_ID_LEN + 1)).is_none()); + } + + #[test] + fn extracts_existing_header_or_generates_new_id() { + let mut headers = HeaderMap::new(); + headers.insert(CORRELATION_ID_HEADER, HeaderValue::from_static("incoming-1")); + assert_eq!(correlation_id_from_headers(&headers).as_str(), "incoming-1"); + + headers.insert(CORRELATION_ID_HEADER, HeaderValue::from_static("bad id")); + let generated = correlation_id_from_headers(&headers); + assert_ne!(generated.as_str(), "bad id"); + assert!(Uuid::parse_str(generated.as_str()).is_ok()); + } + + #[tokio::test] + async fn task_scope_exposes_current_correlation_id() { + let id = CorrelationId::parse("req-123").unwrap(); + let current = scope(id, async { current_correlation_id() }).await; + assert_eq!(current.as_deref(), Some("req-123")); + assert!(current_correlation_id().is_none()); + } + + #[tokio::test] + async fn message_headers_roundtrip_current_id() { + let id = CorrelationId::parse("event-456").unwrap(); + let mut headers = HashMap::new(); + scope(id, async { + inject_current_into_message_headers(&mut headers); + }) + .await; + assert_eq!( + extract_from_message_headers(&headers).unwrap().as_str(), + "event-456" + ); + } +} diff --git a/src/tracing/kafka_propagator.rs b/src/tracing/kafka_propagator.rs index ad78bf1..cb91389 100644 --- a/src/tracing/kafka_propagator.rs +++ b/src/tracing/kafka_propagator.rs @@ -31,6 +31,7 @@ pub fn inject_into_kafka_headers(cx: &Context) -> HashMap { global::get_text_map_propagator(|propagator| { propagator.inject_context(cx, &mut KafkaHeaderInjector(&mut headers)); }); + crate::tracing::correlation::inject_current_into_message_headers(&mut headers); headers } diff --git a/src/tracing/logging.rs b/src/tracing/logging.rs index e2eda8b..21750b6 100644 --- a/src/tracing/logging.rs +++ b/src/tracing/logging.rs @@ -198,6 +198,12 @@ where .unwrap() .insert("span_id".to_string(), serde_json::Value::String(s_id)); } + if let Some(correlation_id) = crate::tracing::correlation::current_correlation_id() { + log_record.as_object_mut().unwrap().insert( + "correlation_id".to_string(), + serde_json::Value::String(correlation_id), + ); + } if !visitor.attributes.is_empty() { log_record.as_object_mut().unwrap().insert( diff --git a/src/tracing/mod.rs b/src/tracing/mod.rs index ffcff86..5ba3de3 100644 --- a/src/tracing/mod.rs +++ b/src/tracing/mod.rs @@ -1,3 +1,4 @@ +pub mod correlation; pub mod db_tracing; pub mod exporters; pub mod kafka_propagator; diff --git a/tests/distributed_tracing_test.rs b/tests/distributed_tracing_test.rs index 94792ae..17318dc 100644 --- a/tests/distributed_tracing_test.rs +++ b/tests/distributed_tracing_test.rs @@ -17,6 +17,9 @@ use utility_backend::gateway::telemetry::{ use utility_backend::tracing::kafka_propagator::{ extract_from_kafka_headers, inject_into_kafka_headers, }; +use utility_backend::tracing::correlation::{ + extract_from_message_headers, scope, CorrelationId, CORRELATION_ID_HEADER, +}; // ── OTLP Exporter & TracerProvider ────────────────────────────────── @@ -114,6 +117,24 @@ async fn test_kafka_headers_inject_extract_roundtrip() { ); } +#[tokio::test] +async fn test_kafka_headers_include_current_correlation_id() { + let _ = init_open_telemetry("kafka-correlation-test"); + let id = CorrelationId::parse("corr-789").unwrap(); + let parent = opentelemetry::Context::new(); + + let headers = scope(id, async { inject_into_kafka_headers(&parent) }).await; + + assert_eq!( + headers.get(CORRELATION_ID_HEADER).map(String::as_str), + Some("corr-789") + ); + assert_eq!( + extract_from_message_headers(&headers).unwrap().as_str(), + "corr-789" + ); +} + #[tokio::test] async fn test_kafka_headers_empty_input_produces_empty_context() { let _ = init_open_telemetry("kafka-empty-test"); diff --git a/tests/logging_test.rs b/tests/logging_test.rs index c4c9386..2873c3d 100644 --- a/tests/logging_test.rs +++ b/tests/logging_test.rs @@ -5,6 +5,7 @@ use std::io; use std::sync::{Arc, Mutex}; use tracing::{info, span, Level}; use tracing_subscriber::{filter::LevelFilter, layer::SubscriberExt, Registry}; +use utility_backend::tracing::correlation::{scope, CorrelationId}; use utility_backend::tracing::logging::OtelJsonFormatter; #[derive(Clone)] @@ -62,13 +63,17 @@ async fn test_otel_json_logging_format() { .with(otel_layer) .with(fmt_layer); - // Run within our custom subscriber - tracing::subscriber::with_default(subscriber, || { - let test_span = span!(Level::INFO, "test_operation"); - let _enter = test_span.enter(); + let correlation_id = CorrelationId::parse("log-correlation-1").unwrap(); + scope(correlation_id, async { + // Run within our custom subscriber + tracing::subscriber::with_default(subscriber, || { + let test_span = span!(Level::INFO, "test_operation"); + let _enter = test_span.enter(); - info!("structured log inside span"); - }); + info!("structured log inside span"); + }); + }) + .await; let output_bytes = buffer.lock().unwrap().clone(); let output_str = String::from_utf8(output_bytes).unwrap(); @@ -90,4 +95,5 @@ async fn test_otel_json_logging_format() { assert!(log2["span_id"].is_string(), "Missing span_id in span log"); assert!(!log2["trace_id"].as_str().unwrap().is_empty()); assert!(!log2["span_id"].as_str().unwrap().is_empty()); + assert_eq!(log2["correlation_id"], "log-correlation-1"); } diff --git a/tests/rate_limit_integration.rs b/tests/rate_limit_integration.rs index 692c1bd..f639888 100644 --- a/tests/rate_limit_integration.rs +++ b/tests/rate_limit_integration.rs @@ -1,5 +1,5 @@ use axum::{ - body::Body, + body::{to_bytes, Body}, extract::ConnectInfo, http::{Request, StatusCode}, middleware as axum_mw, @@ -8,7 +8,11 @@ use axum::{ }; use std::net::SocketAddr; use utility_backend::api::middleware::{ - rate_limit_layer, tenant_rate_limit_layer, DynamicRateLimiter, TenantRateLimiter, + correlation_id_layer, rate_limit_layer, tenant_rate_limit_layer, DynamicRateLimiter, + TenantRateLimiter, +}; +use utility_backend::tracing::correlation::{ + current_correlation_id, CORRELATION_ID_HEADER, }; #[tokio::test] @@ -113,3 +117,31 @@ async fn test_tenant_rate_limit_anonymous() { let res = send_no_header(app.clone()).await; assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); } + +#[tokio::test] +async fn test_correlation_id_middleware_propagates_context_and_response_header() { + let app = Router::new() + .route( + "/", + get(|| async { current_correlation_id().unwrap_or_else(|| "missing".to_string()) }), + ) + .layer(axum_mw::from_fn(correlation_id_layer)); + + let req = Request::builder() + .uri("/") + .header(CORRELATION_ID_HEADER, "request-abc") + .body(Body::empty()) + .unwrap(); + + let res = tower::ServiceExt::oneshot(app, req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + assert_eq!( + res.headers() + .get(CORRELATION_ID_HEADER) + .and_then(|value| value.to_str().ok()), + Some("request-abc") + ); + + let body = to_bytes(res.into_body(), usize::MAX).await.unwrap(); + assert_eq!(body.as_ref(), b"request-abc"); +}