Skip to main content

rmcp_traces/
trace_context.rs

1use std::{collections::BTreeSet, error::Error, fmt};
2
3use rmcp::model::MetaObject;
4use serde_json::Value;
5
6pub const TRACEPARENT_KEY: &str = "traceparent";
7pub const TRACESTATE_KEY: &str = "tracestate";
8pub const BAGGAGE_KEY: &str = "baggage";
9
10const TRACEPARENT_V00_LEN: usize = 55;
11const TRACEPARENT_VERSION_END: usize = 2;
12const TRACEPARENT_TRACE_ID_START: usize = 3;
13const TRACEPARENT_TRACE_ID_END: usize = 35;
14const TRACEPARENT_SPAN_ID_START: usize = 36;
15const TRACEPARENT_SPAN_ID_END: usize = 52;
16const TRACEPARENT_FLAGS_START: usize = 53;
17const TRACEPARENT_FLAGS_END: usize = 55;
18const TRACEPARENT_NEXT_SEPARATOR: usize = 55;
19const MAX_TRACESTATE_MEMBERS: usize = 32;
20
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub enum TraceTrust {
23    Untrusted,
24    Trusted,
25}
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub struct TraceLimits {
29    pub max_traceparent_len: usize,
30    pub max_tracestate_len: usize,
31    pub max_baggage_len: usize,
32    pub max_baggage_members: usize,
33}
34
35impl Default for TraceLimits {
36    fn default() -> Self {
37        Self {
38            max_traceparent_len: 512,
39            max_tracestate_len: 512,
40            max_baggage_len: 8 * 1024,
41            max_baggage_members: 64,
42        }
43    }
44}
45
46#[derive(Clone, PartialEq, Eq)]
47pub(crate) struct TraceParent {
48    trace_id: String,
49    span_id: String,
50    sampled: bool,
51}
52
53impl TraceParent {
54    #[cfg(test)]
55    fn parse(value: &str) -> Result<Self, TraceParseError> {
56        let limits = TraceLimits::default();
57        if value.len() > limits.max_traceparent_len {
58            return Err(TraceParseError::ValueTooLong {
59                field: TRACEPARENT_KEY,
60                actual: value.len(),
61                max: limits.max_traceparent_len,
62            });
63        }
64        parse_traceparent(value)
65    }
66
67    #[cfg(test)]
68    fn trace_id(&self) -> &str {
69        &self.trace_id
70    }
71
72    #[cfg(test)]
73    fn span_id(&self) -> &str {
74        &self.span_id
75    }
76
77    fn sampled(&self) -> bool {
78        self.sampled
79    }
80
81    fn trace_id_short(&self) -> &str {
82        &self.trace_id[..8]
83    }
84
85    fn span_id_short(&self) -> &str {
86        &self.span_id[..8]
87    }
88}
89
90impl fmt::Debug for TraceParent {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        f.debug_struct("TraceParent")
93            .field("trace_id", &self.trace_id_short())
94            .field("span_id", &self.span_id_short())
95            .field("sampled", &self.sampled)
96            .finish()
97    }
98}
99
100#[cfg(test)]
101#[derive(Clone, PartialEq, Eq)]
102struct TraceContext {
103    traceparent: TraceParent,
104    tracestate: Option<String>,
105    baggage: Option<String>,
106    trust: TraceTrust,
107}
108
109#[cfg(test)]
110impl TraceContext {
111    fn from_meta(meta: &MetaObject, trust: TraceTrust) -> Result<Option<Self>, TraceParseError> {
112        Self::from_meta_with_limits(meta, trust, TraceLimits::default())
113    }
114
115    fn from_meta_with_limits(
116        meta: &MetaObject,
117        trust: TraceTrust,
118        limits: TraceLimits,
119    ) -> Result<Option<Self>, TraceParseError> {
120        let Some(traceparent) = parse_meta_traceparent(meta, limits)? else {
121            return Ok(None);
122        };
123        let tracestate =
124            bounded_optional_meta_string(meta, TRACESTATE_KEY, limits.max_tracestate_len)?;
125        validate_tracestate(tracestate.as_deref())?;
126        let baggage = bounded_optional_meta_string(meta, BAGGAGE_KEY, limits.max_baggage_len)?;
127        validate_baggage(baggage.as_deref(), limits.max_baggage_members)?;
128        Ok(Some(Self {
129            traceparent,
130            tracestate,
131            baggage,
132            trust,
133        }))
134    }
135
136    fn summary(&self) -> TraceSummary {
137        TraceSummary::from_context(self)
138    }
139}
140
141#[cfg(test)]
142impl fmt::Debug for TraceContext {
143    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144        self.summary().fmt(f)
145    }
146}
147
148#[derive(Clone, Debug, PartialEq, Eq)]
149pub struct TraceSummary {
150    trace_id_prefix: Option<String>,
151    span_id_prefix: Option<String>,
152    sampled: Option<bool>,
153    trust: TraceTrust,
154    has_tracestate: bool,
155    baggage_member_count: usize,
156    sensitive_baggage_member_count: usize,
157    invalid_reasons: Vec<String>,
158}
159
160impl TraceSummary {
161    pub fn absent() -> Self {
162        Self::absent_with_trust(TraceTrust::Untrusted)
163    }
164
165    pub(crate) fn absent_with_trust(trust: TraceTrust) -> Self {
166        Self {
167            trace_id_prefix: None,
168            span_id_prefix: None,
169            sampled: None,
170            trust,
171            has_tracestate: false,
172            baggage_member_count: 0,
173            sensitive_baggage_member_count: 0,
174            invalid_reasons: Vec::new(),
175        }
176    }
177
178    pub fn invalid(error: &TraceParseError) -> Self {
179        let mut summary = Self::absent();
180        summary.record_invalid(error);
181        summary
182    }
183
184    pub fn from_meta(meta: &MetaObject, trust: TraceTrust) -> Self {
185        Self::from_meta_with_limits(meta, trust, TraceLimits::default())
186    }
187
188    pub fn from_meta_with_limits(
189        meta: &MetaObject,
190        trust: TraceTrust,
191        limits: TraceLimits,
192    ) -> Self {
193        let (mut summary, has_valid_traceparent) = match parse_meta_traceparent(meta, limits) {
194            Ok(Some(traceparent)) => (Self::from_traceparent(&traceparent, trust), true),
195            Ok(None) => (Self::absent_with_trust(trust), false),
196            Err(error) => {
197                let mut summary = Self::absent_with_trust(trust);
198                summary.record_invalid(&error);
199                (summary, false)
200            }
201        };
202        match bounded_optional_meta_str(meta, TRACESTATE_KEY, limits.max_tracestate_len) {
203            Ok(Some(tracestate)) if has_valid_traceparent => {
204                match validate_tracestate(Some(tracestate)) {
205                    Ok(()) => summary.has_tracestate = true,
206                    Err(error) => summary.record_invalid(&error),
207                }
208            }
209            Ok(Some(_)) => {
210                summary.record_invalid(&TraceParseError::TraceStateRequiresTraceParent);
211            }
212            Ok(None) => {}
213            Err(error) => summary.record_invalid(&error),
214        };
215        match bounded_optional_meta_str(meta, BAGGAGE_KEY, limits.max_baggage_len) {
216            Ok(baggage) => match validate_baggage(baggage, limits.max_baggage_members) {
217                Ok(()) => {
218                    let (baggage_member_count, sensitive_baggage_member_count) =
219                        summarize_baggage(baggage);
220                    summary.baggage_member_count = baggage_member_count;
221                    summary.sensitive_baggage_member_count = sensitive_baggage_member_count;
222                }
223                Err(error) => summary.record_invalid(&error),
224            },
225            Err(error) => {
226                summary.record_invalid(&error);
227            }
228        };
229        summary
230    }
231
232    #[cfg(test)]
233    fn from_context(context: &TraceContext) -> Self {
234        let (baggage_member_count, sensitive_baggage_member_count) =
235            summarize_baggage(context.baggage.as_deref());
236        let mut summary = Self::from_traceparent(&context.traceparent, context.trust);
237        summary.has_tracestate = context.tracestate.is_some();
238        summary.baggage_member_count = baggage_member_count;
239        summary.sensitive_baggage_member_count = sensitive_baggage_member_count;
240        summary
241    }
242
243    fn from_traceparent(traceparent: &TraceParent, trust: TraceTrust) -> Self {
244        Self {
245            trace_id_prefix: Some(traceparent.trace_id_short().to_owned()),
246            span_id_prefix: Some(traceparent.span_id_short().to_owned()),
247            sampled: Some(traceparent.sampled()),
248            trust,
249            has_tracestate: false,
250            baggage_member_count: 0,
251            sensitive_baggage_member_count: 0,
252            invalid_reasons: Vec::new(),
253        }
254    }
255
256    #[cfg(feature = "http")]
257    pub(crate) fn from_valid_traceparent(traceparent: &TraceParent, trust: TraceTrust) -> Self {
258        Self::from_traceparent(traceparent, trust)
259    }
260
261    fn record_invalid(&mut self, error: &TraceParseError) {
262        self.invalid_reasons.push(error.safe_reason());
263    }
264
265    #[cfg(feature = "http")]
266    pub(crate) fn record_invalid_reason(&mut self, error: &TraceParseError) {
267        self.record_invalid(error);
268    }
269
270    #[cfg(feature = "http")]
271    pub(crate) fn set_has_tracestate_for_http(&mut self) {
272        self.has_tracestate = true;
273    }
274
275    #[cfg(feature = "http")]
276    pub(crate) fn set_baggage_counts_for_http(&mut self, baggage: Option<&str>) {
277        let (baggage_member_count, sensitive_baggage_member_count) = summarize_baggage(baggage);
278        self.baggage_member_count = baggage_member_count;
279        self.sensitive_baggage_member_count = sensitive_baggage_member_count;
280    }
281
282    pub fn trace_id_prefix(&self) -> Option<&str> {
283        self.trace_id_prefix.as_deref()
284    }
285
286    pub fn span_id_prefix(&self) -> Option<&str> {
287        self.span_id_prefix.as_deref()
288    }
289
290    pub fn sampled(&self) -> Option<bool> {
291        self.sampled
292    }
293
294    pub fn trust(&self) -> TraceTrust {
295        self.trust
296    }
297
298    pub fn has_tracestate(&self) -> bool {
299        self.has_tracestate
300    }
301
302    pub fn baggage_member_count(&self) -> usize {
303        self.baggage_member_count
304    }
305
306    pub fn sensitive_baggage_member_count(&self) -> usize {
307        self.sensitive_baggage_member_count
308    }
309
310    pub fn invalid_reasons(&self) -> &[String] {
311        &self.invalid_reasons
312    }
313
314    pub fn invalid_count(&self) -> usize {
315        self.invalid_reasons.len()
316    }
317}
318
319#[derive(Clone, Debug, PartialEq, Eq)]
320pub enum TraceParseError {
321    NonStringMeta {
322        field: &'static str,
323    },
324    ValueTooLong {
325        field: &'static str,
326        actual: usize,
327        max: usize,
328    },
329    TooManyBaggageMembers {
330        actual: usize,
331        max: usize,
332    },
333    TooManyTraceStateMembers {
334        actual: usize,
335        max: usize,
336    },
337    InvalidTraceParentLength {
338        actual: usize,
339    },
340    InvalidTraceParentFormat,
341    UnsupportedVersion,
342    InvalidTraceId,
343    InvalidSpanId,
344    TraceStateRequiresTraceParent,
345    InvalidTraceState,
346    DuplicateTraceStateKey,
347    InvalidBaggageMember,
348    #[cfg(feature = "http")]
349    MultipleHeaderValues {
350        field: &'static str,
351    },
352    #[cfg(feature = "http")]
353    InvalidHeaderValue {
354        field: &'static str,
355    },
356}
357
358impl TraceParseError {
359    pub fn safe_reason(&self) -> String {
360        match self {
361            Self::NonStringMeta { field } => format!("{field} was not a string"),
362            Self::ValueTooLong { field, actual, max } => {
363                format!("{field} exceeded {max} bytes (actual {actual})")
364            }
365            Self::TooManyBaggageMembers { actual, max } => {
366                format!("baggage exceeded {max} members (actual at least {actual})")
367            }
368            Self::TooManyTraceStateMembers { actual, max } => {
369                format!("tracestate exceeded {max} members (actual at least {actual})")
370            }
371            Self::InvalidTraceParentLength { actual } => {
372                format!("traceparent length was {actual}, expected 55")
373            }
374            Self::InvalidTraceParentFormat => "traceparent format was invalid".to_owned(),
375            Self::UnsupportedVersion => "traceparent version was unsupported".to_owned(),
376            Self::InvalidTraceId => "traceparent trace id was invalid".to_owned(),
377            Self::InvalidSpanId => "traceparent span id was invalid".to_owned(),
378            Self::TraceStateRequiresTraceParent => {
379                "tracestate requires a valid traceparent".to_owned()
380            }
381            Self::InvalidTraceState => "tracestate format was invalid".to_owned(),
382            Self::DuplicateTraceStateKey => "tracestate contained a duplicate key".to_owned(),
383            Self::InvalidBaggageMember => "baggage member format was invalid".to_owned(),
384            #[cfg(feature = "http")]
385            Self::MultipleHeaderValues { field } => {
386                format!("{field} had multiple header values")
387            }
388            #[cfg(feature = "http")]
389            Self::InvalidHeaderValue { field } => {
390                format!("{field} header value was not visible ASCII")
391            }
392        }
393    }
394}
395
396impl fmt::Display for TraceParseError {
397    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
398        f.write_str(&self.safe_reason())
399    }
400}
401
402impl Error for TraceParseError {}
403
404fn optional_meta_str<'a>(
405    meta: &'a MetaObject,
406    field: &'static str,
407) -> Result<Option<&'a str>, TraceParseError> {
408    match meta.get(field) {
409        None => Ok(None),
410        Some(Value::String(value)) => Ok(Some(value.as_str())),
411        Some(_) => Err(TraceParseError::NonStringMeta { field }),
412    }
413}
414
415#[cfg(test)]
416fn bounded_optional_meta_string(
417    meta: &MetaObject,
418    field: &'static str,
419    max: usize,
420) -> Result<Option<String>, TraceParseError> {
421    bounded_optional_meta_str(meta, field, max).map(|value| value.map(str::to_owned))
422}
423
424fn bounded_optional_meta_str<'a>(
425    meta: &'a MetaObject,
426    field: &'static str,
427    max: usize,
428) -> Result<Option<&'a str>, TraceParseError> {
429    let Some(value) = optional_meta_str(meta, field)? else {
430        return Ok(None);
431    };
432    if value.len() > max {
433        return Err(TraceParseError::ValueTooLong {
434            field,
435            actual: value.len(),
436            max,
437        });
438    }
439    Ok(Some(value))
440}
441
442fn parse_meta_traceparent(
443    meta: &MetaObject,
444    limits: TraceLimits,
445) -> Result<Option<TraceParent>, TraceParseError> {
446    let Some(traceparent) = optional_meta_str(meta, TRACEPARENT_KEY)? else {
447        return Ok(None);
448    };
449    if traceparent.len() > limits.max_traceparent_len {
450        return Err(TraceParseError::ValueTooLong {
451            field: TRACEPARENT_KEY,
452            actual: traceparent.len(),
453            max: limits.max_traceparent_len,
454        });
455    }
456    parse_traceparent(traceparent).map(Some)
457}
458
459#[cfg(feature = "http")]
460pub(crate) fn parse_traceparent_value(
461    value: &str,
462    limits: TraceLimits,
463) -> Result<TraceParent, TraceParseError> {
464    if value.len() > limits.max_traceparent_len {
465        return Err(TraceParseError::ValueTooLong {
466            field: TRACEPARENT_KEY,
467            actual: value.len(),
468            max: limits.max_traceparent_len,
469        });
470    }
471    parse_traceparent(value)
472}
473
474fn parse_traceparent(value: &str) -> Result<TraceParent, TraceParseError> {
475    if value.len() < TRACEPARENT_V00_LEN {
476        return Err(TraceParseError::InvalidTraceParentLength {
477            actual: value.len(),
478        });
479    }
480    if !value.is_ascii() {
481        return Err(TraceParseError::InvalidTraceParentFormat);
482    }
483    let bytes = value.as_bytes();
484    if bytes[TRACEPARENT_VERSION_END] != b'-'
485        || bytes[TRACEPARENT_TRACE_ID_END] != b'-'
486        || bytes[TRACEPARENT_SPAN_ID_END] != b'-'
487    {
488        return Err(TraceParseError::InvalidTraceParentFormat);
489    }
490    let version = &value[..TRACEPARENT_VERSION_END];
491    let trace_id = &value[TRACEPARENT_TRACE_ID_START..TRACEPARENT_TRACE_ID_END];
492    let span_id = &value[TRACEPARENT_SPAN_ID_START..TRACEPARENT_SPAN_ID_END];
493    let flags = &value[TRACEPARENT_FLAGS_START..TRACEPARENT_FLAGS_END];
494    if !is_lower_hex(version)
495        || !is_lower_hex(trace_id)
496        || !is_lower_hex(span_id)
497        || !is_lower_hex(flags)
498    {
499        return Err(TraceParseError::InvalidTraceParentFormat);
500    }
501    if version == "ff" {
502        return Err(TraceParseError::UnsupportedVersion);
503    }
504    if version == "00" && value.len() != TRACEPARENT_V00_LEN {
505        return Err(TraceParseError::InvalidTraceParentLength {
506            actual: value.len(),
507        });
508    }
509    if version != "00"
510        && value.len() > TRACEPARENT_V00_LEN
511        && bytes[TRACEPARENT_NEXT_SEPARATOR] != b'-'
512    {
513        return Err(TraceParseError::InvalidTraceParentFormat);
514    }
515    if trace_id.bytes().all(|b| b == b'0') {
516        return Err(TraceParseError::InvalidTraceId);
517    }
518    if span_id.bytes().all(|b| b == b'0') {
519        return Err(TraceParseError::InvalidSpanId);
520    }
521    let flag_byte =
522        u8::from_str_radix(flags, 16).map_err(|_| TraceParseError::InvalidTraceParentFormat)?;
523    Ok(TraceParent {
524        trace_id: trace_id.to_owned(),
525        span_id: span_id.to_owned(),
526        sampled: flag_byte & 0x01 == 0x01,
527    })
528}
529
530fn is_lower_hex(value: &str) -> bool {
531    value
532        .bytes()
533        .all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase())
534}
535
536fn summarize_baggage(baggage: Option<&str>) -> (usize, usize) {
537    let Some(baggage) = baggage else {
538        return (0, 0);
539    };
540    let mut total = 0;
541    let mut sensitive = 0;
542    for key in baggage_keys(baggage) {
543        total += 1;
544        if is_sensitive_key(key) {
545            sensitive += 1;
546        }
547    }
548    (total, sensitive)
549}
550
551fn validate_tracestate(tracestate: Option<&str>) -> Result<(), TraceParseError> {
552    let Some(tracestate) = tracestate else {
553        return Ok(());
554    };
555    let mut keys = BTreeSet::new();
556    let mut total = 0;
557    for member in tracestate.split(',') {
558        let member = member.trim();
559        if member.is_empty() {
560            return Err(TraceParseError::InvalidTraceState);
561        }
562        total += 1;
563        if total > MAX_TRACESTATE_MEMBERS {
564            return Err(TraceParseError::TooManyTraceStateMembers {
565                actual: total,
566                max: MAX_TRACESTATE_MEMBERS,
567            });
568        }
569        let Some((key, value)) = member.split_once('=') else {
570            return Err(TraceParseError::InvalidTraceState);
571        };
572        let key = key.trim();
573        if !is_valid_tracestate_key(key) || !is_valid_tracestate_value(value) {
574            return Err(TraceParseError::InvalidTraceState);
575        }
576        if !keys.insert(key) {
577            return Err(TraceParseError::DuplicateTraceStateKey);
578        }
579    }
580    Ok(())
581}
582
583fn validate_baggage(baggage: Option<&str>, max: usize) -> Result<(), TraceParseError> {
584    let Some(baggage) = baggage else {
585        return Ok(());
586    };
587    let mut total = 0;
588    for member in baggage
589        .split(',')
590        .map(str::trim)
591        .filter(|member| !member.is_empty())
592    {
593        total += 1;
594        if total > max {
595            return Err(TraceParseError::TooManyBaggageMembers { actual: total, max });
596        }
597        let Some((key, rest)) = member.split_once('=') else {
598            return Err(TraceParseError::InvalidBaggageMember);
599        };
600        if !is_valid_baggage_key(key.trim()) {
601            return Err(TraceParseError::InvalidBaggageMember);
602        }
603        let mut value_and_props = rest.split(';');
604        let value = value_and_props.next().unwrap_or("").trim();
605        if !is_valid_baggage_value(value) {
606            return Err(TraceParseError::InvalidBaggageMember);
607        }
608        for property in value_and_props {
609            let property = property.trim();
610            if property.is_empty() {
611                return Err(TraceParseError::InvalidBaggageMember);
612            }
613            if let Some((property_key, property_value)) = property.split_once('=') {
614                if !is_valid_baggage_key(property_key.trim())
615                    || !is_valid_baggage_value(property_value.trim())
616                {
617                    return Err(TraceParseError::InvalidBaggageMember);
618                }
619            } else if !is_valid_baggage_key(property) {
620                return Err(TraceParseError::InvalidBaggageMember);
621            }
622        }
623    }
624    Ok(())
625}
626
627#[cfg(feature = "http")]
628pub(crate) fn validate_tracestate_value(tracestate: Option<&str>) -> Result<(), TraceParseError> {
629    validate_tracestate(tracestate)
630}
631
632#[cfg(feature = "http")]
633pub(crate) fn validate_baggage_value(
634    baggage: Option<&str>,
635    max_members: usize,
636) -> Result<(), TraceParseError> {
637    validate_baggage(baggage, max_members)
638}
639
640fn baggage_keys(baggage: &str) -> impl Iterator<Item = &str> {
641    baggage.split(',').filter_map(|member| {
642        let (key, _) = member.split_once('=')?;
643        let key = key.trim();
644        (!key.is_empty()).then_some(key)
645    })
646}
647
648fn is_sensitive_key(key: &str) -> bool {
649    const SENSITIVE_KEYS: &[&str] = &[
650        "authorization",
651        "cookie",
652        "setcookie",
653        "password",
654        "secret",
655        "token",
656        "accesstoken",
657        "refreshtoken",
658        "apikey",
659        "xapikey",
660        "privatekey",
661        "session",
662        "sessionid",
663    ];
664    SENSITIVE_KEYS
665        .iter()
666        .any(|sensitive_key| normalized_ascii_key_eq(key, sensitive_key))
667}
668
669fn normalized_ascii_key_eq(key: &str, expected: &str) -> bool {
670    let mut key_bytes = key
671        .bytes()
672        .filter(|byte| byte.is_ascii_alphanumeric())
673        .map(|byte| byte.to_ascii_lowercase());
674    let mut expected_bytes = expected.bytes();
675    loop {
676        match (key_bytes.next(), expected_bytes.next()) {
677            (Some(left), Some(right)) if left == right => {}
678            (None, None) => return true,
679            _ => return false,
680        }
681    }
682}
683
684fn is_valid_tracestate_key(key: &str) -> bool {
685    !key.is_empty()
686        && key.is_ascii()
687        && key.bytes().all(|byte| {
688            byte.is_ascii_lowercase()
689                || byte.is_ascii_digit()
690                || matches!(byte, b'_' | b'-' | b'*' | b'/' | b'@')
691        })
692}
693
694fn is_valid_tracestate_value(value: &str) -> bool {
695    !value.is_empty()
696        && value.is_ascii()
697        && value.trim() == value
698        && value
699            .bytes()
700            .all(|byte| matches!(byte, 0x20..=0x2b | 0x2d..=0x3c | 0x3e..=0x7e))
701}
702
703fn is_valid_baggage_key(key: &str) -> bool {
704    !key.is_empty()
705        && key.is_ascii()
706        && key.bytes().all(|byte| {
707            byte.is_ascii_alphanumeric()
708                || matches!(
709                    byte,
710                    b'!' | b'#'
711                        | b'$'
712                        | b'%'
713                        | b'&'
714                        | b'\''
715                        | b'*'
716                        | b'+'
717                        | b'-'
718                        | b'.'
719                        | b'^'
720                        | b'_'
721                        | b'`'
722                        | b'|'
723                        | b'~'
724                )
725        })
726}
727
728fn is_valid_baggage_value(value: &str) -> bool {
729    value
730        .bytes()
731        .all(|byte| matches!(byte, 0x21 | 0x23..=0x2b | 0x2d..=0x3a | 0x3c..=0x5b | 0x5d..=0x7e))
732}
733
734#[cfg(test)]
735mod tests {
736    use super::*;
737
738    const VALID_TRACEPARENT: &str = "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01";
739
740    #[test]
741    fn trace_context_parses_meta_and_summarizes_safely() {
742        let mut meta = MetaObject::new();
743        meta.set_traceparent(VALID_TRACEPARENT);
744        meta.set_tracestate("vendor=value");
745        meta.set_baggage("region=us-east-1,accessToken=super-secret-token");
746
747        let context = TraceContext::from_meta(&meta, TraceTrust::Untrusted)
748            .expect("valid trace metadata")
749            .expect("trace context exists");
750
751        let summary = context.summary();
752
753        assert_eq!(summary.trace_id_prefix(), Some("0af76519"));
754        assert_eq!(summary.span_id_prefix(), Some("00f067aa"));
755        assert_eq!(summary.sampled(), Some(true));
756        assert_eq!(summary.trust(), TraceTrust::Untrusted);
757        assert!(summary.has_tracestate());
758        assert_eq!(summary.baggage_member_count(), 2);
759        assert_eq!(summary.sensitive_baggage_member_count(), 1);
760        assert_eq!(summary.invalid_count(), 0);
761    }
762
763    #[test]
764    fn malformed_traceparents_are_rejected() {
765        for value in [
766            "",
767            "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7",
768            "ff-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01",
769            "00-00000000000000000000000000000000-00f067aa0ba902b7-01",
770            "00-0af7651916cd43dd8448eb211c80319c-0000000000000000-01",
771            "00-0AF7651916CD43DD8448EB211C80319C-00f067aa0ba902b7-01",
772            "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-zz",
773        ] {
774            assert!(
775                TraceParent::parse(value).is_err(),
776                "{value} should be rejected"
777            );
778        }
779    }
780
781    #[test]
782    fn non_ascii_traceparents_are_rejected_without_panicking() {
783        let value = format!("{}\u{00e9}", &VALID_TRACEPARENT[..54]);
784        let result = std::panic::catch_unwind(|| TraceParent::parse(&value));
785
786        assert!(result.is_ok(), "non-ASCII input must not panic");
787        assert!(matches!(
788            result.unwrap(),
789            Err(TraceParseError::InvalidTraceParentFormat)
790        ));
791
792        let higher_version = "01-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01-\u{00e9}";
793        assert!(matches!(
794            TraceParent::parse(higher_version),
795            Err(TraceParseError::InvalidTraceParentFormat)
796        ));
797    }
798
799    #[test]
800    fn traceparent_version_rules_cover_v00_and_higher_version_bounds() {
801        assert!(matches!(
802            TraceParent::parse(&format!("{VALID_TRACEPARENT}-extra")),
803            Err(TraceParseError::InvalidTraceParentLength { actual }) if actual == 61
804        ));
805
806        let higher_base = "01-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01";
807        let max_extra_len = 512 - higher_base.len() - 1;
808        let max_len_value = format!("{higher_base}-{}", "a".repeat(max_extra_len));
809        assert_eq!(max_len_value.len(), 512);
810        TraceParent::parse(&max_len_value).expect("512-byte higher version should be accepted");
811
812        let too_long = format!("{higher_base}-{}", "a".repeat(max_extra_len + 1));
813        assert_eq!(too_long.len(), 513);
814        assert!(matches!(
815            TraceParent::parse(&too_long),
816            Err(TraceParseError::ValueTooLong {
817                field: TRACEPARENT_KEY,
818                actual: 513,
819                max: 512,
820            })
821        ));
822    }
823
824    #[test]
825    fn higher_version_traceparents_preserve_stable_fields() {
826        let traceparent =
827            TraceParent::parse("01-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01-extra")
828                .expect("higher versions can carry additive fields");
829
830        assert_eq!(traceparent.trace_id(), "0af7651916cd43dd8448eb211c80319c");
831        assert_eq!(traceparent.span_id(), "00f067aa0ba902b7");
832        assert!(traceparent.sampled());
833    }
834
835    #[test]
836    fn strict_context_rejects_oversized_optional_metadata() {
837        let mut meta = MetaObject::new();
838        meta.set_traceparent("x".repeat(4096));
839        assert!(TraceContext::from_meta(&meta, TraceTrust::Untrusted).is_err());
840
841        let mut meta = MetaObject::new();
842        meta.set_traceparent(VALID_TRACEPARENT);
843        meta.set_tracestate("v".repeat(20));
844        let limits = TraceLimits {
845            max_tracestate_len: 8,
846            ..TraceLimits::default()
847        };
848        assert!(TraceContext::from_meta_with_limits(&meta, TraceTrust::Untrusted, limits).is_err());
849
850        let mut meta = MetaObject::new();
851        meta.set_traceparent(VALID_TRACEPARENT);
852        meta.set_baggage("a".repeat(20));
853        let limits = TraceLimits {
854            max_baggage_len: 8,
855            ..TraceLimits::default()
856        };
857        assert!(TraceContext::from_meta_with_limits(&meta, TraceTrust::Untrusted, limits).is_err());
858    }
859
860    #[test]
861    fn strict_context_rejects_invalid_optional_metadata() {
862        let mut meta = MetaObject::new();
863        meta.set_traceparent(VALID_TRACEPARENT);
864        meta.set_tracestate("vendor=value,vendor=other");
865        let error = TraceContext::from_meta(&meta, TraceTrust::Untrusted)
866            .expect_err("duplicate tracestate keys should be rejected");
867        assert!(matches!(error, TraceParseError::DuplicateTraceStateKey));
868
869        let mut meta = MetaObject::new();
870        meta.set_traceparent(VALID_TRACEPARENT);
871        meta.set_baggage("a=1,b=2,c=3");
872        let limits = TraceLimits {
873            max_baggage_members: 2,
874            ..TraceLimits::default()
875        };
876        let error = TraceContext::from_meta_with_limits(&meta, TraceTrust::Untrusted, limits)
877            .expect_err("baggage member cap should be enforced");
878
879        assert!(matches!(
880            error,
881            TraceParseError::TooManyBaggageMembers { actual: 3, max: 2 }
882        ));
883        assert!(!error.safe_reason().contains("a=1"));
884    }
885}