Skip to main content

libdd_trace_utils/msgpack_decoder/decode/
string.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::msgpack_decoder::decode::buffer::Buffer;
5use crate::msgpack_decoder::decode::error::DecodeError;
6use crate::span::vec_map::VecMap;
7use crate::span::DeserializableTraceData;
8use rmp::decode;
9
10// https://docs.rs/rmp/latest/rmp/enum.Marker.html#variant.Null (0xc0 == 192)
11const NULL_MARKER: &u8 = &0xc0;
12
13/// Read a nullable string from the slices `buf`.
14///
15/// # Errors
16/// Fails if the buffer doesn't contain a valid utf8 msgpack string or a null marker.
17#[inline]
18pub fn read_nullable_string<T: DeserializableTraceData>(
19    buf: &mut Buffer<T>,
20) -> Result<T::Text, DecodeError> {
21    if handle_null_marker(buf) {
22        Ok(T::Text::default())
23    } else {
24        buf.read_string()
25    }
26}
27
28/// Read a [VecMap] of `(String, String)` pairs from the slices `buf`.
29///
30/// # Errors
31///
32/// Fails if the buffer does not contain a valid map length prefix, or if any key or value is not a
33/// valid utf8 msgpack string.
34/// Null values are skipped.
35#[inline]
36pub fn read_str_map_to_vecmap<T: DeserializableTraceData>(
37    buf: &mut Buffer<T>,
38) -> Result<VecMap<T::Text, T::Text>, DecodeError> {
39    let len = decode::read_map_len(buf.as_mut_slice())
40        .map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
41
42    let mut map = VecMap::with_capacity(len.try_into().unwrap_or_default());
43    for _ in 0..len {
44        let key = buf.read_string()?;
45        // Only insert if value is not null
46        if !handle_null_marker(buf) {
47            let value = buf.read_string()?;
48            map.insert(key, value);
49        }
50    }
51    Ok(map)
52}
53
54/// Read a nullable vec of (string, string) pairs from the slices `buf`.
55///
56/// # Errors
57/// Fails if the buffer does not contain a valid map length prefix,
58/// or if any key or value is not a valid utf8 msgpack string.
59/// Null values are skipped (key not inserted into vec).
60#[inline]
61pub fn read_nullable_str_map_to_strings<T: DeserializableTraceData>(
62    buf: &mut Buffer<T>,
63) -> Result<VecMap<T::Text, T::Text>, DecodeError> {
64    if handle_null_marker(buf) {
65        return Ok(VecMap::new());
66    }
67
68    read_str_map_to_vecmap(buf)
69}
70
71/// Read a hashmap of (string, string) from the slices `buf`.
72/// Used for SpanLink/SpanEvent attributes which remain as HashMap.
73#[inline]
74pub fn read_str_map_to_hashmap<T: DeserializableTraceData>(
75    buf: &mut Buffer<T>,
76) -> Result<std::collections::HashMap<T::Text, T::Text>, DecodeError>
77where
78    T::Text: std::hash::Hash + Eq,
79{
80    let len = decode::read_map_len(buf.as_mut_slice())
81        .map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
82
83    let mut map = std::collections::HashMap::with_capacity(len.try_into().unwrap_or_default());
84    for _ in 0..len {
85        let key = buf.read_string()?;
86        if !handle_null_marker(buf) {
87            let value = buf.read_string()?;
88            map.insert(key, value);
89        }
90    }
91    Ok(map)
92}
93
94/// Handle the null value by peeking if the next value is a null marker, and will only advance the
95/// buffer if it is null. If it is not null, you can continue to decode as expected.
96///
97/// # Returns
98/// A boolean indicating whether the next value is null or not.
99#[inline]
100pub fn handle_null_marker<T: DeserializableTraceData>(buf: &mut Buffer<T>) -> bool {
101    let slice = buf.as_mut_slice();
102    if slice.first() == Some(NULL_MARKER) {
103        *slice = &slice[1..];
104        true
105    } else {
106        false
107    }
108}