1use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19
20#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
25pub struct TraceCtx {
26 pub trace_id: String,
28 #[serde(skip_serializing_if = "Option::is_none")]
30 pub parent_id: Option<String>,
31 #[serde(default, skip_serializing_if = "Vec::is_empty")]
34 pub baggage: Vec<BaggageEntry>,
35}
36
37pub const MAX_BAGGAGE_ENTRIES: usize = 16;
39
40pub const MAX_BAGGAGE_ITEM_BYTES: usize = 256;
42
43#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
45pub struct BaggageEntry {
46 pub key: String,
47 pub value: String,
48}
49
50impl TraceCtx {
51 pub fn generate() -> Self {
53 let trace_id = uuid::Uuid::new_v4().as_simple().to_string();
54 Self {
55 trace_id,
56 parent_id: None,
57 baggage: Vec::new(),
58 }
59 }
60
61 pub fn from_trace_id(trace_id: impl Into<String>) -> Self {
63 Self {
64 trace_id: trace_id.into(),
65 parent_id: None,
66 baggage: Vec::new(),
67 }
68 }
69
70 pub fn from_legacy_trace_id(legacy_id: impl Into<String>) -> Self {
74 Self::from_trace_id(legacy_id)
75 }
76
77 pub fn to_legacy_trace_id(&self) -> &str {
81 &self.trace_id
82 }
83
84 pub fn with_parent(mut self, parent_id: impl Into<String>) -> Self {
86 self.parent_id = Some(parent_id.into());
87 self
88 }
89
90 pub fn child(&self, span_id: impl Into<String>) -> Self {
92 Self {
93 trace_id: self.trace_id.clone(),
94 parent_id: Some(span_id.into()),
95 baggage: self.baggage.clone(),
96 }
97 }
98
99 pub fn add_baggage(
101 &mut self,
102 key: impl Into<String>,
103 value: impl Into<String>,
104 ) -> Result<(), TraceError> {
105 let key = key.into();
106 let value = value.into();
107
108 if key.len() > MAX_BAGGAGE_ITEM_BYTES {
109 return Err(TraceError::BaggageItemTooLarge {
110 field: "key".into(),
111 len: key.len(),
112 max: MAX_BAGGAGE_ITEM_BYTES,
113 });
114 }
115 if value.len() > MAX_BAGGAGE_ITEM_BYTES {
116 return Err(TraceError::BaggageItemTooLarge {
117 field: "value".into(),
118 len: value.len(),
119 max: MAX_BAGGAGE_ITEM_BYTES,
120 });
121 }
122 if let Some(existing) = self.baggage.iter_mut().find(|entry| entry.key == key) {
123 existing.value = value;
124 return Ok(());
125 }
126 if self.baggage.len() >= MAX_BAGGAGE_ENTRIES {
127 return Err(TraceError::BaggageLimitExceeded {
128 max: MAX_BAGGAGE_ENTRIES,
129 });
130 }
131
132 self.baggage.push(BaggageEntry { key, value });
133 Ok(())
134 }
135
136 pub fn baggage_value(&self, key: &str) -> Option<&str> {
138 self.baggage
139 .iter()
140 .find(|e| e.key == key)
141 .map(|e| e.value.as_str())
142 }
143
144 pub fn to_traceparent(&self) -> Result<String, TraceError> {
155 let trace_id = if is_w3c_trace_id(&self.trace_id) {
156 self.trace_id.clone()
157 } else {
158 hash_to_w3c_trace_id(&self.trace_id)
159 };
160 let parent_id = match &self.parent_id {
161 Some(p) if is_w3c_span_id(p) => p.clone(),
162 Some(p) => {
163 return Err(TraceError::InvalidTraceparent {
164 reason: format!(
165 "parent_id must be 16 hex chars, got {} chars: '{}'",
166 p.len(),
167 p
168 ),
169 });
170 }
171 None => "0000000000000000".to_string(),
172 };
173 Ok(format!("00-{trace_id}-{parent_id}-01"))
174 }
175
176 pub fn from_traceparent(header: &str) -> Result<Self, TraceError> {
180 let parts: Vec<&str> = header.split('-').collect();
181 if parts.len() != 4 {
182 return Err(TraceError::InvalidTraceparent {
183 reason: format!("expected 4 dash-separated parts, got {}", parts.len()),
184 });
185 }
186
187 let version = parts[0];
188 if version != "00" {
189 return Err(TraceError::InvalidTraceparent {
190 reason: format!("unsupported version: {version}"),
191 });
192 }
193
194 let trace_id = parts[1].to_string();
195 if trace_id.len() != 32 || !trace_id.chars().all(|c| c.is_ascii_hexdigit()) {
196 return Err(TraceError::InvalidTraceparent {
197 reason: "trace-id must be 32 hex characters".into(),
198 });
199 }
200
201 let parent_id = parts[2].to_string();
202 if parent_id.len() != 16 || !parent_id.chars().all(|c| c.is_ascii_hexdigit()) {
203 return Err(TraceError::InvalidTraceparent {
204 reason: "parent-id must be 16 hex characters".into(),
205 });
206 }
207 let flags = parts[3];
208 if flags.len() != 2 || !flags.chars().all(|c| c.is_ascii_hexdigit()) {
209 return Err(TraceError::InvalidTraceparent {
210 reason: "trace-flags must be 2 hex characters".into(),
211 });
212 }
213
214 let parent = if parent_id == "0000000000000000" {
215 None
216 } else {
217 Some(parent_id)
218 };
219
220 Ok(Self {
221 trace_id,
222 parent_id: parent,
223 baggage: Vec::new(),
224 })
225 }
226}
227
228#[derive(Debug, Clone, PartialEq, Eq)]
230pub enum TraceError {
231 BaggageLimitExceeded { max: usize },
233 BaggageItemTooLarge {
235 field: String,
236 len: usize,
237 max: usize,
238 },
239 InvalidTraceparent { reason: String },
241}
242
243impl TraceError {
244 pub fn kind(&self) -> &'static str {
245 match self {
246 Self::BaggageLimitExceeded { .. } => "baggage_limit_exceeded",
247 Self::BaggageItemTooLarge { .. } => "baggage_item_too_large",
248 Self::InvalidTraceparent { .. } => "invalid_traceparent",
249 }
250 }
251}
252
253impl std::fmt::Display for TraceError {
254 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
255 match self {
256 Self::BaggageLimitExceeded { max } => {
257 write!(f, "baggage limit exceeded (max {max} entries)")
258 }
259 Self::BaggageItemTooLarge { field, len, max } => {
260 write!(f, "baggage {field} too large ({len} bytes, max {max})")
261 }
262 Self::InvalidTraceparent { reason } => {
263 write!(f, "invalid traceparent: {reason}")
264 }
265 }
266 }
267}
268
269impl std::error::Error for TraceError {}
270
271fn is_w3c_trace_id(id: &str) -> bool {
273 id.len() == 32 && id.chars().all(|c| c.is_ascii_hexdigit())
274}
275
276fn is_w3c_span_id(id: &str) -> bool {
278 id.len() == 16 && id.chars().all(|c| c.is_ascii_hexdigit())
279}
280
281pub fn hash_to_w3c_trace_id(legacy_id: &str) -> String {
291 let hash = blake3::hash(legacy_id.as_bytes());
292 let bytes = hash.as_bytes();
293 bytes[..16].iter().map(|b| format!("{b:02x}")).collect()
295}
296
297#[cfg(test)]
298mod tests {
299 use super::*;
300
301 #[test]
302 fn trace_ctx_generate() {
303 let ctx = TraceCtx::generate();
304 assert_eq!(ctx.trace_id.len(), 32);
305 assert!(ctx.parent_id.is_none());
306 assert!(ctx.baggage.is_empty());
307 }
308
309 #[test]
310 fn trace_ctx_child() {
311 let parent = TraceCtx::generate();
312 let child = parent.child("abcdef0123456789");
313 assert_eq!(child.trace_id, parent.trace_id);
314 assert_eq!(child.parent_id.as_deref(), Some("abcdef0123456789"));
315 }
316
317 #[test]
318 fn traceparent_roundtrip() {
319 let ctx = TraceCtx::from_trace_id("0af7651916cd43dd8448eb211c80319c")
320 .with_parent("b7ad6b7169203331");
321 let header = ctx.to_traceparent().unwrap();
322 assert_eq!(
323 header,
324 "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
325 );
326
327 let parsed = TraceCtx::from_traceparent(&header).unwrap();
328 assert_eq!(parsed.trace_id, "0af7651916cd43dd8448eb211c80319c");
329 assert_eq!(parsed.parent_id.as_deref(), Some("b7ad6b7169203331"));
330 }
331
332 #[test]
333 fn traceparent_no_parent() {
334 let ctx = TraceCtx::from_trace_id("0af7651916cd43dd8448eb211c80319c");
335 let header = ctx.to_traceparent().unwrap();
336 assert!(header.contains("0000000000000000"));
337
338 let parsed = TraceCtx::from_traceparent(&header).unwrap();
339 assert!(parsed.parent_id.is_none());
340 }
341
342 #[test]
343 fn traceparent_legacy_trace_id_uses_hash() {
344 let ctx = TraceCtx::from_trace_id("old-trace-abc");
345 let header = ctx.to_traceparent().unwrap();
346 let parts: Vec<&str> = header.split('-').collect();
348 assert_eq!(parts[0], "00");
349 assert_eq!(parts[1].len(), 32);
350 assert!(parts[1].chars().all(|c| c.is_ascii_hexdigit()));
351 let header2 = ctx.to_traceparent().unwrap();
353 assert_eq!(header, header2);
354 }
355
356 #[test]
357 fn traceparent_rejects_non_w3c_parent_id() {
358 let ctx = TraceCtx::from_trace_id("0af7651916cd43dd8448eb211c80319c")
359 .with_parent("not-hex-parent");
360 let err = ctx.to_traceparent().unwrap_err();
361 assert!(matches!(err, TraceError::InvalidTraceparent { .. }));
362 }
363
364 #[test]
365 fn hash_to_w3c_trace_id_is_deterministic() {
366 let id1 = super::hash_to_w3c_trace_id("legacy-id-123");
367 let id2 = super::hash_to_w3c_trace_id("legacy-id-123");
368 assert_eq!(id1, id2);
369 assert_eq!(id1.len(), 32);
370 assert!(id1.chars().all(|c| c.is_ascii_hexdigit()));
371 }
372
373 #[test]
374 fn hash_to_w3c_trace_id_different_inputs_differ() {
375 let id1 = super::hash_to_w3c_trace_id("legacy-id-123");
376 let id2 = super::hash_to_w3c_trace_id("legacy-id-456");
377 assert_ne!(id1, id2);
378 }
379
380 #[test]
381 fn traceparent_invalid_format() {
382 let err = TraceCtx::from_traceparent("bad").unwrap_err();
383 assert!(matches!(err, TraceError::InvalidTraceparent { .. }));
384 }
385
386 #[test]
387 fn traceparent_unsupported_version() {
388 let err =
389 TraceCtx::from_traceparent("01-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
390 .unwrap_err();
391 assert!(matches!(err, TraceError::InvalidTraceparent { .. }));
392 }
393
394 #[test]
395 fn traceparent_rejects_malformed_flags() {
396 let err =
397 TraceCtx::from_traceparent("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-zz")
398 .unwrap_err();
399 assert!(matches!(err, TraceError::InvalidTraceparent { .. }));
400 }
401
402 #[test]
403 fn baggage_add_and_get() {
404 let mut ctx = TraceCtx::generate();
405 ctx.add_baggage("env", "prod").unwrap();
406 assert_eq!(ctx.baggage_value("env"), Some("prod"));
407 assert_eq!(ctx.baggage_value("missing"), None);
408 }
409
410 #[test]
411 fn baggage_duplicate_key_updates_existing_entry() {
412 let mut ctx = TraceCtx::generate();
413 ctx.add_baggage("env", "dev").unwrap();
414 ctx.add_baggage("env", "prod").unwrap();
415 assert_eq!(ctx.baggage.len(), 1);
416 assert_eq!(ctx.baggage_value("env"), Some("prod"));
417 }
418
419 #[test]
420 fn baggage_duplicate_key_updates_even_when_entry_limit_is_full() {
421 let mut ctx = TraceCtx::generate();
422 for i in 0..MAX_BAGGAGE_ENTRIES {
423 ctx.add_baggage(format!("k{i}"), "v").unwrap();
424 }
425 ctx.add_baggage("k0", "updated").unwrap();
426 assert_eq!(ctx.baggage.len(), MAX_BAGGAGE_ENTRIES);
427 assert_eq!(ctx.baggage_value("k0"), Some("updated"));
428 }
429
430 #[test]
431 fn baggage_limit_enforced() {
432 let mut ctx = TraceCtx::generate();
433 for i in 0..MAX_BAGGAGE_ENTRIES {
434 ctx.add_baggage(format!("k{i}"), "v").unwrap();
435 }
436 let err = ctx.add_baggage("overflow", "v").unwrap_err();
437 assert!(matches!(err, TraceError::BaggageLimitExceeded { .. }));
438 }
439
440 #[test]
441 fn baggage_size_limit_enforced() {
442 let mut ctx = TraceCtx::generate();
443 let big_key = "x".repeat(MAX_BAGGAGE_ITEM_BYTES + 1);
444 let err = ctx.add_baggage(big_key, "v").unwrap_err();
445 assert!(matches!(err, TraceError::BaggageItemTooLarge { .. }));
446 }
447
448 #[test]
449 fn legacy_trace_id_compat() {
450 let ctx = TraceCtx::from_legacy_trace_id("old-trace-abc");
451 assert_eq!(ctx.to_legacy_trace_id(), "old-trace-abc");
452 }
453
454 #[test]
455 fn trace_ctx_serde_roundtrip() {
456 let mut ctx = TraceCtx::generate().with_parent("abcdef0123456789");
457 ctx.add_baggage("env", "test").unwrap();
458 let json = serde_json::to_string(&ctx).unwrap();
459 let back: TraceCtx = serde_json::from_str(&json).unwrap();
460 assert_eq!(back, ctx);
461 }
462}