libdd_trace_utils/msgpack_decoder/v05/
mod.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::msgpack_decoder::decode::error::DecodeError;
5use crate::msgpack_decoder::decode::{
6    buffer::Buffer, map::read_map_len, number::read_number, string::handle_null_marker,
7};
8use crate::span::v04::{Span, SpanBytes, SpanSlice};
9use crate::span::vec_map::VecMap;
10use crate::span::DeserializableTraceData;
11
12const PAYLOAD_LEN: u32 = 2;
13const SPAN_ELEM_COUNT: u32 = 12;
14
15/// Decodes a Bytes buffer into a `Vec<Vec<SpanBytes>>` object, also represented as a vector of
16/// `TracerPayloadV05` objects.
17///
18/// # Arguments
19///
20/// * `data` - A tinybytes Bytes buffer containing the encoded data. Bytes are expected to be
21///   encoded msgpack data containing a list of a list of v05 spans.
22///
23/// # Returns
24///
25/// * `Ok(Vec<TracerPayloadV05>)` - A vector of decoded `TracerPayloadV05` objects if successful.
26/// * `Err(DecodeError)` - An error if the decoding process fails.
27///
28/// # Errors
29///
30/// This function will return an error if:
31/// - The array length for trace count or span count cannot be read.
32/// - Any span cannot be decoded.
33///
34/// # Examples
35///
36/// ```
37/// use libdd_tinybytes;
38/// use libdd_trace_utils::msgpack_decoder::v05::from_bytes;
39/// use rmp_serde::to_vec;
40/// use std::collections::HashMap;
41///
42/// let data = (
43///     vec!["".to_string()],
44///     vec![vec![(
45///         0,
46///         0,
47///         0,
48///         1,
49///         2,
50///         3,
51///         4,
52///         5,
53///         6,
54///         HashMap::<u32, u32>::new(),
55///         HashMap::<u32, f64>::new(),
56///         0,
57///     )]],
58/// );
59/// let encoded_data = to_vec(&data).unwrap();
60/// let encoded_data_as_tinybytes = libdd_tinybytes::Bytes::from(encoded_data);
61/// let (decoded_traces, _payload_size) =
62///     from_bytes(encoded_data_as_tinybytes).expect("Decoding failed");
63///
64/// assert_eq!(1, decoded_traces.len());
65/// assert_eq!(1, decoded_traces[0].len());
66/// let decoded_span = &decoded_traces[0][0];
67/// assert_eq!("", decoded_span.name.as_str());
68/// ```
69pub fn from_bytes(
70    data: libdd_tinybytes::Bytes,
71) -> Result<(Vec<Vec<SpanBytes>>, usize), DecodeError> {
72    from_buffer(&mut Buffer::new(data))
73}
74
75/// Decodes a slice of bytes into a `Vec<Vec<SpanSlice>>` object.
76/// The resulting spans have the same lifetime as the initial buffer.
77///
78/// # Arguments
79///
80/// * `data` - A slice of bytes containing the encoded data. Bytes are expected to be encoded
81///   msgpack data containing a list of a list of v05 spans.
82///
83/// # Returns
84///
85/// * `Ok(Vec<Vec<SpanSlice>>)` - A vector of decoded `Vec<SpanSlice>` objects if successful.
86/// * `Err(DecodeError)` - An error if the decoding process fails.
87///
88/// # Errors
89///
90/// This function will return an error if:
91/// - The array length for trace count or span count cannot be read.
92/// - Any span cannot be decoded.
93///
94/// # Examples
95///
96/// ```
97/// use libdd_tinybytes;
98/// use libdd_trace_utils::msgpack_decoder::v05::from_slice;
99/// use rmp_serde::to_vec;
100/// use std::collections::HashMap;
101///
102/// let data = (
103///     vec!["".to_string()],
104///     vec![vec![(
105///         0,
106///         0,
107///         0,
108///         1,
109///         2,
110///         3,
111///         4,
112///         5,
113///         6,
114///         HashMap::<u32, u32>::new(),
115///         HashMap::<u32, f64>::new(),
116///         0,
117///     )]],
118/// );
119/// let encoded_data = to_vec(&data).unwrap();
120/// let encoded_data_as_tinybytes = libdd_tinybytes::Bytes::from(encoded_data);
121/// let (decoded_traces, _payload_size) =
122///     from_slice(&encoded_data_as_tinybytes).expect("Decoding failed");
123///
124/// assert_eq!(1, decoded_traces.len());
125/// assert_eq!(1, decoded_traces[0].len());
126/// let decoded_span = &decoded_traces[0][0];
127/// assert_eq!("", decoded_span.name);
128/// ```
129pub fn from_slice(data: &[u8]) -> Result<(Vec<Vec<SpanSlice<'_>>>, usize), DecodeError> {
130    from_buffer(&mut Buffer::new(data))
131}
132
133#[allow(clippy::type_complexity)]
134fn from_buffer<T: DeserializableTraceData>(
135    data: &mut Buffer<T>,
136) -> Result<(Vec<Vec<Span<T>>>, usize), DecodeError>
137where
138    T::Text: Clone,
139{
140    let data_elem = rmp::decode::read_array_len(data.as_mut_slice())
141        .map_err(|_| DecodeError::InvalidFormat("Unable to read payload len".to_string()))?;
142
143    if data_elem != PAYLOAD_LEN {
144        return Err(DecodeError::InvalidFormat(
145            "Invalid payload size".to_string(),
146        ));
147    }
148
149    let dict = deserialize_dict(data)?;
150
151    let trace_count = rmp::decode::read_array_len(data.as_mut_slice())
152        .map_err(|_| DecodeError::InvalidFormat("Unable to read trace len".to_string()))?;
153
154    let mut traces: Vec<Vec<Span<T>>> = Vec::with_capacity(trace_count as usize);
155    let start_len = data.len();
156
157    for _ in 0..trace_count {
158        let span_count = rmp::decode::read_array_len(data.as_mut_slice())
159            .map_err(|_| DecodeError::InvalidFormat("Unable to read span len".to_string()))?;
160        let mut trace: Vec<Span<T>> = Vec::with_capacity(span_count as usize);
161
162        for _ in 0..span_count {
163            let span = deserialize_span(data, &dict)?;
164            trace.push(span);
165        }
166        traces.push(trace);
167    }
168    Ok((traces, start_len - data.len()))
169}
170
171fn deserialize_dict<T: DeserializableTraceData>(
172    data: &mut Buffer<T>,
173) -> Result<Vec<T::Text>, DecodeError> {
174    let dict_len = rmp::decode::read_array_len(data.as_mut_slice())
175        .map_err(|_| DecodeError::InvalidFormat("Unable to read dictionary len".to_string()))?;
176
177    let mut dict: Vec<T::Text> = Vec::with_capacity(dict_len as usize);
178    for _ in 0..dict_len {
179        let str = data.read_string()?;
180        dict.push(str);
181    }
182    Ok(dict)
183}
184
185fn deserialize_span<T: DeserializableTraceData>(
186    data: &mut Buffer<T>,
187    dict: &[T::Text],
188) -> Result<Span<T>, DecodeError>
189where
190    T::Text: Clone,
191{
192    let mut span = Span::default();
193    let span_len = rmp::decode::read_array_len(data.as_mut_slice())
194        .map_err(|_| DecodeError::InvalidFormat("Unable to read dictionary len".to_string()))?;
195
196    if span_len != SPAN_ELEM_COUNT {
197        return Err(DecodeError::InvalidFormat(
198            "Invalid number of span fields".to_string(),
199        ));
200    }
201
202    span.service = get_from_dict(data, dict)?;
203    span.name = get_from_dict(data, dict)?;
204    span.resource = get_from_dict(data, dict)?;
205    span.trace_id = read_number::<_, u64>(data)? as u128;
206    span.span_id = read_number(data)?;
207    span.parent_id = read_number(data)?;
208    span.start = read_number(data)?;
209    span.duration = read_number(data)?;
210    span.error = read_number(data)?;
211    span.meta = read_indexed_map_to_bytes_strings(data, dict)?;
212    span.metrics = read_metrics(data, dict)?;
213    span.r#type = get_from_dict(data, dict)?;
214
215    // Decoded from msgpack maps (unique keys) or empty defaults: no duplicates, so mark deduped to
216    // skip the defensive dedup (and its warning) at encoding time. A later mutation re-dirties it.
217    span.meta.mark_deduped();
218    span.metrics.mark_deduped();
219    span.meta_struct.mark_deduped();
220
221    Ok(span)
222}
223
224fn get_from_dict<T: DeserializableTraceData>(
225    data: &mut Buffer<T>,
226    dict: &[T::Text],
227) -> Result<T::Text, DecodeError>
228where
229    T::Text: Clone,
230{
231    let index: u32 = read_number(data)?;
232    match dict.get(index as usize) {
233        Some(value) => Ok(value.clone()),
234        None => Err(DecodeError::InvalidFormat(
235            "Unable to locate string in the dictionary".to_string(),
236        )),
237    }
238}
239
240fn read_indexed_map_to_bytes_strings<T: DeserializableTraceData>(
241    buf: &mut Buffer<T>,
242    dict: &[T::Text],
243) -> Result<VecMap<T::Text, T::Text>, DecodeError>
244where
245    T::Text: Clone,
246{
247    let len = rmp::decode::read_map_len(buf.as_mut_slice())
248        .map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
249
250    let mut map = VecMap::with_capacity(len.try_into().unwrap_or_default());
251    for _ in 0..len {
252        let key = get_from_dict(buf, dict)?;
253        let value = get_from_dict(buf, dict)?;
254        map.insert(key, value);
255    }
256    Ok(map)
257}
258
259fn read_metrics<T: DeserializableTraceData>(
260    buf: &mut Buffer<T>,
261    dict: &[T::Text],
262) -> Result<VecMap<T::Text, f64>, DecodeError>
263where
264    T::Text: Clone,
265{
266    if handle_null_marker(buf) {
267        return Ok(VecMap::new());
268    }
269
270    let len = read_map_len(buf)?;
271
272    let mut map = VecMap::with_capacity(len);
273    for _ in 0..len {
274        let k = get_from_dict(buf, dict)?;
275        let v = read_number(buf)?;
276        map.insert(k, v);
277    }
278    Ok(map)
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284    use crate::span::SliceData;
285    use std::collections::HashMap;
286
287    type V05Span = (
288        u8,
289        u8,
290        u8,
291        u64,
292        u64,
293        u64,
294        i64,
295        i64,
296        i32,
297        HashMap<u8, u8>,
298        HashMap<u8, f64>,
299        u8,
300    );
301
302    type V05SpanMalformed = (
303        u8,
304        u8,
305        u8,
306        u64,
307        u64,
308        u64,
309        i64,
310        i64,
311        i32,
312        HashMap<u8, u8>,
313        HashMap<u8, f64>,
314    );
315
316    type V05Payload = (Vec<String>, Vec<Vec<V05Span>>);
317    type V05PayloadMalformed = (Vec<String>, Vec<Vec<V05SpanMalformed>>);
318
319    #[test]
320    fn deserialize_dict_test() {
321        let dict = vec!["foo", "bar", "baz"];
322        let mpack = rmp_serde::to_vec(&dict).unwrap();
323        let mut payload = Buffer::<SliceData>::new(mpack.as_ref());
324
325        let result = deserialize_dict(&mut payload).unwrap();
326        assert_eq!(dict, result);
327    }
328
329    #[test]
330    fn from_bytes_invalid_size_test() {
331        // 3 empty array.
332        let empty_three: [u8; 3] = [0x93, 0x90, 0x90];
333        let payload = unsafe { std::mem::transmute::<&'_ [u8], &'static [u8]>(&empty_three) };
334        let bytes = libdd_tinybytes::Bytes::from_static(payload);
335        let result = from_bytes(bytes);
336
337        assert!(result.is_err());
338        matches!(result.err().unwrap(), DecodeError::InvalidFormat(_));
339
340        // 1 empty array
341        let empty_one: [u8; 2] = [0x91, 0x90];
342        let payload = unsafe { std::mem::transmute::<&'_ [u8], &'static [u8]>(&empty_one) };
343        let bytes = libdd_tinybytes::Bytes::from_static(payload);
344        let result = from_bytes(bytes);
345
346        assert!(result.is_err());
347        matches!(result.err().unwrap(), DecodeError::InvalidFormat(_));
348    }
349
350    #[test]
351    fn from_bytes_test() {
352        let data: V05Payload = (
353            vec![
354                "".to_string(),
355                "item".to_string(),
356                "version".to_string(),
357                "7.0".to_string(),
358                "my-name".to_string(),
359                "X".to_string(),
360                "my-service".to_string(),
361                "my-resource".to_string(),
362                "_dd.sampling_rate_whatever".to_string(),
363                "value whatever".to_string(),
364                "sql".to_string(),
365            ],
366            vec![vec![(
367                6,
368                4,
369                7,
370                1,
371                2,
372                3,
373                123,
374                456,
375                1,
376                HashMap::from([(8, 9), (0, 1), (2, 3)]),
377                HashMap::from([(5, 1.2)]),
378                10,
379            )]],
380        );
381        let msgpack = rmp_serde::to_vec(&data).unwrap();
382        let (traces, _) = from_bytes(libdd_tinybytes::Bytes::from(msgpack)).unwrap();
383
384        let span = &traces[0][0];
385        assert_eq!(span.service.as_str(), "my-service");
386        assert_eq!(span.name.as_str(), "my-name");
387        assert_eq!(span.resource.as_str(), "my-resource");
388        assert_eq!(span.trace_id, 1);
389        assert_eq!(span.span_id, 2);
390        assert_eq!(span.parent_id, 3);
391        assert_eq!(span.start, 123);
392        assert_eq!(span.duration, 456);
393        assert_eq!(span.error, 1);
394        assert_eq!(span.meta.len(), 3);
395        assert_eq!(
396            span.meta
397                .get("_dd.sampling_rate_whatever")
398                .unwrap()
399                .as_str(),
400            "value whatever"
401        );
402        assert_eq!(span.meta.get("").unwrap().as_str(), "item");
403        assert_eq!(span.meta.get("version").unwrap().as_str(), "7.0");
404        assert_eq!(span.metrics.len(), 1);
405        assert_eq!(*span.metrics.get("X").unwrap(), 1.2_f64);
406        assert_eq!(span.r#type.as_str(), "sql");
407    }
408
409    #[test]
410    fn missing_dict_elements_test() {
411        let data: V05Payload = (
412            vec![
413                "".to_string(),
414                "item".to_string(),
415                "version".to_string(),
416                "7.0".to_string(),
417                "my-name".to_string(),
418                "X".to_string(),
419                "my-service".to_string(),
420                "my-resource".to_string(),
421                "_dd.sampling_rate_whatever".to_string(),
422                "value whatever".to_string(),
423            ],
424            vec![vec![(
425                6,
426                4,
427                7,
428                1,
429                2,
430                3,
431                123,
432                456,
433                1,
434                HashMap::from([(8, 9), (0, 1), (2, 3)]),
435                HashMap::from([(5, 1.2)]),
436                10,
437            )]],
438        );
439        let payload = rmp_serde::to_vec(&data).unwrap();
440        let payload = unsafe { std::mem::transmute::<&'_ [u8], &'static [u8]>(&payload) };
441        let result = from_bytes(libdd_tinybytes::Bytes::from_static(payload));
442
443        assert!(result.is_err());
444
445        // Unable to locate string in the dictionary
446        matches!(result.err().unwrap(), DecodeError::InvalidFormat(_));
447    }
448
449    #[test]
450    fn missing_span_elements_test() {
451        let data: V05PayloadMalformed = (
452            vec![
453                "".to_string(),
454                "item".to_string(),
455                "version".to_string(),
456                "7.0".to_string(),
457                "my-name".to_string(),
458                "X".to_string(),
459                "my-service".to_string(),
460                "my-resource".to_string(),
461                "_dd.sampling_rate_whatever".to_string(),
462                "value whatever".to_string(),
463            ],
464            vec![vec![(
465                6,
466                4,
467                7,
468                1,
469                2,
470                3,
471                123,
472                456,
473                1,
474                HashMap::from([(8, 9), (0, 1), (2, 3)]),
475                HashMap::from([(5, 1.2)]),
476            )]],
477        );
478
479        let payload = rmp_serde::to_vec(&data).unwrap();
480        let payload = unsafe { std::mem::transmute::<&'_ [u8], &'static [u8]>(&payload) };
481        let result = from_bytes(libdd_tinybytes::Bytes::from_static(payload));
482
483        assert!(result.is_err());
484
485        // Invalid number of span fields.
486        matches!(result.err().unwrap(), DecodeError::InvalidFormat(_));
487    }
488}
libdd_trace_utils/msgpack_decoder/v05/mod.rs

libdd_trace_utils/msgpack_decoder/v05/
mod.rs
