1use anyhow::{bail, ensure, Context, Result};
3use serde::{Deserialize, Serialize};
4
5pub const MAX_DATAGRAM_BYTES: usize = 8192;
6pub const MAX_RECORDS: usize = 256;
7pub const DEFAULT_DIALECT: &str = "dogstatsd";
8pub const DEFAULT_LLM_FALLBACK: bool = false;
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum Dialect {
12 Statsd,
13 Dogstatsd,
14}
15impl Dialect {
16 pub fn parse(value: &str) -> Result<Self> {
17 match value {
18 "statsd" => Ok(Self::Statsd),
19 "dogstatsd" => Ok(Self::Dogstatsd),
20 _ => bail!("dialect must be statsd or dogstatsd"),
21 }
22 }
23 pub fn as_str(self) -> &'static str {
24 match self {
25 Self::Statsd => "statsd",
26 Self::Dogstatsd => "dogstatsd",
27 }
28 }
29}
30
31#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
33#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
34pub enum Record {
35 Metric {
36 name: String,
37 value: String,
38 metric_type: String,
39 #[serde(default, skip_serializing_if = "Option::is_none")]
40 sample_rate: Option<f64>,
41 #[serde(default, skip_serializing_if = "Vec::is_empty")]
42 tags: Vec<String>,
43 },
44 Event {
45 title: String,
46 text: String,
47 #[serde(default, skip_serializing_if = "Option::is_none")]
48 timestamp: Option<u64>,
49 #[serde(default, skip_serializing_if = "Option::is_none")]
50 hostname: Option<String>,
51 #[serde(default, skip_serializing_if = "Option::is_none")]
52 aggregation_key: Option<String>,
53 #[serde(default, skip_serializing_if = "Option::is_none")]
54 priority: Option<String>,
55 #[serde(default, skip_serializing_if = "Option::is_none")]
56 source_type: Option<String>,
57 #[serde(default, skip_serializing_if = "Option::is_none")]
58 alert_type: Option<String>,
59 #[serde(default, skip_serializing_if = "Vec::is_empty")]
60 tags: Vec<String>,
61 },
62 ServiceCheck {
63 name: String,
64 status: u8,
65 #[serde(default, skip_serializing_if = "Option::is_none")]
66 timestamp: Option<u64>,
67 #[serde(default, skip_serializing_if = "Option::is_none")]
68 hostname: Option<String>,
69 #[serde(default, skip_serializing_if = "Vec::is_empty")]
70 tags: Vec<String>,
71 #[serde(default, skip_serializing_if = "Option::is_none")]
72 message: Option<String>,
73 },
74}
75
76fn field(value: &str) -> Result<()> {
77 ensure!(
78 !value.is_empty() && !value.contains(['|', '\n', '\r', '\0']),
79 "empty or invalid field delimiter"
80 );
81 Ok(())
82}
83fn tags_suffix(tags: &[String], output: &mut String) -> Result<()> {
84 if !tags.is_empty() {
85 for tag in tags {
86 field(tag)?;
87 ensure!(!tag.contains(','), "comma in tag");
88 }
89 output.push_str("|#");
90 output.push_str(&tags.join(","));
91 }
92 Ok(())
93}
94fn opt_suffix(key: &str, value: &Option<String>, output: &mut String) -> Result<()> {
95 if let Some(value) = value {
96 field(value)?;
97 output.push('|');
98 output.push_str(key);
99 output.push(':');
100 output.push_str(value);
101 }
102 Ok(())
103}
104fn escaped(value: &str) -> Result<String> {
105 ensure!(
106 !value.contains(['\r', '\0']),
107 "carriage return or NUL in text"
108 );
109 Ok(value.replace('\n', "\\n"))
110}
111
112pub fn encode_record(record: &Record, dialect: Dialect) -> Result<String> {
113 let out = match record {
114 Record::Metric {
115 name,
116 value,
117 metric_type,
118 sample_rate,
119 tags,
120 } => {
121 ensure!(
122 !name.is_empty()
123 && name
124 .bytes()
125 .all(|c| c.is_ascii_alphanumeric() || c == b'_' || c == b'.'),
126 "metric name must contain ASCII letters, digits, underscores or periods"
127 );
128 field(value)?;
129 ensure!(
130 !value.contains(':'),
131 "packed metric values are not supported"
132 );
133 ensure!(
134 matches!(metric_type.as_str(), "c" | "g" | "ms" | "s")
135 || dialect == Dialect::Dogstatsd && matches!(metric_type.as_str(), "h" | "d"),
136 "unsupported metric type for dialect"
137 );
138 if metric_type != "s" {
139 let number: f64 = value.parse().context("metric value must be numeric")?;
140 ensure!(
141 number.is_finite() && !value.chars().any(char::is_whitespace),
142 "metric value must be finite without whitespace"
143 );
144 }
145 ensure!(
146 dialect == Dialect::Dogstatsd || tags.is_empty(),
147 "tags require dogstatsd"
148 );
149 let mut out = format!("{name}:{value}|{metric_type}");
150 if let Some(rate) = sample_rate {
151 ensure!(
152 rate.is_finite() && (0.0..=1.0).contains(rate),
153 "sample rate must be between zero and one"
154 );
155 ensure!(
156 dialect == Dialect::Dogstatsd || *rate > 0.0,
157 "classic StatsD sample rate must be positive"
158 );
159 ensure!(
160 matches!(metric_type.as_str(), "c" | "ms" | "h" | "d"),
161 "sampling is only supported on counts, timers, histograms and distributions"
162 );
163 out.push_str(&format!("|@{rate}"));
164 }
165 tags_suffix(tags, &mut out)?;
166 out
167 }
168 Record::Event {
169 title,
170 text,
171 timestamp,
172 hostname,
173 aggregation_key,
174 priority,
175 source_type,
176 alert_type,
177 tags,
178 } => {
179 ensure!(dialect == Dialect::Dogstatsd, "events require dogstatsd");
180 ensure!(!title.is_empty(), "event title must not be empty");
181 if let Some(priority) = priority {
182 ensure!(
183 matches!(priority.as_str(), "low" | "normal"),
184 "invalid event priority"
185 );
186 }
187 if let Some(alert) = alert_type {
188 ensure!(
189 matches!(alert.as_str(), "error" | "warning" | "info" | "success"),
190 "invalid event alert type"
191 );
192 }
193 let title = escaped(title)?;
194 let text = escaped(text)?;
195 let mut out = format!("_e{{{},{}}}:{}|{}", title.len(), text.len(), title, text);
196 if let Some(ts) = timestamp {
197 out.push_str(&format!("|d:{ts}"));
198 }
199 opt_suffix("h", hostname, &mut out)?;
200 opt_suffix("k", aggregation_key, &mut out)?;
201 opt_suffix("p", priority, &mut out)?;
202 opt_suffix("s", source_type, &mut out)?;
203 opt_suffix("t", alert_type, &mut out)?;
204 tags_suffix(tags, &mut out)?;
205 out
206 }
207 Record::ServiceCheck {
208 name,
209 status,
210 timestamp,
211 hostname,
212 tags,
213 message,
214 } => {
215 ensure!(
216 dialect == Dialect::Dogstatsd,
217 "service checks require dogstatsd"
218 );
219 field(name)?;
220 ensure!(*status <= 3, "service check status must be 0..3");
221 let mut out = format!("_sc|{name}|{status}");
222 if let Some(ts) = timestamp {
223 out.push_str(&format!("|d:{ts}"));
224 }
225 opt_suffix("h", hostname, &mut out)?;
226 tags_suffix(tags, &mut out)?;
227 if let Some(message) = message {
228 out.push_str("|m:");
229 out.push_str(&escaped(message)?.replace("m:", "m\\:"));
230 }
231 out
232 }
233 };
234 ensure!(
235 out.len() <= MAX_DATAGRAM_BYTES,
236 "record exceeds datagram size limit"
237 );
238 Ok(out)
239}
240
241pub fn encode_datagram(records: &[Record], dialect: Dialect) -> Result<Vec<u8>> {
242 ensure!(
243 !records.is_empty() && records.len() <= MAX_RECORDS,
244 "batch must contain 1..={MAX_RECORDS} records"
245 );
246 let mut output = Vec::new();
247 for record in records {
248 let line = encode_record(record, dialect)?;
249 let separator = usize::from(!output.is_empty());
250 ensure!(
251 output.len() + separator + line.len() <= MAX_DATAGRAM_BYTES,
252 "batch exceeds {MAX_DATAGRAM_BYTES} bytes"
253 );
254 if separator != 0 {
255 output.push(b'\n');
256 }
257 output.extend_from_slice(line.as_bytes());
258 }
259 Ok(output)
260}
261
262fn options<'a>(
263 suffix: &'a str,
264 allowed: &[&str],
265) -> Result<std::collections::HashMap<&'a str, &'a str>> {
266 let mut out = std::collections::HashMap::new();
267 if suffix.is_empty() {
268 return Ok(out);
269 }
270 ensure!(suffix.starts_with('|'), "missing option delimiter");
271 for part in suffix[1..].split('|') {
272 let (key, value) = if let Some(value) = part.strip_prefix('#') {
273 ("#", value)
274 } else {
275 part.split_once(':').context("malformed option")?
276 };
277 ensure!(allowed.contains(&key), "unsupported option {key}");
278 field(value)?;
279 ensure!(out.insert(key, value).is_none(), "duplicate option {key}");
280 }
281 Ok(out)
282}
283fn read_tags(value: Option<&&str>) -> Vec<String> {
284 value
285 .map(|v| v.split(',').map(str::to_owned).collect())
286 .unwrap_or_default()
287}
288fn read_timestamp(value: Option<&&str>) -> Result<Option<u64>> {
289 value
290 .map(|v| v.parse().context("invalid timestamp"))
291 .transpose()
292}
293fn read_string(value: Option<&&str>) -> Option<String> {
294 value.map(|v| (*v).to_owned())
295}
296
297pub fn parse_record(line: &str, dialect: Dialect) -> Result<Record> {
298 let record = if let Some(rest) = line.strip_prefix("_e{") {
299 let (lengths, body) = rest.split_once("}:").context("invalid event header")?;
300 let (title_len, text_len) = lengths.split_once(',').context("missing event lengths")?;
301 let title_len: usize = title_len.parse().context("invalid title length")?;
302 let text_len: usize = text_len.parse().context("invalid text length")?;
303 let title = body
304 .get(..title_len)
305 .context("invalid UTF-8 title byte length")?;
306 let rest = body
307 .get(title_len..)
308 .and_then(|s| s.strip_prefix('|'))
309 .context("missing title delimiter")?;
310 let text = rest
311 .get(..text_len)
312 .context("invalid UTF-8 text byte length")?;
313 let opts = options(&rest[text_len..], &["d", "h", "k", "p", "s", "t", "#"])?;
314 Record::Event {
315 title: title.replace("\\n", "\n"),
316 text: text.replace("\\n", "\n"),
317 timestamp: read_timestamp(opts.get("d"))?,
318 hostname: read_string(opts.get("h")),
319 aggregation_key: read_string(opts.get("k")),
320 priority: read_string(opts.get("p")),
321 source_type: read_string(opts.get("s")),
322 alert_type: read_string(opts.get("t")),
323 tags: read_tags(opts.get("#")),
324 }
325 } else if let Some(rest) = line.strip_prefix("_sc|") {
326 let (name, rest) = rest
327 .split_once('|')
328 .context("missing service check status")?;
329 let (status, suffix) = rest
330 .split_once('|')
331 .map(|(a, b)| (a, format!("|{b}")))
332 .unwrap_or((rest, String::new()));
333 let (metadata, message) = suffix
335 .split_once("|m:")
336 .map(|(a, b)| (a, Some(b.replace("\\n", "\n").replace("m\\:", "m:"))))
337 .unwrap_or((&suffix, None));
338 let opts = options(metadata, &["d", "h", "#"])?;
339 Record::ServiceCheck {
340 name: name.to_owned(),
341 status: status.parse().context("invalid service check status")?,
342 timestamp: read_timestamp(opts.get("d"))?,
343 hostname: read_string(opts.get("h")),
344 tags: read_tags(opts.get("#")),
345 message,
346 }
347 } else {
348 let (name, rest) = line.split_once(':').context("missing metric colon")?;
349 let mut parts = rest.split('|');
350 let value = parts.next().context("missing metric value")?.to_owned();
351 let metric_type = parts.next().context("missing metric type")?.to_owned();
352 let mut sample_rate = None;
353 let mut tags = None;
354 for part in parts {
355 if let Some(rate) = part.strip_prefix('@') {
356 ensure!(sample_rate.is_none(), "duplicate sample rate");
357 sample_rate = Some(rate.parse().context("invalid sample rate")?);
358 } else if let Some(value) = part.strip_prefix('#') {
359 ensure!(
360 tags.is_none() && !value.is_empty(),
361 "empty or duplicate tags"
362 );
363 tags = Some(value.split(',').map(str::to_owned).collect());
364 } else {
365 bail!("unsupported metric option");
366 }
367 }
368 Record::Metric {
369 name: name.to_owned(),
370 value,
371 metric_type,
372 sample_rate,
373 tags: tags.unwrap_or_default(),
374 }
375 };
376 encode_record(&record, dialect)?;
378 Ok(record)
379}
380
381pub fn parse_datagram(bytes: &[u8], dialect: Dialect) -> Result<Vec<Record>> {
383 ensure!(
384 !bytes.is_empty() && bytes.len() <= MAX_DATAGRAM_BYTES,
385 "datagram must contain 1..={MAX_DATAGRAM_BYTES} bytes"
386 );
387 let text = std::str::from_utf8(bytes).context("datagram is not UTF-8")?;
388 let text = text.strip_suffix('\n').unwrap_or(text);
389 let mut records = Vec::new();
390 for line in text.split('\n') {
391 ensure!(
392 records.len() < MAX_RECORDS,
393 "datagram exceeds {MAX_RECORDS} records"
394 );
395 records.push(parse_record(line, dialect)?);
396 }
397 Ok(records)
398}