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 any tags exceed the maximum length, they're dropped. If the number
105/// of tags exceeds the maximum count, the excess tags are dropped. The remaining slice doesn't contain any dropped tags.
106///
107/// # Errors
108///
109/// If the input slice isn't at least one byte long, or if it's not valid UTF-8, an error is returned.
110#[inline]
111pub fn tags(config: &DogStatsDCodecConfiguration) -> impl Fn(&[u8]) -> IResult<&[u8], RawTags<'_>> {
112    let max_tag_count = config.maximum_tag_count;
113    let max_tag_len = config.maximum_tag_length;
114
115    move |input| match simdutf8::basic::from_utf8(input) {
116        Ok(tags) => Ok((&[], RawTags::new(tags, max_tag_count, max_tag_len))),
117        Err(_) => Err(nom::Err::Error(Error::new(input, ErrorKind::Verify))),
118    }
119}
120
121/// Parses a Unix timestamp from the input slice.
122///
123/// # Errors
124///
125/// If the input slice isn't a valid unsigned 64-bit integer, an error is returned.
126#[inline]
127pub fn unix_timestamp(input: &[u8]) -> IResult<&[u8], u64> {
128    parse_u64(input)
129}
130
131/// Parses Local Data from the input slice.
132///
133/// # Errors
134///
135/// If the input slice doesn't contain at least one byte of valid characters, an error is returned.
136#[inline]
137pub fn local_data(input: &[u8]) -> IResult<&[u8], &str> {
138    // Local Data is only meant to be able to represent container IDs (which arelong hexadecimal strings), or in special
139    // cases, the inode number of the cgroup controller that contains the container sending the metrics, where the value
140    // will look like `in-<integer value>`.
141    //
142    // In some cases, it might contain _multiple_ of these values, separated by a comma.
143    let valid_char = |c: u8| c.is_ascii_alphanumeric() || c == b'-' || c == b',';
144    map(take_while1(valid_char), |b: &[u8]| {
145        // SAFETY: We know the bytes in `b` can only be comprised of ASCII characters, which ensures that it's valid to
146        // interpret the bytes directly as UTF-8.
147        saluki_antithesis::always!(
148            b.is_ascii(),
149            "DogStatsD local-data bytes are ASCII before unchecked UTF-8 conversion"
150        );
151        unsafe { std::str::from_utf8_unchecked(b) }
152    })
153    .parse(input)
154}
155
156/// Parses External Data from the input slice.
157///
158/// # Errors
159///
160/// If the input slice doesn't contain at least one byte of valid characters, an error is returned.
161#[inline]
162pub fn external_data(input: &[u8]) -> IResult<&[u8], &str> {
163    // External Data is only meant to be able to represent origin information, which includes container names, pod UIDs,
164    // and the like... which are constrained by the RFC 1123 definition of a DNS label: lowercase ASCII letters,
165    // numbers, and hyphens.
166    //
167    // We don't go the full nine yards with enforcing the "starts with a letter and number" bit.. but we _do_ allow
168    // commas since individual items in the External Data string are comma-separated.
169    let valid_char = |c: u8| c.is_ascii_lowercase() || c.is_ascii_digit() || c == b'-' || c == b',';
170    map(take_while1(valid_char), |b: &[u8]| {
171        // SAFETY: We know the bytes in `b` can only be comprised of ASCII characters, which ensures that it's valid to
172        // interpret the bytes directly as UTF-8.
173        saluki_antithesis::always!(
174            b.is_ascii(),
175            "DogStatsD external-data bytes are ASCII before unchecked UTF-8 conversion"
176        );
177        unsafe { std::str::from_utf8_unchecked(b) }
178    })
179    .parse(input)
180}
181
182/// Parses `OriginTagCardinality` from the input slice.
183///
184/// Unknown cardinality values are accepted and returned as `None` rather than failing the parse.
185/// This matches the behavior of the core Datadog Agent, which silently ignores unrecognized values.
186#[inline]
187pub fn cardinality(input: &[u8]) -> IResult<&[u8], Option<OriginTagCardinality>> {
188    let (remaining, raw_bytes) = all_consuming(preceded(tag(CARDINALITY_PREFIX), rest)).parse(input)?;
189
190    // Use simdutf8 (consistent with other UTF-8 checks in this codec) for checked conversion.
191    // Non-UTF-8 bytes are treated as an unrecognized value — return None so the frame continues
192    // processing rather than hard-failing.
193    let cardinality = simdutf8::basic::from_utf8(raw_bytes)
194        .ok()
195        .and_then(|s| OriginTagCardinality::try_from(s).ok());
196
197    Ok((remaining, cardinality))
198}
199
200#[cfg(test)]
201mod tests {
202    use saluki_context::origin::OriginTagCardinality;
203
204    use super::{cardinality, CARDINALITY_PREFIX};
205
206    fn card(s: &str) -> Vec<u8> {
207        format!("{}{}", simdutf8::basic::from_utf8(CARDINALITY_PREFIX).unwrap(), s).into_bytes()
208    }
209
210    #[test]
211    fn cardinality_known_values() {
212        let cases = [
213            ("none", Some(OriginTagCardinality::None)),
214            ("low", Some(OriginTagCardinality::Low)),
215            ("orchestrator", Some(OriginTagCardinality::Orchestrator)),
216            ("high", Some(OriginTagCardinality::High)),
217        ];
218        for (value, expected) in cases {
219            let (_, result) = cardinality(&card(value)).expect("parse should succeed");
220            assert_eq!(result, expected, "failed for '{}'", value);
221        }
222    }
223
224    #[test]
225    fn cardinality_unknown_value_returns_none() {
226        // An unrecognized value should parse successfully and return None rather than
227        // failing the parse and dropping the whole metric frame.
228        let (_, result) = cardinality(&card("not-a-valid-cardinality")).expect("parse should succeed");
229        assert_eq!(result, None);
230    }
231
232    #[test]
233    fn cardinality_case_insensitive() {
234        // Matching is case-insensitive to align with the core Datadog Agent (StringToTagCardinality
235        // uses strings.ToLower). Wrong-case values should resolve to the correct cardinality.
236        let cases = [
237            ("LOW", Some(OriginTagCardinality::Low)),
238            ("HIGH", Some(OriginTagCardinality::High)),
239            ("Orchestrator", Some(OriginTagCardinality::Orchestrator)),
240            ("NONE", Some(OriginTagCardinality::None)),
241        ];
242        for (value, expected) in cases {
243            let (_, result) = cardinality(&card(value)).expect("parse should succeed");
244            assert_eq!(result, expected, "failed for '{}'", value);
245        }
246    }
247
248    #[test]
249    fn cardinality_non_utf8_bytes_returns_none() {
250        // Non-UTF-8 bytes after the prefix must not invoke undefined behavior; they should
251        // be treated as an unrecognized value and return None. This is the bug that was fixed:
252        // the previous implementation used from_utf8_unchecked which would cause UB here.
253        let mut input = CARDINALITY_PREFIX.to_vec();
254        input.extend_from_slice(&[0xff, 0xfe]);
255        let (_, result) = cardinality(&input).expect("parse should succeed");
256        assert_eq!(result, None);
257    }
258}