helix_driver_host/trace/
carrier.rs1use 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 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 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}