feat: harden community production foundation through story 1.5

apps/admin-api/src/request_context.rs

@@ -4,20 +4,30 @@ use axum::{
     middleware::Next,
     response::Response,
 };
-use crank_observability::{RequestId, set_remote_trace_parent, with_request_correlation};
+use crank_core::{CorrelationContext, RequestId, TraceContext};
+use crank_observability::{set_remote_trace_parent, with_request_correlation};
 use tracing::{Instrument, info, info_span};
 
 pub const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
+pub const TRACE_ID_HEADER: HeaderName = HeaderName::from_static("x-trace-id");
 
 #[derive(Clone, Debug)]
 pub struct RequestContext {
-    pub request_id: String,
+    pub correlation: CorrelationContext,
+}
+
+impl RequestContext {
+    pub fn request_id(&self) -> &str {
+        self.correlation.request_id().as_str()
+    }
+
+    pub fn trace_id(&self) -> &str {
+        self.correlation.trace_id().as_str()
+    }
 }
 
 pub async fn apply_request_context(mut request: Request, next: Next) -> Response {
-    let context = RequestContext {
-        request_id: RequestId::resolve_from_headers(request.headers()).into_string(),
-    };
+    let (request_id, remote_parent) = resolve_correlation(request.headers());
     let method = request.method().clone();
     let route = request
         .extensions()
@@ -27,41 +37,105 @@ pub async fn apply_request_context(mut request: Request, next: Next) -> Response
     let span = info_span!(
         target: "crank::trace",
         "http.request",
-        request_id = %context.request_id,
+        request_id = %request_id,
+        trace_id = tracing::field::Empty,
     );
-    set_remote_trace_parent(&span, request.headers());
+    if let Some(remote_parent) = remote_parent.as_ref() {
+        set_canonical_parent(&span, remote_parent);
+    }
+    let trace_context = crank_trace::trace_context_for_span(&span).unwrap_or_else(|| {
+        remote_parent
+            .as_ref()
+            .map_or_else(TraceContext::generate, TraceContext::continue_local)
+    });
+    span.record("trace_id", trace_context.trace_id().as_str());
+    let context = RequestContext {
+        correlation: CorrelationContext::new(request_id, trace_context),
+    };
     request.extensions_mut().insert(context.clone());
 
-    with_request_correlation(context.request_id.clone(), async move {
-        let mut response = next.run(request).instrument(span).await;
-        info!(
-            name: "admin.request.completed",
-            request_id = %context.request_id,
-            method = %method,
-            route,
-            status = response.status().as_u16(),
-            "admin request completed"
-        );
-        if let Ok(value) = HeaderValue::from_str(&context.request_id) {
-            response.headers_mut().insert(REQUEST_ID_HEADER, value);
-        }
-        response
-    })
+    with_request_correlation(
+        context.correlation.request_id().to_string(),
+        context.correlation.trace_id().to_string(),
+        async move {
+            let mut response = next.run(request).instrument(span).await;
+            info!(
+                name: "admin.request.completed",
+                request_id = %context.correlation.request_id(),
+                trace_id = %context.correlation.trace_id(),
+                method = %method,
+                route,
+                status = response.status().as_u16(),
+                "admin request completed"
+            );
+            if let Ok(value) = HeaderValue::from_str(context.correlation.request_id().as_str()) {
+                response.headers_mut().insert(REQUEST_ID_HEADER, value);
+            }
+            if let Ok(value) = HeaderValue::from_str(context.correlation.trace_id().as_str()) {
+                response.headers_mut().insert(TRACE_ID_HEADER, value);
+            }
+            response
+        },
+    )
     .await
 }
 
+fn resolve_correlation(headers: &axum::http::HeaderMap) -> (RequestId, Option<TraceContext>) {
+    let _tracestate_accepted = one_auxiliary_header_within_budget(
+        headers,
+        "tracestate",
+        TraceContext::tracestate_within_budget,
+    );
+    let _baggage_accepted =
+        one_auxiliary_header_within_budget(headers, "baggage", TraceContext::baggage_within_budget);
+    let mut request_ids = headers.get_all(REQUEST_ID_HEADER).iter();
+    let request_id = request_ids.next().and_then(|value| value.to_str().ok());
+    let request_id = if request_ids.next().is_some() {
+        RequestId::generate()
+    } else {
+        RequestId::resolve(request_id)
+    };
+
+    let mut traceparents = headers.get_all("traceparent").iter();
+    let traceparent = traceparents.next().and_then(|value| value.to_str().ok());
+    let remote_parent = if traceparents.next().is_some() {
+        None
+    } else {
+        traceparent.and_then(|value| TraceContext::parse(value).ok())
+    };
+    (request_id, remote_parent)
+}
+
+fn one_auxiliary_header_within_budget(
+    headers: &axum::http::HeaderMap,
+    name: &'static str,
+    validate: fn(&str) -> bool,
+) -> bool {
+    let mut values = headers.get_all(name).iter();
+    let value = values.next().and_then(|value| value.to_str().ok());
+    values.next().is_none() && value.is_some_and(validate)
+}
+
+fn set_canonical_parent(span: &tracing::Span, context: &TraceContext) {
+    let mut headers = axum::http::HeaderMap::new();
+    if let Ok(value) = HeaderValue::from_str(context.traceparent()) {
+        headers.insert("traceparent", value);
+        set_remote_trace_parent(span, &headers);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     #[test]
     fn accepts_visible_ascii_request_ids() {
-        assert!(crank_observability::RequestId::is_valid("req_test_123"));
-        assert!(crank_observability::RequestId::is_valid("trace-123/abc"));
+        assert!(crank_core::RequestId::is_valid("req_test_123"));
+        assert!(crank_core::RequestId::is_valid("trace-123/abc"));
     }
 
     #[test]
     fn rejects_empty_or_control_request_ids() {
-        assert!(!crank_observability::RequestId::is_valid(""));
-        assert!(!crank_observability::RequestId::is_valid("bad value"));
-        assert!(!crank_observability::RequestId::is_valid("bad\nvalue"));
+        assert!(!crank_core::RequestId::is_valid(""));
+        assert!(!crank_core::RequestId::is_valid("bad value"));
+        assert!(!crank_core::RequestId::is_valid("bad\nvalue"));
     }
 }