provide_telemetry/
propagation.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 provide.io llc
2// SPDX-License-Identifier: Apache-2.0
3// SPDX-Comment: Part of provide-telemetry.
4//
5
6use std::collections::BTreeMap;
7
8use serde_json::Value;
9
10use crate::context::{bind_context, ContextGuard};
11use crate::tracer::set_trace_context;
12
13const MAX_HEADER_LENGTH: usize = 512;
14const MAX_TRACESTATE_PAIRS: usize = 32;
15const MAX_BAGGAGE_LENGTH: usize = 8192;
16
17#[derive(Clone, Debug, Default, PartialEq, Eq)]
18pub struct PropagationContext {
19    pub traceparent: Option<String>,
20    pub tracestate: Option<String>,
21    pub baggage: Option<String>,
22    pub trace_id: Option<String>,
23    pub span_id: Option<String>,
24}
25
26pub struct PropagationGuard {
27    trace_guard: Option<ContextGuard>,
28    context_guard: Option<ContextGuard>,
29}
30
31impl Drop for PropagationGuard {
32    #[cfg_attr(test, mutants::skip)] // Equivalent mutant: fields still drop after an empty body.
33    fn drop(&mut self) {
34        drop(self.trace_guard.take());
35        drop(self.context_guard.take());
36    }
37}
38
39/// Parse a W3C baggage header into key-value pairs.
40/// Properties after `;` are stripped. Empty keys are skipped.
41pub fn parse_baggage(raw: &str) -> BTreeMap<String, String> {
42    let mut result = BTreeMap::new();
43    for member in raw.split(',') {
44        let kv = member.split(';').next().unwrap_or("");
45        if let Some(eq_idx) = kv.find('=') {
46            let key = kv[..eq_idx].trim();
47            if !key.is_empty() {
48                let value = kv[eq_idx + 1..].trim();
49                result.insert(key.to_string(), value.to_string());
50            }
51        }
52    }
53    result
54}
55
56fn parse_traceparent(value: Option<&str>) -> (Option<String>, Option<String>, Option<String>) {
57    let Some(raw) = value else {
58        return (None, None, None);
59    };
60    let parts = raw.split('-').collect::<Vec<_>>();
61    if parts.len() != 4 {
62        return (None, None, None);
63    }
64    let version = parts[0];
65    let trace_id = parts[1];
66    let span_id = parts[2];
67    let flags = parts[3];
68    let valid = version.len() == 2
69        && trace_id.len() == 32
70        && span_id.len() == 16
71        && flags.len() == 2
72        && !version.eq_ignore_ascii_case("ff")
73        && trace_id != "00000000000000000000000000000000"
74        && span_id != "0000000000000000"
75        && [version, trace_id, span_id, flags]
76            .iter()
77            .all(|part| part.chars().all(|ch| ch.is_ascii_hexdigit()));
78
79    if !valid {
80        return (None, None, None);
81    }
82
83    (
84        Some(raw.to_string()),
85        Some(trace_id.to_ascii_lowercase()),
86        Some(span_id.to_ascii_lowercase()),
87    )
88}
89
90pub fn extract_w3c_context(
91    traceparent: Option<&str>,
92    tracestate: Option<&str>,
93    baggage: Option<&str>,
94) -> PropagationContext {
95    let traceparent = traceparent.filter(|value| value.len() <= MAX_HEADER_LENGTH);
96    let tracestate = tracestate.and_then(|value| {
97        if value.len() > MAX_HEADER_LENGTH || value.split(',').count() > MAX_TRACESTATE_PAIRS {
98            None
99        } else {
100            Some(value.to_string())
101        }
102    });
103    let baggage = baggage.and_then(|value| {
104        if value.len() > MAX_BAGGAGE_LENGTH {
105            None
106        } else {
107            Some(value.to_string())
108        }
109    });
110    let (traceparent, trace_id, span_id) = parse_traceparent(traceparent);
111
112    PropagationContext {
113        traceparent,
114        tracestate,
115        baggage,
116        trace_id,
117        span_id,
118    }
119}
120
121pub fn bind_propagation_context(context: PropagationContext) -> PropagationGuard {
122    let mut fields = Vec::new();
123    if let Some(traceparent) = context.traceparent.clone() {
124        fields.push(("traceparent".to_string(), Value::String(traceparent)));
125    }
126    if let Some(tracestate) = context.tracestate.clone() {
127        fields.push(("tracestate".to_string(), Value::String(tracestate)));
128    }
129    if let Some(ref baggage) = context.baggage {
130        fields.push(("baggage".to_string(), Value::String(baggage.clone())));
131        for (k, v) in parse_baggage(baggage) {
132            fields.push((format!("baggage.{k}"), Value::String(v)));
133        }
134    }
135
136    let context_guard = if fields.is_empty() {
137        None
138    } else {
139        Some(bind_context(fields))
140    };
141    let trace_guard = if context.trace_id.is_some() || context.span_id.is_some() {
142        Some(set_trace_context(context.trace_id, context.span_id))
143    } else {
144        None
145    };
146
147    PropagationGuard {
148        trace_guard,
149        context_guard,
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    use serde_json::json;
158
159    use crate::context::get_context;
160    use crate::testing::acquire_test_state_lock;
161    use crate::tracer::get_trace_context;
162
163    #[test]
164    fn propagation_test_a_parse_baggage_keeps_pairs_and_strips_parameters() {
165        let baggage = parse_baggage("user=alice;prop=x,env=prod;ttl=100,invalid,=skip");
166
167        assert_eq!(baggage.get("user").map(String::as_str), Some("alice"));
168        assert_eq!(baggage.get("env").map(String::as_str), Some("prod"));
169        assert_eq!(baggage.len(), 2);
170    }
171
172    #[test]
173    fn propagation_test_a_bind_propagation_context_roundtrip_restores_state() {
174        let _guard = acquire_test_state_lock();
175        let context = PropagationContext {
176            traceparent: Some(
177                "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string(),
178            ),
179            tracestate: Some("k=v".to_string()),
180            baggage: Some("user=alice,env=prod".to_string()),
181            trace_id: Some("4bf92f3577b34da6a3ce929d0e0e4736".to_string()),
182            span_id: Some("00f067aa0ba902b7".to_string()),
183        };
184
185        {
186            let _propagation = bind_propagation_context(context);
187            let trace = get_trace_context();
188            let fields = get_context();
189            assert_eq!(
190                trace.get("trace_id").and_then(std::clone::Clone::clone),
191                Some("4bf92f3577b34da6a3ce929d0e0e4736".to_string())
192            );
193            assert_eq!(
194                trace.get("span_id").and_then(std::clone::Clone::clone),
195                Some("00f067aa0ba902b7".to_string())
196            );
197            assert_eq!(
198                fields.get("traceparent"),
199                Some(&json!(
200                    "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
201                ))
202            );
203            assert_eq!(fields.get("tracestate"), Some(&json!("k=v")));
204            assert_eq!(fields.get("baggage"), Some(&json!("user=alice,env=prod")));
205            assert_eq!(fields.get("baggage.user"), Some(&json!("alice")));
206            assert_eq!(fields.get("baggage.env"), Some(&json!("prod")));
207        }
208
209        assert!(get_context().is_empty());
210        let trace = get_trace_context();
211        assert_eq!(trace.get("trace_id"), Some(&None));
212        assert_eq!(trace.get("span_id"), Some(&None));
213    }
214}
provide_telemetry/propagation.rs

provide_telemetry/
propagation.rs
