Skip to main content

product_os_router/
axum_trace_context.rs

1//! W3C Trace Context extractor for framework_axum.
2//!
3//! Standalone version that implements `axum::extract::FromRequestParts`.
4
5use std::prelude::v1::*;
6
7use axum::extract::FromRequestParts;
8use axum::http::request::Parts;
9use axum::response::{IntoResponse, Response};
10
11/// Parsed W3C `traceparent` header.
12#[derive(Debug, Clone)]
13pub struct TraceContext {
14    pub version: u8,
15    pub trace_id: String,
16    pub parent_id: String,
17    pub trace_flags: u8,
18    pub trace_state: Option<String>,
19}
20
21#[derive(Debug)]
22pub struct TraceContextRejection(pub &'static str);
23
24impl IntoResponse for TraceContextRejection {
25    fn into_response(self) -> Response {
26        (axum::http::StatusCode::BAD_REQUEST, self.0).into_response()
27    }
28}
29
30impl<S> FromRequestParts<S> for TraceContext
31where
32    S: Send + Sync,
33{
34    type Rejection = TraceContextRejection;
35
36    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
37        let header = parts
38            .headers
39            .get("traceparent")
40            .and_then(|v| v.to_str().ok())
41            .ok_or(TraceContextRejection("Missing traceparent header"))?;
42
43        let segments: Vec<&str> = header.split('-').collect();
44        if segments.len() != 4 {
45            return Err(TraceContextRejection("Invalid traceparent format"));
46        }
47
48        let version = u8::from_str_radix(segments[0], 16)
49            .map_err(|_| TraceContextRejection("Invalid version"))?;
50        let trace_id = segments[1].to_string();
51        let parent_id = segments[2].to_string();
52        let trace_flags = u8::from_str_radix(segments[3], 16)
53            .map_err(|_| TraceContextRejection("Invalid trace flags"))?;
54
55        if trace_id.len() != 32 || parent_id.len() != 16 {
56            return Err(TraceContextRejection("Invalid trace/parent ID length"));
57        }
58
59        let trace_state = parts
60            .headers
61            .get("tracestate")
62            .and_then(|v| v.to_str().ok())
63            .map(|s| s.to_string());
64
65        Ok(TraceContext {
66            version,
67            trace_id,
68            parent_id,
69            trace_flags,
70            trace_state,
71        })
72    }
73}