architect_sdk/middleware/
trace.rs

1//! W3C Trace Context (`traceparent`) middleware.
2//!
3//! On each request it reads the incoming `traceparent` header, extracts the
4//! `trace-id`, and opens a `tracing` span carrying it — so every log event
5//! emitted while handling the request automatically includes `trace_id` with no
6//! change at the log call site. When the header is absent or malformed a fresh
7//! **root** trace id is generated, so every request is traceable end-to-end.
8//!
9//! The resolved trace id is also placed in a task-local for the duration of the
10//! request, so the SDK's own outbound clients (events, authrs) can continue the
11//! trace via [`outbound_traceparent`]. Finally it is echoed back on the response
12//! as a `traceparent` value so callers can correlate.
13//!
14//! Header format (W3C): `00-<32-hex trace-id>-<16-hex span-id>-<2-hex flags>`.
15
16use axum::{
17    extract::Request,
18    http::{header::HeaderValue, HeaderName},
19    middleware::{from_fn, FromFnLayer, Next},
20    response::Response,
21};
22use std::future::Future;
23use std::pin::Pin;
24use tracing::Instrument;
25use uuid::Uuid;
26
27/// Header name carrying the W3C trace context.
28pub const TRACEPARENT_HEADER: &str = "traceparent";
29
30tokio::task_local! {
31    /// The current request's 32-hex trace id, set by the ingress middleware for
32    /// the duration of the request. Read via [`current_trace_id`].
33    static CURRENT_TRACE_ID: String;
34}
35
36/// The trace id resolved for the current request, stored as a request extension
37/// so handlers can read it explicitly. Always a 32-hex trace id (never empty).
38#[derive(Clone, Debug)]
39pub struct TraceId(pub String);
40
41type BoxFut = Pin<Box<dyn Future<Output = Response> + Send>>;
42type TraceFn = fn(Request, Next) -> BoxFut;
43
44/// Middleware layer that resolves a `traceparent` trace id and attaches it to
45/// every log line for the request.
46///
47/// Apply once on the top-level router:
48/// ```ignore
49/// let app = router.layer(architect_sdk::middleware::trace_id_layer());
50/// ```
51pub fn trace_id_layer() -> FromFnLayer<TraceFn, (), (Request,)> {
52    from_fn(trace_id_mw as TraceFn)
53}
54
55/// The current request's trace id (32 lowercase hex), when a request handled
56/// through [`trace_id_layer`] is on the current task. Returns `None` outside a
57/// request, or on a task `spawn`ed away from the request task (task-locals do
58/// not propagate across `tokio::spawn`).
59pub fn current_trace_id() -> Option<String> {
60    CURRENT_TRACE_ID.try_with(|t| t.clone()).ok()
61}
62
63/// A `traceparent` header value that continues the current request's trace on an
64/// outbound call, with a fresh span id for this hop. `None` when no trace is in
65/// scope (see [`current_trace_id`]).
66pub fn outbound_traceparent() -> Option<String> {
67    current_trace_id().map(|tid| format_traceparent(&tid))
68}
69
70/// Format `00-<trace-id>-<span-id>-01` with a fresh random span id. `flags=01`
71/// (sampled) marks the trace as recorded.
72fn format_traceparent(trace_id: &str) -> String {
73    let span_id = &Uuid::new_v4().simple().to_string()[..16];
74    format!("00-{trace_id}-{span_id}-01")
75}
76
77/// Extract and validate the `trace-id` from a W3C `traceparent` value. Returns
78/// `None` when the shape is wrong, the version is `ff` (invalid), or the trace id
79/// is all zeros (also invalid) — the caller then starts a fresh root trace.
80fn parse_trace_id(traceparent: &str) -> Option<String> {
81    let parts: [&str; 4] = {
82        let mut it = traceparent.trim().split('-');
83        let a = it.next()?;
84        let b = it.next()?;
85        let c = it.next()?;
86        let d = it.next()?;
87        if it.next().is_some() {
88            return None; // more than 4 segments
89        }
90        [a, b, c, d]
91    };
92    let [version, trace_id, parent_id, flags] = parts;
93    let is_hex = |s: &str, n: usize| s.len() == n && s.bytes().all(|b| b.is_ascii_hexdigit());
94    if !is_hex(version, 2) || !is_hex(trace_id, 32) || !is_hex(parent_id, 16) || !is_hex(flags, 2) {
95        return None;
96    }
97    if version.eq_ignore_ascii_case("ff") {
98        return None;
99    }
100    let tid = trace_id.to_ascii_lowercase();
101    if tid.bytes().all(|b| b == b'0') {
102        return None;
103    }
104    Some(tid)
105}
106
107/// Generate a new root trace id: a random UUID as 32 lowercase hex chars.
108fn new_trace_id() -> String {
109    Uuid::new_v4().simple().to_string()
110}
111
112fn trace_id_mw(mut req: Request, next: Next) -> BoxFut {
113    // Continue an incoming W3C trace, or start a fresh root trace when the header
114    // is missing or malformed.
115    let trace_id = req
116        .headers()
117        .get(TRACEPARENT_HEADER)
118        .and_then(|v| v.to_str().ok())
119        .and_then(parse_trace_id)
120        .unwrap_or_else(new_trace_id);
121
122    let method = req.method().clone();
123    let path = req.uri().path().to_string();
124
125    // Make the id available to handlers that want it explicitly.
126    req.extensions_mut().insert(TraceId(trace_id.clone()));
127
128    let span = tracing::info_span!(
129        "request",
130        trace_id = %trace_id,
131        method = %method,
132        path = %path,
133    );
134    // Echo a traceparent (this server's span) back so clients can correlate.
135    let response_tp = format_traceparent(&trace_id);
136
137    Box::pin(
138        CURRENT_TRACE_ID
139            .scope(trace_id, async move {
140                let mut response = next.run(req).await;
141                if let Ok(value) = HeaderValue::from_str(&response_tp) {
142                    response
143                        .headers_mut()
144                        .insert(HeaderName::from_static(TRACEPARENT_HEADER), value);
145                }
146                response
147            })
148            .instrument(span),
149    )
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use axum::{body::Body, http::Request as HttpRequest, routing::get, Router};
156    use std::sync::{Arc, Mutex};
157    use tower::ServiceExt;
158
159    #[test]
160    fn parse_trace_id_accepts_valid_traceparent() {
161        let tp = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
162        assert_eq!(
163            parse_trace_id(tp).as_deref(),
164            Some("4bf92f3577b34da6a3ce929d0e0e4736")
165        );
166    }
167
168    #[test]
169    fn parse_trace_id_rejects_malformed_and_zero() {
170        assert_eq!(parse_trace_id("garbage"), None);
171        assert_eq!(parse_trace_id("00-abc-00f067aa0ba902b7-01"), None); // short trace-id
172        assert_eq!(
173            parse_trace_id("ff-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"),
174            None
175        ); // invalid version
176        assert_eq!(
177            parse_trace_id("00-00000000000000000000000000000000-00f067aa0ba902b7-01"),
178            None
179        ); // all-zero trace-id
180    }
181
182    /// A `MakeWriter` that appends all output to a shared buffer, so a test can
183    /// inspect the JSON log lines the subscriber emitted.
184    #[derive(Clone)]
185    struct BufWriter(Arc<Mutex<Vec<u8>>>);
186
187    impl std::io::Write for BufWriter {
188        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
189            self.0.lock().unwrap().extend_from_slice(buf);
190            Ok(buf.len())
191        }
192        fn flush(&mut self) -> std::io::Result<()> {
193            Ok(())
194        }
195    }
196
197    impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for BufWriter {
198        type Writer = BufWriter;
199        fn make_writer(&'a self) -> Self::Writer {
200            self.clone()
201        }
202    }
203
204    async fn handler() -> &'static str {
205        // Emits a log with no explicit trace_id; the request span must supply it.
206        tracing::info!("handling request");
207        "ok"
208    }
209
210    fn app() -> Router {
211        Router::new()
212            .route("/x", get(handler))
213            .layer(trace_id_layer())
214    }
215
216    /// Parse the trace-id out of a response `traceparent` header value.
217    fn resp_trace_id(header: Option<&str>) -> Option<String> {
218        header.and_then(parse_trace_id)
219    }
220
221    #[test]
222    fn incoming_traceparent_flows_into_logs_and_response() {
223        let buf = Arc::new(Mutex::new(Vec::new()));
224        let subscriber = tracing_subscriber::fmt()
225            .json()
226            .flatten_event(true)
227            .with_current_span(true)
228            .with_span_list(true)
229            .with_max_level(tracing::Level::TRACE)
230            .with_writer(BufWriter(buf.clone()))
231            .finish();
232
233        // Pin this as the process-wide default so the `tracing` macro level gate
234        // stays open regardless of other tests' thread-local subscribers (a
235        // thread with no subscriber otherwise drags the global max level to OFF,
236        // short-circuiting our `info!` before it reaches the buffer). This test
237        // owns the only global-default call in the crate, so it always succeeds.
238        tracing::subscriber::set_global_default(subscriber)
239            .expect("no other test should set a global subscriber");
240
241        let rt = tokio::runtime::Builder::new_current_thread()
242            .enable_all()
243            .build()
244            .unwrap();
245
246        let incoming = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
247        let echoed = rt.block_on(async {
248            let req = HttpRequest::builder()
249                .uri("/x")
250                .header(TRACEPARENT_HEADER, incoming)
251                .body(Body::empty())
252                .unwrap();
253            let resp = app().oneshot(req).await.unwrap();
254            resp.headers()
255                .get(TRACEPARENT_HEADER)
256                .and_then(|v| v.to_str().ok())
257                .map(String::from)
258        });
259
260        // Echoed back as a valid traceparent carrying the same trace-id.
261        assert_eq!(
262            resp_trace_id(echoed.as_deref()).as_deref(),
263            Some("4bf92f3577b34da6a3ce929d0e0e4736")
264        );
265
266        // Present as a top-level field in the JSON log emitted inside the handler.
267        let logs = String::from_utf8(buf.lock().unwrap().clone()).unwrap();
268        assert!(
269            logs.contains("\"trace_id\":\"4bf92f3577b34da6a3ce929d0e0e4736\""),
270            "expected trace_id in JSON logs, got: {logs}"
271        );
272        assert!(logs.contains("handling request"));
273    }
274
275    #[test]
276    fn missing_header_generates_root_trace() {
277        let rt = tokio::runtime::Builder::new_current_thread()
278            .enable_all()
279            .build()
280            .unwrap();
281        let echoed = rt.block_on(async {
282            let req = HttpRequest::builder()
283                .uri("/x")
284                .body(Body::empty())
285                .unwrap();
286            let resp = app().oneshot(req).await.unwrap();
287            resp.headers()
288                .get(TRACEPARENT_HEADER)
289                .and_then(|v| v.to_str().ok())
290                .map(String::from)
291        });
292        // A fresh root trace id is generated and echoed as a valid traceparent.
293        let tid = resp_trace_id(echoed.as_deref());
294        assert!(
295            tid.as_deref().map(|t| t.len() == 32).unwrap_or(false),
296            "expected a generated 32-hex trace id, got: {echoed:?}"
297        );
298    }
299}
architect_sdk/middleware/trace.rs

architect_sdk/middleware/
trace.rs
