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::error::DecodeError;
5use rmp::decode;
6use rmp::decode::DecodeStringError;
7use std::collections::HashMap;
8
9// https://docs.rs/rmp/latest/rmp/enum.Marker.html#variant.Null (0xc0 == 192)
10const NULL_MARKER: &u8 = &0xc0;
11
12/// Read a string from `buf`.
13///
14/// # Errors
15/// Fails if the buffer doesn't contain a valid utf8 msgpack string.
16#[inline]
17pub fn read_string_ref_nomut(buf: &[u8]) -> Result<(&str, &[u8]), DecodeError> {
18    decode::read_str_from_slice(buf).map_err(|e| match e {
19        DecodeStringError::InvalidMarkerRead(e) => DecodeError::InvalidFormat(e.to_string()),
20        DecodeStringError::InvalidDataRead(e) => DecodeError::InvalidConversion(e.to_string()),
21        DecodeStringError::TypeMismatch(marker) => {
22            DecodeError::InvalidType(format!("Type mismatch at marker {marker:?}"))
23        }
24        DecodeStringError::InvalidUtf8(_, e) => DecodeError::Utf8Error(e.to_string()),
25        _ => DecodeError::IOError,
26    })
27}
28
29/// Read a string from the slices `buf`.
30///
31/// # Errors
32/// Fails if the buffer doesn't contain a valid utf8 msgpack string.
33#[inline]
34pub fn read_string_ref<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
35    read_string_ref_nomut(buf).map(|(str, newbuf)| {
36        *buf = newbuf;
37        str
38    })
39}
40
41/// Read a nullable string from the slices `buf`.
42///
43/// # Errors
44/// Fails if the buffer doesn't contain a valid utf8 msgpack string or a null marker.
45#[inline]
46pub fn read_nullable_string<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
47    if handle_null_marker(buf) {
48        Ok("")
49    } else {
50        read_string_ref(buf)
51    }
52}
53
54/// Read a hashmap of (string, string) 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#[inline]
60pub fn read_str_map_to_strings<'a>(
61    buf: &mut &'a [u8],
62) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
63    let len = decode::read_map_len(buf)
64        .map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
65
66    #[allow(clippy::expect_used)]
67    let mut map = HashMap::with_capacity(len.try_into().expect("Unable to cast map len to usize"));
68    for _ in 0..len {
69        let key = read_string_ref(buf)?;
70        let value = read_string_ref(buf)?;
71        map.insert(key, value);
72    }
73    Ok(map)
74}
75
76/// Read a nullable hashmap of (string, string) from the slices `buf`.
77///
78/// # Errors
79/// Fails if the buffer does not contain a valid map length prefix,
80/// or if any key or value is not a valid utf8 msgpack string.
81#[inline]
82pub fn read_nullable_str_map_to_strings<'a>(
83    buf: &mut &'a [u8],
84) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
85    if handle_null_marker(buf) {
86        return Ok(HashMap::default());
87    }
88
89    read_str_map_to_strings(buf)
90}
91
92/// Handle the null value by peeking if the next value is a null marker, and will only advance the
93/// buffer if it is null. If it is not null, you can continue to decode as expected.
94///
95/// # Returns
96/// A boolean indicating whether the next value is null or not.
97#[inline]
98pub fn handle_null_marker(buf: &mut &[u8]) -> bool {
99    if buf.first() == Some(NULL_MARKER) {
100        *buf = &buf[1..];
101        true
102    } else {
103        false
104    }
105}