rust_zero_core/
trace.rs

1use std::{
2    fmt,
3    sync::atomic::{AtomicU64, Ordering},
4    time::{SystemTime, UNIX_EPOCH},
5};
6
7static NEXT_ID: AtomicU64 = AtomicU64::new(1);
8
9/// W3C trace flags carried by a [`TraceContext`].
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct TraceFlags(u8);
12
13impl TraceFlags {
14    pub const NONE: Self = Self(0);
15    pub const SAMPLED: Self = Self(1);
16
17    pub fn is_sampled(self) -> bool {
18        self.0 & 1 == 1
19    }
20}
21
22/// A parsed W3C `traceparent` value.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct TraceContext {
25    trace_id: [u8; 16],
26    span_id: [u8; 8],
27    parent_span_id: Option<[u8; 8]>,
28    flags: TraceFlags,
29}
30
31impl TraceContext {
32    pub fn root(flags: TraceFlags) -> Self {
33        Self {
34            trace_id: next_trace_id(),
35            span_id: next_span_id(),
36            parent_span_id: None,
37            flags,
38        }
39    }
40
41    pub fn parse(value: &str) -> Result<Self, TraceContextError> {
42        let mut parts = value.split('-');
43        let version = parts.next().ok_or(TraceContextError)?;
44        let trace = parts.next().ok_or(TraceContextError)?;
45        let span = parts.next().ok_or(TraceContextError)?;
46        let flags = parts.next().ok_or(TraceContextError)?;
47        if parts.next().is_some()
48            || version != "00"
49            || trace.len() != 32
50            || span.len() != 16
51            || flags.len() != 2
52        {
53            return Err(TraceContextError);
54        }
55
56        let trace_id = decode_hex::<16>(trace)?;
57        let span_id = decode_hex::<8>(span)?;
58        let flags = TraceFlags(u8::from_str_radix(flags, 16).map_err(|_| TraceContextError)?);
59        if trace_id == [0; 16] || span_id == [0; 8] {
60            return Err(TraceContextError);
61        }
62
63        Ok(Self {
64            trace_id,
65            span_id,
66            parent_span_id: None,
67            flags,
68        })
69    }
70
71    pub fn child(&self) -> Self {
72        Self {
73            trace_id: self.trace_id,
74            span_id: next_span_id(),
75            parent_span_id: Some(self.span_id),
76            flags: self.flags,
77        }
78    }
79
80    pub fn trace_id(&self) -> String {
81        encode_hex(&self.trace_id)
82    }
83
84    pub fn span_id(&self) -> String {
85        encode_hex(&self.span_id)
86    }
87
88    pub fn parent_span_id(&self) -> Option<String> {
89        self.parent_span_id.as_ref().map(|id| encode_hex(id))
90    }
91
92    pub fn flags(&self) -> TraceFlags {
93        self.flags
94    }
95
96    pub fn traceparent(&self) -> String {
97        format!(
98            "00-{}-{}-{:02x}",
99            encode_hex(&self.trace_id),
100            encode_hex(&self.span_id),
101            self.flags.0
102        )
103    }
104}
105
106/// Returned when a W3C trace context is malformed or uses an unsupported version.
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub struct TraceContextError;
109
110impl fmt::Display for TraceContextError {
111    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
112        formatter.write_str("invalid W3C traceparent value")
113    }
114}
115
116impl std::error::Error for TraceContextError {}
117
118fn next_trace_id() -> [u8; 16] {
119    let mut id = [0; 16];
120    let time = SystemTime::now()
121        .duration_since(UNIX_EPOCH)
122        .unwrap_or_default()
123        .as_nanos();
124    id.copy_from_slice(&time.to_be_bytes());
125    let sequence = NEXT_ID.fetch_add(1, Ordering::Relaxed).to_be_bytes();
126    for (target, source) in id[8..].iter_mut().zip(sequence) {
127        *target ^= source;
128    }
129    if id == [0; 16] {
130        id[15] = 1;
131    }
132    id
133}
134
135fn next_span_id() -> [u8; 8] {
136    let mut id = NEXT_ID.fetch_add(1, Ordering::Relaxed).to_be_bytes();
137    if id == [0; 8] {
138        id[7] = 1;
139    }
140    id
141}
142
143fn decode_hex<const N: usize>(value: &str) -> Result<[u8; N], TraceContextError> {
144    let mut output = [0; N];
145    for (index, byte) in output.iter_mut().enumerate() {
146        *byte = u8::from_str_radix(&value[index * 2..index * 2 + 2], 16)
147            .map_err(|_| TraceContextError)?;
148    }
149    Ok(output)
150}
151
152fn encode_hex(bytes: &[u8]) -> String {
153    use fmt::Write;
154    let mut output = String::with_capacity(bytes.len() * 2);
155    for byte in bytes {
156        write!(&mut output, "{byte:02x}").expect("writing to a String cannot fail");
157    }
158    output
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    #[test]
166    fn parses_and_creates_child_trace_contexts() {
167        let parent =
168            TraceContext::parse("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01").unwrap();
169        let child = parent.child();
170
171        assert_eq!(child.trace_id(), parent.trace_id());
172        assert_eq!(child.parent_span_id(), Some(parent.span_id()));
173        assert!(child.flags().is_sampled());
174        assert!(TraceContext::parse(&child.traceparent()).is_ok());
175    }
176
177    #[test]
178    fn rejects_zero_and_malformed_identifiers() {
179        assert!(
180            TraceContext::parse("00-00000000000000000000000000000000-00f067aa0ba902b7-01").is_err()
181        );
182        assert!(TraceContext::parse("not-a-trace").is_err());
183    }
184}
rust_zero_core/trace.rs

rust_zero_core/
trace.rs
