saluki_io/deser/codec/dogstatsd/
helpers.rs

1use nom::{
2    bytes::complete::{tag, take_while1},
3    character::complete::u64 as parse_u64,
4    combinator::{all_consuming, map, rest},
5    error::{Error, ErrorKind},
6    sequence::preceded,
7    IResult, Parser as _,
8};
9use saluki_context::{origin::OriginTagCardinality, tags::RawTags};
10
11use super::DogStatsDCodecConfiguration;
12
13/// DogStatsD message type.
14#[derive(Eq, PartialEq)]
15pub enum MessageType {
16    MetricSample,
17    Event,
18    ServiceCheck,
19}
20
21pub const EVENT_PREFIX: &[u8] = b"_e{";
22pub const SERVICE_CHECK_PREFIX: &[u8] = b"_sc|";
23
24pub const TIMESTAMP_PREFIX: &[u8] = b"d:";
25pub const HOSTNAME_PREFIX: &[u8] = b"h:";
26pub const AGGREGATION_KEY_PREFIX: &[u8] = b"k:";
27pub const PRIORITY_PREFIX: &[u8] = b"p:";
28pub const SOURCE_TYPE_PREFIX: &[u8] = b"s:";
29pub const ALERT_TYPE_PREFIX: &[u8] = b"t:";
30pub const TAGS_PREFIX: &[u8] = b"#";
31pub const SERVICE_CHECK_MESSAGE_PREFIX: &[u8] = b"m:";
32pub const LOCAL_DATA_PREFIX: &[u8] = b"c:";
33pub const EXTERNAL_DATA_PREFIX: &[u8] = b"e:";
34pub const CARDINALITY_PREFIX: &[u8] = b"card:";
35
36/// Parses the given raw payload and returns the DogStatsD message type.
37///
38/// If the payload isn't an event or service check, it's assumed to be a metric.
39#[inline]
40pub fn parse_message_type(data: &[u8]) -> MessageType {
41    if data.starts_with(EVENT_PREFIX) {
42        return MessageType::Event;
43    } else if data.starts_with(SERVICE_CHECK_PREFIX) {
44        return MessageType::ServiceCheck;
45    }
46    MessageType::MetricSample
47}
48
49/// Splits the input buffer at the given delimiter.
50///
51/// If the delimiter isn't found, or the input buffer is empty, `None` is returned. Otherwise, the buffer is
52/// split into two parts at the delimiter, and the delimiter is _not_ included.
53#[inline]
54pub fn split_at_delimiter(input: &[u8], delimiter: u8) -> Option<(&[u8], &[u8])> {
55    match memchr::memchr(delimiter, input) {
56        Some(index) => Some((&input[0..index], &input[index + 1..input.len()])),
57        None => {
58            if input.is_empty() {
59                None
60            } else {
61                Some((input, &[]))
62            }
63        }
64    }
65}
66
67/// Maps the input slice as a UTF-8 string.
68///
69/// # Errors
70///
71/// If the input slice isn't valid UTF-8, an error is returned.
72#[inline]
73pub fn utf8(input: &[u8]) -> IResult<&[u8], &str> {
74    match simdutf8::basic::from_utf8(input) {
75        Ok(s) => Ok((&[], s)),
76        Err(_) => Err(nom::Err::Error(Error::new(input, ErrorKind::Verify))),
77    }
78}
79
80/// Returns the longest input slice that contains only ASCII alphanumeric characters and "separators" as a UTF-8 string.
81///
82/// Separators are defined as spaces, underscores, hyphens, and periods.
83///
84/// # Errors
85///
86/// If the input slice doesn't at least one byte of valid characters, an error is returned.
87#[inline]
88pub fn ascii_alphanum_and_seps(input: &[u8]) -> IResult<&[u8], &str> {
89    let valid_char = |c: u8| c.is_ascii_alphanumeric() || c == b' ' || c == b'_' || c == b'-' || c == b'.';
90    map(take_while1(valid_char), |b: &[u8]| {
91        // SAFETY: We know the bytes in `b` can only be comprised of ASCII characters, which ensures that it's valid to
92        // interpret the bytes directly as UTF-8.
93        saluki_antithesis::always!(
94            b.is_ascii(),
95            "DogStatsD name bytes are ASCII before unchecked UTF-8 conversion"
96        );
97        unsafe { std::str::from_utf8_unchecked(b) }
98    })
99    .parse(input)
100}
101
102/// Extracts as many raw tags from the input slice as possible, up to the configured limit.
103///
104/// Tags can be limited by length as well as count. If a tag exceeds the maximum length, it is truncated. If the number
105/// of tags exceeds the maximum count, the excess tags are dropped. Invalid UTF-8 is replaced with one U+FFFD character
106/// per contiguous invalid byte run so downstream tag handling operates exclusively on valid UTF-8.
107#[inline]
108pub fn tags(config: &DogStatsDCodecConfiguration) -> impl Fn(&[u8]) -> IResult<&[u8], RawTags<'_>> {
109    let max_tag_count = config.maximum_tag_count;
110    let max_tag_len = config.maximum_tag_length;
111
112    move |input| match simdutf8::basic::from_utf8(input) {
113        Ok(tags) => Ok((&[], RawTags::new(tags, max_tag_count, max_tag_len))),
114        Err(_) => Ok((
115            &[],
116            RawTags::from_owned(to_valid_utf8(input), max_tag_count, max_tag_len),
117        )),
118    }
119}
120
121fn to_valid_utf8(bytes: &[u8]) -> String {
122    let mut output = String::with_capacity(bytes.len());
123    let mut previous_chunk_was_invalid = false;
124
125    for chunk in bytes.utf8_chunks() {
126        if !chunk.valid().is_empty() {
127            output.push_str(chunk.valid());
128            previous_chunk_was_invalid = false;
129        }
130        if !chunk.invalid().is_empty() {
131            if !previous_chunk_was_invalid {
132                output.push('\u{FFFD}');
133            }
134            previous_chunk_was_invalid = true;
135        }
136    }
137
138    output
139}
140
141/// Parses a Unix timestamp from the input slice.
142///
143/// # Errors
144///
145/// If the input slice isn't a valid unsigned 64-bit integer, an error is returned.
146#[inline]
147pub fn unix_timestamp(input: &[u8]) -> IResult<&[u8], u64> {
148    parse_u64(input)
149}
150
151/// Parses Local Data from the input slice.
152///
153/// # Errors
154///
155/// If the input slice doesn't contain at least one byte of valid characters, an error is returned.
156#[inline]
157pub fn local_data(input: &[u8]) -> IResult<&[u8], &str> {
158    // Local Data is only meant to be able to represent container IDs (which arelong hexadecimal strings), or in special
159    // cases, the inode number of the cgroup controller that contains the container sending the metrics, where the value
160    // will look like `in-<integer value>`.
161    //
162    // In some cases, it might contain _multiple_ of these values, separated by a comma.
163    let valid_char = |c: u8| c.is_ascii_alphanumeric() || c == b'-' || c == b',';
164    map(take_while1(valid_char), |b: &[u8]| {
165        // SAFETY: We know the bytes in `b` can only be comprised of ASCII characters, which ensures that it's valid to
166        // interpret the bytes directly as UTF-8.
167        saluki_antithesis::always!(
168            b.is_ascii(),
169            "DogStatsD local-data bytes are ASCII before unchecked UTF-8 conversion"
170        );
171        unsafe { std::str::from_utf8_unchecked(b) }
172    })
173    .parse(input)
174}
175
176/// Parses External Data from the input slice.
177///
178/// # Errors
179///
180/// If the input slice doesn't contain at least one byte of valid characters, an error is returned.
181#[inline]
182pub fn external_data(input: &[u8]) -> IResult<&[u8], &str> {
183    // External Data is only meant to be able to represent origin information, which includes container names, pod UIDs,
184    // and the like... which are constrained by the RFC 1123 definition of a DNS label: lowercase ASCII letters,
185    // numbers, and hyphens.
186    //
187    // We don't go the full nine yards with enforcing the "starts with a letter and number" bit.. but we _do_ allow
188    // commas since individual items in the External Data string are comma-separated.
189    let valid_char = |c: u8| c.is_ascii_lowercase() || c.is_ascii_digit() || c == b'-' || c == b',';
190    map(take_while1(valid_char), |b: &[u8]| {
191        // SAFETY: We know the bytes in `b` can only be comprised of ASCII characters, which ensures that it's valid to
192        // interpret the bytes directly as UTF-8.
193        saluki_antithesis::always!(
194            b.is_ascii(),
195            "DogStatsD external-data bytes are ASCII before unchecked UTF-8 conversion"
196        );
197        unsafe { std::str::from_utf8_unchecked(b) }
198    })
199    .parse(input)
200}
201
202/// Parses `OriginTagCardinality` from the input slice.
203///
204/// Unknown cardinality values are accepted and returned as `None` rather than failing the parse.
205/// This matches the behavior of the core Datadog Agent, which silently ignores unrecognized values.
206#[inline]
207pub fn cardinality(input: &[u8]) -> IResult<&[u8], Option<OriginTagCardinality>> {
208    let (remaining, raw_bytes) = all_consuming(preceded(tag(CARDINALITY_PREFIX), rest)).parse(input)?;
209
210    // Use simdutf8 (consistent with other UTF-8 checks in this codec) for checked conversion.
211    // Non-UTF-8 bytes are treated as an unrecognized value — return None so the frame continues
212    // processing rather than hard-failing.
213    let cardinality = simdutf8::basic::from_utf8(raw_bytes)
214        .ok()
215        .and_then(|s| OriginTagCardinality::try_from(s).ok());
216
217    Ok((remaining, cardinality))
218}
219
220#[cfg(test)]
221mod tests {
222    use saluki_context::origin::OriginTagCardinality;
223
224    use super::{cardinality, to_valid_utf8, CARDINALITY_PREFIX};
225
226    fn card(s: &str) -> Vec<u8> {
227        format!("{}{}", simdutf8::basic::from_utf8(CARDINALITY_PREFIX).unwrap(), s).into_bytes()
228    }
229
230    #[test]
231    fn to_valid_utf8_replaces_each_contiguous_invalid_run_once() {
232        let cases = [
233            (b"ok".as_slice(), "ok"),
234            (b"a\xff\xfeb".as_slice(), "a\u{FFFD}b"),
235            (b"\xff\xff\xff".as_slice(), "\u{FFFD}"),
236            (b"a\xffb\xffc".as_slice(), "a\u{FFFD}b\u{FFFD}c"),
237            ("café".as_bytes(), "café"),
238        ];
239
240        for (input, expected) in cases {
241            assert_eq!(to_valid_utf8(input), expected, "failed for {input:?}");
242        }
243    }
244
245    #[test]
246    fn cardinality_known_values() {
247        let cases = [
248            ("none", Some(OriginTagCardinality::None)),
249            ("low", Some(OriginTagCardinality::Low)),
250            ("orchestrator", Some(OriginTagCardinality::Orchestrator)),
251            ("high", Some(OriginTagCardinality::High)),
252        ];
253        for (value, expected) in cases {
254            let (_, result) = cardinality(&card(value)).expect("parse should succeed");
255            assert_eq!(result, expected, "failed for '{}'", value);
256        }
257    }
258
259    #[test]
260    fn cardinality_unknown_value_returns_none() {
261        // An unrecognized value should parse successfully and return None rather than
262        // failing the parse and dropping the whole metric frame.
263        let (_, result) = cardinality(&card("not-a-valid-cardinality")).expect("parse should succeed");
264        assert_eq!(result, None);
265    }
266
267    #[test]
268    fn cardinality_case_insensitive() {
269        // Matching is case-insensitive to align with the core Datadog Agent (StringToTagCardinality
270        // uses strings.ToLower). Wrong-case values should resolve to the correct cardinality.
271        let cases = [
272            ("LOW", Some(OriginTagCardinality::Low)),
273            ("HIGH", Some(OriginTagCardinality::High)),
274            ("Orchestrator", Some(OriginTagCardinality::Orchestrator)),
275            ("NONE", Some(OriginTagCardinality::None)),
276        ];
277        for (value, expected) in cases {
278            let (_, result) = cardinality(&card(value)).expect("parse should succeed");
279            assert_eq!(result, expected, "failed for '{}'", value);
280        }
281    }
282
283    #[test]
284    fn cardinality_non_utf8_bytes_returns_none() {
285        // Non-UTF-8 bytes after the prefix must not invoke undefined behavior; they should
286        // be treated as an unrecognized value and return None. This is the bug that was fixed:
287        // the previous implementation used from_utf8_unchecked which would cause UB here.
288        let mut input = CARDINALITY_PREFIX.to_vec();
289        input.extend_from_slice(&[0xff, 0xfe]);
290        let (_, result) = cardinality(&input).expect("parse should succeed");
291        assert_eq!(result, None);
292    }
293}
saluki_io/deser/codec/dogstatsd/helpers.rs

saluki_io/deser/codec/dogstatsd/
helpers.rs
