Skip to main content

stack_ids/
trace.rs

1//! Trace context primitives for cross-crate and cross-boundary correlation.
2//!
3//! ## Trace law (from MASTER_SUPPORTING_DELTA ยง4.3)
4//!
5//! - `TraceCtx` is canonical in-process.
6//! - Across queue/network boundaries, preserve W3C trace context plus bounded baggage only.
7//! - No large opaque blobs. No sensitive payloads in baggage.
8//!
9//! ## W3C Trace Context
10//!
11//! The `traceparent` header format is:
12//! ```text
13//! {version}-{trace-id}-{parent-id}-{trace-flags}
14//! 00-{32 hex chars}-{16 hex chars}-{2 hex chars}
15//! ```
16
17use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19
20/// In-process trace context for cross-crate correlation.
21///
22/// Wraps a W3C-compatible trace ID and optional parent span ID.
23/// Additional bounded baggage can be attached for cross-boundary metadata.
24#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
25pub struct TraceCtx {
26    /// The trace ID (W3C: 32 hex chars, or any opaque string for legacy compat).
27    pub trace_id: String,
28    /// Optional parent span ID (W3C: 16 hex chars).
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub parent_id: Option<String>,
31    /// Bounded baggage for cross-boundary metadata.
32    /// Keys and values must be short ASCII strings. No sensitive data.
33    #[serde(default, skip_serializing_if = "Vec::is_empty")]
34    pub baggage: Vec<BaggageEntry>,
35}
36
37/// Maximum number of baggage entries allowed.
38pub const MAX_BAGGAGE_ENTRIES: usize = 16;
39
40/// Maximum byte length for a single baggage key or value.
41pub const MAX_BAGGAGE_ITEM_BYTES: usize = 256;
42
43/// A single baggage entry (key-value pair).
44#[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    /// Create a new trace context with a generated trace ID (UUID v4 hex).
52    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    /// Create from a raw trace ID string.
62    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    /// Create from a legacy `TraceId(String)` value for migration compatibility.
71    ///
72    /// Phase status: compatibility / migration-only
73    pub fn from_legacy_trace_id(legacy_id: impl Into<String>) -> Self {
74        Self::from_trace_id(legacy_id)
75    }
76
77    /// Extract the trace ID as a string for legacy interop.
78    ///
79    /// Phase status: compatibility / migration-only
80    pub fn to_legacy_trace_id(&self) -> &str {
81        &self.trace_id
82    }
83
84    /// Set the parent span ID.
85    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    /// Create a child context: same trace ID, new parent.
91    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    /// Add a baggage entry. Returns Err if limits are exceeded.
100    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    /// Get a baggage value by key.
137    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    /// Format as a W3C traceparent header.
145    ///
146    /// Format: `00-{trace_id}-{parent_id}-01`
147    ///
148    /// If the trace ID is not exactly 32 hex chars (e.g., a legacy opaque string),
149    /// it is converted via deterministic BLAKE3 hash truncation using
150    /// [`hash_to_w3c_trace_id`]. Padding and truncation are forbidden.
151    ///
152    /// If no parent ID exists, a zero span ID is used.
153    /// Non-compliant parent IDs (not 16 hex chars) are rejected.
154    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    /// Parse from a W3C traceparent header.
177    ///
178    /// Format: `{version}-{trace_id}-{parent_id}-{flags}`
179    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/// Errors related to trace context operations.
229#[derive(Debug, Clone, PartialEq, Eq)]
230pub enum TraceError {
231    /// Too many baggage entries.
232    BaggageLimitExceeded { max: usize },
233    /// A single baggage key or value exceeds the size limit.
234    BaggageItemTooLarge {
235        field: String,
236        len: usize,
237        max: usize,
238    },
239    /// Invalid traceparent header format.
240    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
271/// Check if a string is a valid W3C trace ID (exactly 32 hex chars).
272fn is_w3c_trace_id(id: &str) -> bool {
273    id.len() == 32 && id.chars().all(|c| c.is_ascii_hexdigit())
274}
275
276/// Check if a string is a valid W3C span ID (exactly 16 hex chars).
277fn is_w3c_span_id(id: &str) -> bool {
278    id.len() == 16 && id.chars().all(|c| c.is_ascii_hexdigit())
279}
280
281/// Convert a non-W3C trace ID to a W3C-compliant 32-hex-char wire ID
282/// via deterministic BLAKE3 hash truncation.
283///
284/// The same input always produces the same output.
285/// The original legacy identifier should be preserved in baggage
286/// under the key `legacy_trace_id` by the caller if round-trip is needed.
287///
288/// Padding and truncation are forbidden. This is the only canonical
289/// conversion path for non-W3C trace IDs.
290pub 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    // Take first 16 bytes (128 bits) = 32 hex chars, matching W3C trace-id length.
294    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        // The legacy ID is hashed, not padded/truncated.
347        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        // Deterministic: same input always produces same output.
352        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}