product_os_router/
axum_trace_context.rs1use std::prelude::v1::*;
6
7use axum::extract::FromRequestParts;
8use axum::http::request::Parts;
9use axum::response::{IntoResponse, Response};
10
11#[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}