ullun/
trace.rs

1//! W3C Trace Context (`traceparent`) propagation
2//!
3//! A request that crosses several services is only debuggable when every service
4//! logs the same trace id. [`TraceLayer`] reads the incoming `traceparent`
5//! header (or starts a new trace), gives the request its own span, and makes the
6//! ids available to handlers, the access log and outgoing calls:
7//!
8//! ```no_run
9//! use ullun::client::HttpClient;
10//! use ullun::prelude::*;
11//! use ullun::trace::TraceLayer;
12//!
13//! # async fn run() {
14//! App::new()
15//!     .middleware(TraceLayer::new())
16//!     .get("/orders", |req: Request| async move {
17//!         // the call below carries the same trace id; the callee's parent is this span
18//!         let upstream = HttpClient::new()
19//!             .get("http://inventory.internal/stock")
20//!             .trace(&req)
21//!             .send()
22//!             .await?;
23//!         Ok(Response::text(upstream.text()?))
24//!     })
25//!     .run("127.0.0.1:3000")
26//!     .await
27//!     .unwrap();
28//! # }
29//! ```
30//!
31//! [`AccessLog`](crate::observability::AccessLog) adds `trace=<id>` (text) or
32//! `trace_id` (JSON) to its lines when the layer is installed. This is the
33//! propagation half of distributed tracing, without an exporter: ship the ids
34//! with your logs, or read [`Request::trace`] from your own OpenTelemetry glue.
35//! Facing the internet, use [`TraceLayer::ignore_incoming`] so callers cannot
36//! choose your trace ids.
37
38use crate::error::Result;
39use crate::middleware::{Middleware, Next};
40use crate::request::Request;
41use crate::response::Response;
42use crate::security::random_bytes;
43use async_trait::async_trait;
44
45/// Trace identity of the current request, stored in `Request::extensions`
46#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct TraceContext {
48    /// 32 lowercase hex digits shared by every span of the trace
49    pub trace_id: String,
50    /// 16 hex digits identifying this service's span
51    pub span_id: String,
52    /// The caller's span, when the request arrived with a `traceparent`
53    pub parent_span_id: Option<String>,
54    /// The sampled flag from the caller (new traces are sampled)
55    pub sampled: bool,
56    /// Vendor state, passed on unchanged
57    pub tracestate: Option<String>,
58}
59
60fn hex(bytes: &[u8]) -> String {
61    bytes.iter().map(|b| format!("{b:02x}")).collect()
62}
63
64fn random_hex(len: usize) -> Result<String> {
65    // all-zero ids are invalid; vanishingly unlikely, but cheap to rule out
66    loop {
67        let bytes = random_bytes(len)?;
68        if bytes.iter().any(|b| *b != 0) {
69            return Ok(hex(&bytes));
70        }
71    }
72}
73
74fn is_hex(s: &str, len: usize) -> bool {
75    s.len() == len && s.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f'))
76}
77
78impl TraceContext {
79    /// Parse a `traceparent` header value (`00-<trace>-<span>-<flags>`)
80    pub fn parse(header: &str) -> Option<(String, String, bool)> {
81        let mut parts = header.trim().split('-');
82        let (version, trace, span, flags) =
83            (parts.next()?, parts.next()?, parts.next()?, parts.next()?);
84        // version ff is forbidden; future versions may append fields
85        if !is_hex(version, 2) || version == "ff" || (version == "00" && parts.next().is_some()) {
86            return None;
87        }
88        if !is_hex(trace, 32) || trace.bytes().all(|b| b == b'0') {
89            return None;
90        }
91        if !is_hex(span, 16) || span.bytes().all(|b| b == b'0') {
92            return None;
93        }
94        if !is_hex(flags, 2) {
95            return None;
96        }
97        let sampled = u8::from_str_radix(flags, 16).ok()? & 1 == 1;
98        Some((trace.to_string(), span.to_string(), sampled))
99    }
100
101    /// `traceparent` value naming this request's own span
102    pub fn traceparent(&self) -> String {
103        format!(
104            "00-{}-{}-{}",
105            self.trace_id,
106            self.span_id,
107            if self.sampled { "01" } else { "00" }
108        )
109    }
110}
111
112impl Request {
113    /// Trace ids assigned by [`TraceLayer`], if installed
114    pub fn trace(&self) -> Option<&TraceContext> {
115        self.extensions.get::<TraceContext>()
116    }
117}
118
119/// Middleware that continues or starts a W3C trace
120#[derive(Clone, Debug)]
121pub struct TraceLayer {
122    trust_incoming: bool,
123    echo: bool,
124}
125
126impl Default for TraceLayer {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132impl TraceLayer {
133    pub fn new() -> Self {
134        Self {
135            trust_incoming: true,
136            echo: false,
137        }
138    }
139
140    /// Always start a new trace, whatever the caller sent (public edge)
141    pub fn ignore_incoming(mut self) -> Self {
142        self.trust_incoming = false;
143        self
144    }
145
146    /// Return this request's `traceparent` in the response so callers can
147    /// find the trace of a request they made
148    pub fn echo_response(mut self) -> Self {
149        self.echo = true;
150        self
151    }
152}
153
154#[async_trait]
155impl Middleware for TraceLayer {
156    async fn handle(&self, mut req: Request, next: Next) -> Result<Response> {
157        let incoming = self
158            .trust_incoming
159            .then(|| {
160                req.headers
161                    .get("traceparent")
162                    .and_then(|v| v.to_str().ok())
163                    .and_then(TraceContext::parse)
164            })
165            .flatten();
166        let tracestate = incoming
167            .as_ref()
168            .and_then(|_| req.headers.get("tracestate"))
169            .and_then(|v| v.to_str().ok())
170            .filter(|v| v.len() <= 512)
171            .map(str::to_string);
172        let ctx = match incoming {
173            Some((trace_id, parent, sampled)) => TraceContext {
174                trace_id,
175                span_id: random_hex(8)?,
176                parent_span_id: Some(parent),
177                sampled,
178                tracestate,
179            },
180            None => TraceContext {
181                trace_id: random_hex(16)?,
182                span_id: random_hex(8)?,
183                parent_span_id: None,
184                sampled: true,
185                tracestate: None,
186            },
187        };
188        let own = ctx.traceparent();
189        req.extensions.insert(ctx);
190        let res = next.run(req).await.unwrap_or_else(Response::from_error);
191        Ok(if self.echo {
192            res.with_header("traceparent", own)
193        } else {
194            res
195        })
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    const TRACE: &str = "4bf92f3577b34da6a3ce929d0e0e4736";
204    const SPAN: &str = "00f067aa0ba902b7";
205
206    #[test]
207    fn parses_valid_headers() {
208        let (t, s, sampled) = TraceContext::parse(&format!("00-{TRACE}-{SPAN}-01")).unwrap();
209        assert_eq!((t.as_str(), s.as_str(), sampled), (TRACE, SPAN, true));
210        let (_, _, sampled) = TraceContext::parse(&format!("00-{TRACE}-{SPAN}-00")).unwrap();
211        assert!(!sampled);
212        // a future version may carry extra fields
213        assert!(TraceContext::parse(&format!("01-{TRACE}-{SPAN}-01-extra")).is_some());
214    }
215
216    #[test]
217    fn rejects_malformed_headers() {
218        for bad in [
219            "".to_string(),
220            "garbage".to_string(),
221            format!("ff-{TRACE}-{SPAN}-01"),
222            format!("00-{TRACE}-{SPAN}"),
223            format!("00-{TRACE}-{SPAN}-01-extra"),
224            format!("00-{}-{SPAN}-01", "0".repeat(32)),
225            format!("00-{TRACE}-{}-01", "0".repeat(16)),
226            format!("00-{}-{SPAN}-01", TRACE.to_uppercase()),
227            format!("00-{}-{SPAN}-01", &TRACE[..30]),
228            format!("00-{TRACE}-{SPAN}-0g"),
229        ] {
230            assert!(TraceContext::parse(&bad).is_none(), "{bad}");
231        }
232    }
233
234    #[test]
235    fn own_traceparent_names_the_span() {
236        let ctx = TraceContext {
237            trace_id: TRACE.into(),
238            span_id: SPAN.into(),
239            parent_span_id: None,
240            sampled: true,
241            tracestate: None,
242        };
243        assert_eq!(ctx.traceparent(), format!("00-{TRACE}-{SPAN}-01"));
244        let unsampled = TraceContext {
245            sampled: false,
246            ..ctx
247        };
248        assert_eq!(unsampled.traceparent(), format!("00-{TRACE}-{SPAN}-00"));
249    }
250}
ullun/trace.rs

ullun/
trace.rs
