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.and_then(|value| {
96        if value.len() > MAX_HEADER_LENGTH {
97            None
98        } else {
99            Some(value)
100        }
101    });
102    let tracestate = tracestate.and_then(|value| {
103        if value.len() > MAX_HEADER_LENGTH || value.split(',').count() > MAX_TRACESTATE_PAIRS {
104            None
105        } else {
106            Some(value.to_string())
107        }
108    });
109    let baggage = baggage.and_then(|value| {
110        if value.len() > MAX_BAGGAGE_LENGTH {
111            None
112        } else {
113            Some(value.to_string())
114        }
115    });
116    let (traceparent, trace_id, span_id) = parse_traceparent(traceparent);
117
118    PropagationContext {
119        traceparent,
120        tracestate,
121        baggage,
122        trace_id,
123        span_id,
124    }
125}
126
127pub fn bind_propagation_context(context: PropagationContext) -> PropagationGuard {
128    let mut fields = Vec::new();
129    if let Some(traceparent) = context.traceparent.clone() {
130        fields.push(("traceparent".to_string(), Value::String(traceparent)));
131    }
132    if let Some(tracestate) = context.tracestate.clone() {
133        fields.push(("tracestate".to_string(), Value::String(tracestate)));
134    }
135    if let Some(ref baggage) = context.baggage {
136        fields.push(("baggage".to_string(), Value::String(baggage.clone())));
137        for (k, v) in parse_baggage(baggage) {
138            fields.push((format!("baggage.{k}"), Value::String(v)));
139        }
140    }
141
142    let context_guard = if fields.is_empty() {
143        None
144    } else {
145        Some(bind_context(fields))
146    };
147    let trace_guard = if context.trace_id.is_some() || context.span_id.is_some() {
148        Some(set_trace_context(context.trace_id, context.span_id))
149    } else {
150        None
151    };
152
153    PropagationGuard {
154        trace_guard,
155        context_guard,
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    use serde_json::json;
164
165    use crate::context::get_context;
166    use crate::testing::acquire_test_state_lock;
167    use crate::tracer::get_trace_context;
168
169    #[test]
170    fn propagation_test_a_parse_baggage_keeps_pairs_and_strips_parameters() {
171        let baggage = parse_baggage("user=alice;prop=x,env=prod;ttl=100,invalid,=skip");
172
173        assert_eq!(baggage.get("user").map(String::as_str), Some("alice"));
174        assert_eq!(baggage.get("env").map(String::as_str), Some("prod"));
175        assert_eq!(baggage.len(), 2);
176    }
177
178    #[test]
179    fn propagation_test_a_bind_propagation_context_roundtrip_restores_state() {
180        let _guard = acquire_test_state_lock();
181        let context = PropagationContext {
182            traceparent: Some(
183                "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string(),
184            ),
185            tracestate: Some("k=v".to_string()),
186            baggage: Some("user=alice,env=prod".to_string()),
187            trace_id: Some("4bf92f3577b34da6a3ce929d0e0e4736".to_string()),
188            span_id: Some("00f067aa0ba902b7".to_string()),
189        };
190
191        {
192            let _propagation = bind_propagation_context(context);
193            let trace = get_trace_context();
194            let fields = get_context();
195            assert_eq!(
196                trace.get("trace_id").and_then(std::clone::Clone::clone),
197                Some("4bf92f3577b34da6a3ce929d0e0e4736".to_string())
198            );
199            assert_eq!(
200                trace.get("span_id").and_then(std::clone::Clone::clone),
201                Some("00f067aa0ba902b7".to_string())
202            );
203            assert_eq!(
204                fields.get("traceparent"),
205                Some(&json!(
206                    "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
207                ))
208            );
209            assert_eq!(fields.get("tracestate"), Some(&json!("k=v")));
210            assert_eq!(fields.get("baggage"), Some(&json!("user=alice,env=prod")));
211            assert_eq!(fields.get("baggage.user"), Some(&json!("alice")));
212            assert_eq!(fields.get("baggage.env"), Some(&json!("prod")));
213        }
214
215        assert!(get_context().is_empty());
216        let trace = get_trace_context();
217        assert_eq!(trace.get("trace_id"), Some(&None));
218        assert_eq!(trace.get("span_id"), Some(&None));
219    }
220}
provide_telemetry/propagation.rs

provide_telemetry/
propagation.rs
