Skip to main content

netget/server/statsd/
codec.rs

1//! Bounded StatsD / DogStatsD UDP codec. Unknown extensions fail explicitly.
2use 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/// Values stay textual: signed gauges are deltas and sets can contain non-numeric members.
32#[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        // A service-check message consumes the remainder, including any pipe characters.
334        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    // One validator for both directions, including dialect and numeric restrictions.
377    encode_record(&record, dialect)?;
378    Ok(record)
379}
380
381/// Entire datagram is rejected if any line is malformed; no partial batch is emitted.
382pub 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}