1use std::sync::OnceLock;
9
10use blake3::Hasher;
11use http::HeaderMap;
12
13#[derive(Debug, Clone, Default, PartialEq, Eq)]
20pub struct ObsTraceCtx {
21 pub trace_id: String,
23 pub span_id: String,
25 pub flags: String,
27 pub tracestate: String,
29}
30
31impl ObsTraceCtx {
32 #[must_use]
34 pub fn sampled(&self) -> bool {
35 self.flags.ends_with('1')
36 }
37
38 #[must_use]
41 pub fn fresh(sampled: bool) -> Self {
42 Self {
43 trace_id: fresh_trace_id(),
44 span_id: fresh_span_id(),
45 flags: if sampled {
46 "01".to_string()
47 } else {
48 "00".to_string()
49 },
50 tracestate: String::new(),
51 }
52 }
53
54 #[must_use]
57 pub fn child_of(&self) -> Self {
58 Self {
59 trace_id: self.trace_id.clone(),
60 span_id: fresh_span_id(),
61 flags: self.flags.clone(),
62 tracestate: self.tracestate.clone(),
63 }
64 }
65}
66
67#[derive(Debug, Clone, Copy, Default)]
69pub struct W3cPropagator;
70
71impl W3cPropagator {
72 #[must_use]
74 pub fn new() -> Self {
75 Self
76 }
77
78 #[must_use]
81 pub fn extract(&self, headers: &HeaderMap) -> Option<ObsTraceCtx> {
82 extract_w3c(headers)
83 }
84
85 pub fn inject(&self, headers: &mut HeaderMap, ctx: &ObsTraceCtx) {
87 inject_w3c(headers, ctx);
88 }
89}
90
91#[must_use]
94pub fn extract_w3c(headers: &HeaderMap) -> Option<ObsTraceCtx> {
95 let raw = headers.get("traceparent")?.to_str().ok()?;
96 let mut parts = raw.split('-');
97 let version = parts.next()?;
98 let trace_id = parts.next()?;
99 let span_id = parts.next()?;
100 let flags = parts.next()?;
101 if parts.next().is_some() || version != "00" {
102 return None;
103 }
104 if trace_id.len() != 32 || span_id.len() != 16 || flags.len() != 2 {
105 return None;
106 }
107 if !trace_id.bytes().all(is_hex) || !span_id.bytes().all(is_hex) || !flags.bytes().all(is_hex) {
108 return None;
109 }
110 let tracestate = headers
111 .get("tracestate")
112 .and_then(|v| v.to_str().ok())
113 .unwrap_or("")
114 .to_string();
115 Some(ObsTraceCtx {
116 trace_id: trace_id.to_string(),
117 span_id: span_id.to_string(),
118 flags: flags.to_string(),
119 tracestate,
120 })
121}
122
123pub fn inject_w3c(headers: &mut HeaderMap, ctx: &ObsTraceCtx) {
125 let value = format!("00-{}-{}-{}", ctx.trace_id, ctx.span_id, ctx.flags);
126 if let Ok(v) = http::HeaderValue::from_str(&value) {
127 headers.insert("traceparent", v);
128 }
129 if !ctx.tracestate.is_empty()
130 && let Ok(v) = http::HeaderValue::from_str(&ctx.tracestate)
131 {
132 headers.insert("tracestate", v);
133 }
134}
135
136#[must_use]
140pub fn fresh_trace_id() -> String {
141 let mut buf = [0u8; 32];
142 fill_random(&mut buf);
143 let mut hasher = Hasher::new();
144 hasher.update(&buf);
145 let hash = hasher.finalize();
146 let bytes = hash.as_bytes();
147 let mut out = String::with_capacity(32);
148 for b in &bytes[..16] {
149 use std::fmt::Write;
150 let _ = write!(&mut out, "{b:02x}");
151 }
152 out
153}
154
155#[must_use]
158pub fn fresh_span_id() -> String {
159 let mut buf = [0u8; 16];
160 fill_random(&mut buf);
161 let mut hasher = Hasher::new();
162 hasher.update(&buf);
163 let hash = hasher.finalize();
164 let bytes = hash.as_bytes();
165 let mut out = String::with_capacity(16);
166 for b in &bytes[..8] {
167 use std::fmt::Write;
168 let _ = write!(&mut out, "{b:02x}");
169 }
170 out
171}
172
173#[must_use]
175pub fn status_class(status: u16) -> &'static str {
176 match status {
177 100..=199 => "1xx",
178 200..=299 => "2xx",
179 300..=399 => "3xx",
180 400..=499 => "4xx",
181 500..=599 => "5xx",
182 _ => "err",
183 }
184}
185
186fn is_hex(b: u8) -> bool {
187 b.is_ascii_hexdigit()
188}
189
190fn fill_random(buf: &mut [u8]) {
191 if getrandom::fill(buf).is_ok() {
196 return;
197 }
198 let counter = next_counter();
199 let now = std::time::SystemTime::now()
200 .duration_since(std::time::UNIX_EPOCH)
201 .map(|d| d.as_nanos() as u64)
202 .unwrap_or(0);
203 let mut hasher = Hasher::new();
204 hasher.update(&counter.to_le_bytes());
205 hasher.update(&now.to_le_bytes());
206 let h = hasher.finalize();
207 let bytes = h.as_bytes();
208 for (i, byte) in buf.iter_mut().enumerate() {
209 if let Some(&b) = bytes.get(i % 32) {
210 *byte = b;
211 }
212 }
213}
214
215fn next_counter() -> u64 {
216 static COUNTER: OnceLock<std::sync::atomic::AtomicU64> = OnceLock::new();
217 COUNTER
218 .get_or_init(|| std::sync::atomic::AtomicU64::new(0))
219 .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
220}
221
222#[cfg(test)]
223mod tests {
224 use http::HeaderMap;
225
226 use super::*;
227
228 #[test]
229 fn test_propagator_round_trip() {
230 let mut headers = HeaderMap::new();
231 let ctx_in = ObsTraceCtx {
232 trace_id: "0123456789abcdef0123456789abcdef".to_string(),
233 span_id: "0123456789abcdef".to_string(),
234 flags: "01".to_string(),
235 tracestate: "vendor=value".to_string(),
236 };
237 inject_w3c(&mut headers, &ctx_in);
238 let ctx_out = extract_w3c(&headers).expect("parse");
239 assert_eq!(ctx_in.trace_id, ctx_out.trace_id);
240 assert_eq!(ctx_in.span_id, ctx_out.span_id);
241 assert_eq!(ctx_in.flags, ctx_out.flags);
242 assert_eq!(ctx_in.tracestate, ctx_out.tracestate);
243 assert!(ctx_out.sampled());
244 }
245
246 #[test]
247 fn test_extract_rejects_malformed() {
248 let mut headers = HeaderMap::new();
249 headers.insert("traceparent", "garbage".parse().unwrap());
250 assert!(extract_w3c(&headers).is_none());
251 headers.insert("traceparent", "01-aa-bb-cc".parse().unwrap());
252 assert!(extract_w3c(&headers).is_none());
253 }
254
255 #[test]
256 fn test_fresh_ids_should_be_correct_length_and_hex() {
257 let t = fresh_trace_id();
258 let s = fresh_span_id();
259 assert_eq!(t.len(), 32);
260 assert_eq!(s.len(), 16);
261 assert!(t.bytes().all(is_hex));
262 assert!(s.bytes().all(is_hex));
263 }
264
265 #[test]
266 fn test_fresh_ids_are_unique() {
267 let a = fresh_trace_id();
268 let b = fresh_trace_id();
269 assert_ne!(a, b);
270 }
271
272 #[test]
273 fn test_child_of_inherits_trace_id() {
274 let parent = ObsTraceCtx::fresh(true);
275 let child = parent.child_of();
276 assert_eq!(parent.trace_id, child.trace_id);
277 assert_ne!(parent.span_id, child.span_id);
278 assert_eq!(parent.flags, child.flags);
279 }
280
281 #[test]
282 fn test_status_class_should_classify() {
283 assert_eq!(status_class(200), "2xx");
284 assert_eq!(status_class(404), "4xx");
285 assert_eq!(status_class(503), "5xx");
286 assert_eq!(status_class(0), "err");
287 }
288}