Skip to main content

helix_driver_host/trace/
carrier.rs

1use std::collections::VecDeque;
2use std::sync::Arc;
3
4use parking_lot::Mutex;
5use serde::{Deserialize, Serialize};
6use thiserror::Error;
7
8pub const TRACE_CARRIER_JSON_MAX_BYTES: usize = 16 * 1024;
9pub const TRACEPARENT_BYTES: usize = 55;
10pub const TRACE_BAGGAGE_MAX_BYTES: usize = 8 * 1024;
11
12#[derive(Clone, Debug, Default, PartialEq, Eq)]
13pub struct TraceCarrier {
14    pub traceparent: Option<String>,
15    pub baggage: Option<String>,
16    pub raw_json: Option<String>,
17}
18
19impl TraceCarrier {
20    pub fn from_json_str(raw: &str) -> Result<Self, TraceCarrierError> {
21        if raw.len() > TRACE_CARRIER_JSON_MAX_BYTES {
22            return Err(TraceCarrierError::TooLarge {
23                field: "trace_json",
24                max: TRACE_CARRIER_JSON_MAX_BYTES,
25                actual: raw.len(),
26            });
27        }
28        let wire: TraceCarrierWire =
29            serde_json::from_str(raw).map_err(|e| TraceCarrierError::Parse(e.to_string()))?;
30        let traceparent = match wire.traceparent.filter(|v| !v.is_empty()) {
31            Some(value) if value.len() > TRACEPARENT_BYTES => {
32                return Err(TraceCarrierError::TooLarge {
33                    field: "traceparent",
34                    max: TRACEPARENT_BYTES,
35                    actual: value.len(),
36                });
37            }
38            Some(value) if parse_trace_id(&value).is_none() => {
39                return Err(TraceCarrierError::InvalidTraceparent);
40            }
41            value => value,
42        };
43        let baggage = match wire.baggage.filter(|v| !v.is_empty()) {
44            Some(value) if value.len() > TRACE_BAGGAGE_MAX_BYTES => {
45                return Err(TraceCarrierError::TooLarge {
46                    field: "baggage",
47                    max: TRACE_BAGGAGE_MAX_BYTES,
48                    actual: value.len(),
49                });
50            }
51            value => value,
52        };
53        Ok(Self {
54            traceparent,
55            baggage,
56            raw_json: Some(raw.to_string()),
57        })
58    }
59
60    pub fn from_headers(headers: &[(String, String)]) -> Option<Self> {
61        let traceparent = headers
62            .iter()
63            .find(|(name, _)| name.eq_ignore_ascii_case("traceparent"))
64            .and_then(|(_, value)| bounded_traceparent(value));
65        let baggage = headers
66            .iter()
67            .find(|(name, _)| name.eq_ignore_ascii_case("baggage"))
68            .and_then(|(_, value)| bounded_baggage(value));
69
70        if traceparent.is_none() && baggage.is_none() {
71            return None;
72        }
73
74        Some(Self {
75            traceparent,
76            baggage,
77            raw_json: None,
78        })
79    }
80
81    /// 从 Go HTTP 响应的 `X-CSES-Traceparent` 读取服务端 Span 上下文。
82    pub fn from_response_headers(headers: &[(String, String)]) -> Option<Self> {
83        let traceparent = headers
84            .iter()
85            .find(|(name, _)| name.eq_ignore_ascii_case("x-cses-traceparent"))
86            .and_then(|(_, value)| bounded_traceparent(value));
87        let baggage = headers
88            .iter()
89            .find(|(name, _)| name.eq_ignore_ascii_case("baggage"))
90            .and_then(|(_, value)| bounded_baggage(value));
91
92        if traceparent.is_none() && baggage.is_none() {
93            return None;
94        }
95
96        Some(Self {
97            traceparent,
98            baggage,
99            raw_json: None,
100        })
101    }
102
103    /// 从 Go WS envelope 的顶层 `tracing` carrier 提取 W3C 上下文。
104    ///
105    /// 只解析有界 JSON 和标准 traceparent/baggage;业务 data、track id 与消息正文不进入
106    /// OTel carrier,避免把业务关联字段误当成 trace context。
107    pub fn from_ws_frame(frame: &[u8]) -> Option<Self> {
108        if frame.len() > TRACE_CARRIER_JSON_MAX_BYTES {
109            return None;
110        }
111        let envelope = serde_json::from_slice::<WsEnvelopeWire>(frame).ok()?;
112        let tracing = envelope.tracing?;
113        if tracing.traceparent.is_none() && tracing.baggage.is_none() {
114            return None;
115        }
116        let raw = serde_json::to_string(&tracing).ok()?;
117        Self::from_json_str(&raw).ok()
118    }
119}
120
121#[derive(Debug, Error)]
122pub enum TraceCarrierError {
123    #[error("trace sidecar JSON parse failed: {0}")]
124    Parse(String),
125    #[error("{field} exceeds {max} bytes: {actual}")]
126    TooLarge {
127        field: &'static str,
128        max: usize,
129        actual: usize,
130    },
131    #[error("traceparent is not valid W3C version 00 format")]
132    InvalidTraceparent,
133}
134
135#[derive(Debug, Deserialize, Serialize)]
136struct TraceCarrierWire {
137    #[serde(default)]
138    traceparent: Option<String>,
139    #[serde(default)]
140    baggage: Option<String>,
141}
142
143#[derive(Debug, Deserialize)]
144struct WsEnvelopeWire {
145    #[serde(default)]
146    tracing: Option<TraceCarrierWire>,
147}
148
149pub(crate) fn parse_trace_id(traceparent: &str) -> Option<String> {
150    if !is_valid_w3c_traceparent(traceparent) {
151        return None;
152    }
153    Some(traceparent[3..35].to_string())
154}
155
156fn bounded_traceparent(value: &str) -> Option<String> {
157    if value.len() > TRACEPARENT_BYTES || parse_trace_id(value).is_none() {
158        return None;
159    }
160    Some(value.to_string())
161}
162
163fn bounded_baggage(value: &str) -> Option<String> {
164    if value.is_empty() || value.len() > TRACE_BAGGAGE_MAX_BYTES {
165        return None;
166    }
167    Some(value.to_string())
168}
169
170fn is_valid_w3c_traceparent(value: &str) -> bool {
171    if value.len() != TRACEPARENT_BYTES {
172        return false;
173    }
174    let bytes = value.as_bytes();
175    if &bytes[0..2] != b"00" || bytes[2] != b'-' || bytes[35] != b'-' || bytes[52] != b'-' {
176        return false;
177    }
178
179    let trace_id = &value[3..35];
180    let parent_id = &value[36..52];
181    let flags = &value[53..55];
182    all_lower_hex(trace_id)
183        && all_lower_hex(parent_id)
184        && all_lower_hex(flags)
185        && !all_zero(trace_id)
186        && !all_zero(parent_id)
187}
188
189fn all_lower_hex(value: &str) -> bool {
190    value
191        .bytes()
192        .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
193}
194
195fn all_zero(value: &str) -> bool {
196    value.bytes().all(|b| b == b'0')
197}
198
199#[derive(Clone, Debug, Default)]
200pub struct CommandTraceQueue {
201    inner: Arc<Mutex<VecDeque<Option<TraceCarrier>>>>,
202}
203
204impl CommandTraceQueue {
205    pub fn push_slot(&self, carrier: Option<TraceCarrier>) {
206        self.inner.lock().push_back(carrier);
207    }
208
209    pub fn rollback_last(&self) {
210        let _ = self.inner.lock().pop_back();
211    }
212
213    pub fn pop_next(&self) -> Option<TraceCarrier> {
214        self.inner.lock().pop_front().flatten()
215    }
216
217    #[cfg(test)]
218    pub(super) fn len_for_test(&self) -> usize {
219        self.inner.lock().len()
220    }
221}