featherbit/plugins/util/
trace.rs

1//! Distributed-tracing span context and propagation codecs.
2//!
3//! Shared by the tracing plugins (`opentelemetry`, `zipkin`, `skywalking`),
4//! which are each a **start/end node pair** around `upstream`: the start node
5//! extracts or creates a [`SpanContext`], stores it in `context.message`
6//! (per-request state), and injects the downstream propagation header so the
7//! upstream service continues the trace; the end node loads the span, computes
8//! its duration, and exports it to the collector (fire-and-forget).
9//!
10//! Provides trace/span id generation and codecs for the three wire formats:
11//! W3C `traceparent` (OpenTelemetry), B3 (Zipkin), and SkyWalking `sw8`.
12
13use serde::{Deserialize, Serialize};
14use std::time::{SystemTime, UNIX_EPOCH};
15
16use base64::engine::general_purpose::STANDARD as BASE64;
17use base64::Engine;
18use ring::rand::{SecureRandom, SystemRandom};
19
20use crate::context::Context;
21
22/// The reserved `context.message` key under which the active span is stored.
23pub const SPAN_KEY: &str = "__trace_span";
24
25/// A span in flight: the trace it belongs to, this hop's span, its parent, the
26/// sampling decision, and when it started (unix millis).
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct SpanContext {
29    /// 32-hex-char trace id, shared across the whole distributed trace.
30    pub trace_id: String,
31    /// 16-hex-char span id for this gateway hop.
32    pub span_id: String,
33    /// The caller's span id, when the request arrived already traced.
34    pub parent_span_id: Option<String>,
35    /// Whether this trace is sampled (exported).
36    pub sampled: bool,
37    /// Span start time (unix millis).
38    pub start_ms: u64,
39}
40
41impl SpanContext {
42    /// Milliseconds elapsed since the span started.
43    pub fn duration_ms(&self) -> u64 {
44        now_ms().saturating_sub(self.start_ms)
45    }
46}
47
48/// Stores the span in `context.message`.
49pub fn store_span(ctx: &mut Context, span: &SpanContext) {
50    if let Ok(v) = serde_json::to_value(span) {
51        ctx.message.insert(SPAN_KEY.to_string(), v);
52    }
53}
54
55/// Loads the span previously stored by the start node.
56pub fn load_span(ctx: &Context) -> Option<SpanContext> {
57    serde_json::from_value(ctx.message.get(SPAN_KEY)?.clone()).ok()
58}
59
60/// A random 32-hex-char (128-bit) trace id.
61pub fn new_trace_id() -> String {
62    let mut b = [0u8; 16];
63    SystemRandom::new().fill(&mut b).expect("system RNG");
64    hex(&b)
65}
66
67/// A random 16-hex-char (64-bit) span id.
68pub fn new_span_id() -> String {
69    let mut b = [0u8; 8];
70    SystemRandom::new().fill(&mut b).expect("system RNG");
71    hex(&b)
72}
73
74/// Current unix time in milliseconds.
75pub fn now_ms() -> u64 {
76    SystemTime::now()
77        .duration_since(UNIX_EPOCH)
78        .map(|d| d.as_millis() as u64)
79        .unwrap_or(0)
80}
81
82// ---- W3C traceparent (OpenTelemetry) --------------------------------------
83
84/// Parses a W3C `traceparent` header (`00-<trace>-<parent>-<flags>`) into
85/// `(trace_id, parent_span_id, sampled)`.
86pub fn parse_traceparent(value: &str) -> Option<(String, String, bool)> {
87    let parts: Vec<&str> = value.trim().split('-').collect();
88    if parts.len() != 4 {
89        return None;
90    }
91    let trace_id = parts[1];
92    let parent = parts[2];
93    if trace_id.len() != 32 || parent.len() != 16 {
94        return None;
95    }
96    if trace_id.bytes().all(|b| b == b'0') || parent.bytes().all(|b| b == b'0') {
97        return None;
98    }
99    let sampled = u8::from_str_radix(parts[3], 16)
100        .map(|f| f & 1 == 1)
101        .unwrap_or(false);
102    Some((trace_id.to_string(), parent.to_string(), sampled))
103}
104
105/// Builds a W3C `traceparent` header value for a span.
106pub fn build_traceparent(span: &SpanContext) -> String {
107    format!(
108        "00-{}-{}-{:02x}",
109        span.trace_id,
110        span.span_id,
111        if span.sampled { 1 } else { 0 }
112    )
113}
114
115// ---- B3 (Zipkin) -----------------------------------------------------------
116
117/// Parses B3 propagation from either the single `b3` header
118/// (`trace-span-sampled-parent`) or the multi-header form, given a lookup of
119/// the relevant request headers, into `(trace_id, parent_span_id, sampled)`.
120pub fn parse_b3(
121    b3_single: Option<&str>,
122    trace_id: Option<&str>,
123    span_id: Option<&str>,
124    sampled: Option<&str>,
125) -> Option<(String, String, bool)> {
126    if let Some(single) = b3_single {
127        let parts: Vec<&str> = single.trim().split('-').collect();
128        if parts.len() >= 2 && !parts[0].is_empty() {
129            let s = parts.get(2).map(|v| *v == "1" || *v == "d").unwrap_or(true);
130            return Some((parts[0].to_string(), parts[1].to_string(), s));
131        }
132        return None;
133    }
134    let t = trace_id?;
135    let s = span_id?;
136    if t.is_empty() || s.is_empty() {
137        return None;
138    }
139    let sampled = sampled.map(|v| v == "1" || v == "true").unwrap_or(true);
140    Some((t.to_string(), s.to_string(), sampled))
141}
142
143/// Builds the B3 multi-headers for a span as `(name, value)` pairs to set on
144/// the outgoing request.
145pub fn build_b3_headers(span: &SpanContext) -> Vec<(String, String)> {
146    let mut h = vec![
147        ("x-b3-traceid".to_string(), span.trace_id.clone()),
148        ("x-b3-spanid".to_string(), span.span_id.clone()),
149        (
150            "x-b3-sampled".to_string(),
151            if span.sampled { "1" } else { "0" }.to_string(),
152        ),
153    ];
154    if let Some(parent) = &span.parent_span_id {
155        h.push(("x-b3-parentspanid".to_string(), parent.clone()));
156    }
157    h
158}
159
160// ---- SkyWalking sw8 --------------------------------------------------------
161
162/// Parses a SkyWalking `sw8` header (hyphen-separated, base64 fields) into
163/// `(trace_id, parent_segment_span_id, sampled)`. Field 0 is the sample flag,
164/// field 1 the base64 trace id, field 3 the parent span id.
165pub fn parse_sw8(value: &str) -> Option<(String, String, bool)> {
166    let parts: Vec<&str> = value.trim().split('-').collect();
167    if parts.len() < 8 {
168        return None;
169    }
170    let sampled = parts[0] == "1";
171    let trace_id = String::from_utf8(BASE64.decode(parts[1]).ok()?).ok()?;
172    // Field 3 is the parent span id — a plain integer, not a base64 field
173    // (only the trace/segment ids and the service/instance/endpoint names are
174    // base64-encoded in sw8).
175    let parent_span = parts[3].to_string();
176    Some((trace_id, parent_span, sampled))
177}
178
179/// Base64-encodes a field for an sw8 header.
180pub fn sw8_encode(field: &str) -> String {
181    BASE64.encode(field.as_bytes())
182}
183
184fn hex(bytes: &[u8]) -> String {
185    let mut s = String::with_capacity(bytes.len() * 2);
186    for b in bytes {
187        s.push_str(&format!("{:02x}", b));
188    }
189    s
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    fn span() -> SpanContext {
197        SpanContext {
198            trace_id: "0af7651916cd43dd8448eb211c80319c".to_string(),
199            span_id: "b7ad6b7169203331".to_string(),
200            parent_span_id: Some("0020000000000001".to_string()),
201            sampled: true,
202            start_ms: now_ms(),
203        }
204    }
205
206    #[test]
207    fn test_id_lengths() {
208        assert_eq!(new_trace_id().len(), 32);
209        assert_eq!(new_span_id().len(), 16);
210        // hex and non-zero-ish
211        assert!(new_trace_id().chars().all(|c| c.is_ascii_hexdigit()));
212    }
213
214    #[test]
215    fn test_traceparent_round_trip() {
216        let s = span();
217        let header = build_traceparent(&s);
218        assert_eq!(
219            header,
220            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
221        );
222        let (trace, parent, sampled) = parse_traceparent(&header).unwrap();
223        assert_eq!(trace, s.trace_id);
224        assert_eq!(parent, s.span_id);
225        assert!(sampled);
226
227        // All-zero trace id is invalid.
228        assert!(
229            parse_traceparent("00-00000000000000000000000000000000-b7ad6b7169203331-01").is_none()
230        );
231        // Wrong shape.
232        assert!(parse_traceparent("garbage").is_none());
233    }
234
235    #[test]
236    fn test_b3() {
237        let s = span();
238        let headers = build_b3_headers(&s);
239        assert!(headers
240            .iter()
241            .any(|(k, v)| k == "x-b3-traceid" && v == &s.trace_id));
242        assert!(headers.iter().any(|(k, v)| k == "x-b3-sampled" && v == "1"));
243
244        // single-header form
245        let (t, sp, sampled) = parse_b3(Some("abc-def-1-ghi"), None, None, None).unwrap();
246        assert_eq!(t, "abc");
247        assert_eq!(sp, "def");
248        assert!(sampled);
249        // multi-header form
250        let (t, sp, sampled) = parse_b3(None, Some("t1"), Some("s1"), Some("0")).unwrap();
251        assert_eq!((t.as_str(), sp.as_str(), sampled), ("t1", "s1", false));
252        // missing → None
253        assert!(parse_b3(None, None, Some("s"), None).is_none());
254    }
255
256    #[test]
257    fn test_sw8_round_trip() {
258        let trace = "1f2d4bf47bf711eab794acde48001122";
259        let header = format!(
260            "1-{}-{}-3-5-2-{}-{}-{}",
261            sw8_encode(trace),
262            sw8_encode("segment-id"),
263            sw8_encode("service"),
264            sw8_encode("instance"),
265            sw8_encode("/endpoint")
266        );
267        let (t, parent_span, sampled) = parse_sw8(&header).unwrap();
268        assert_eq!(t, trace);
269        assert_eq!(parent_span, "3");
270        assert!(sampled);
271        assert!(parse_sw8("too-short").is_none());
272    }
273
274    #[test]
275    fn test_store_load_span() {
276        use crate::context::{GatewayRequest, GatewayResponse, Protocol};
277        use bytes::Bytes;
278        use std::collections::HashMap;
279        let mut ctx = Context {
280            request: GatewayRequest {
281                method: "GET".into(),
282                path: "/".into(),
283                host: "h".into(),
284                scheme: "http".into(),
285                headers: HashMap::new(),
286                query_params: HashMap::new(),
287                body: Bytes::new(),
288                remote_addr: "1.2.3.4:5".into(),
289                protocol: Protocol::Http1,
290            },
291            response: GatewayResponse {
292                status_code: 0,
293                headers: HashMap::new(),
294                body: Bytes::new(),
295            },
296            message: HashMap::new(),
297            errors: Vec::new(),
298        };
299        assert!(load_span(&ctx).is_none());
300        let s = span();
301        store_span(&mut ctx, &s);
302        let back = load_span(&ctx).unwrap();
303        assert_eq!(back.trace_id, s.trace_id);
304        assert_eq!(back.span_id, s.span_id);
305    }
306}
featherbit/plugins/util/trace.rs

featherbit/plugins/util/
trace.rs
