libdd_trace_utils/msgpack_decoder/decode/
buffer.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::span::DeserializableTraceData;
6use rmp::decode;
7use rmp::decode::DecodeStringError;
8
9use std::borrow::Borrow;
10use std::ops::Deref;
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/// Internal Buffer used to wrap msgpack data for decoding.
30/// Provides a couple accessors to extract data from the buffer.
31pub struct Buffer<T: DeserializableTraceData>(T::Bytes);
32
33impl<T: DeserializableTraceData> Buffer<T> {
34    pub fn new(data: T::Bytes) -> Self {
35        Buffer(data)
36    }
37
38    /// Returns a mutable reference to the underlying slice.
39    pub fn as_mut_slice(&mut self) -> &mut &'static [u8] {
40        T::get_mut_slice(&mut self.0)
41    }
42
43    /// Returns an immutable reference to the underlying slice, without advancing the buffer.
44    pub fn as_slice(&self) -> &[u8] {
45        self.0.borrow()
46    }
47
48    /// Tries to extract a slice of `bytes` from the buffer and advances the buffer.
49    pub fn try_slice_and_advance(&mut self, bytes: usize) -> Option<T::Bytes> {
50        T::try_slice_and_advance(&mut self.0, bytes)
51    }
52
53    /// Read a string from the slices `buf`.
54    ///
55    /// # Errors
56    /// Fails if the buffer doesn't contain a valid utf8 msgpack string.
57    pub fn read_string(&mut self) -> Result<T::Text, DecodeError> {
58        T::read_string(&mut self.0)
59    }
60
61    /// Caps a decoded element count at the bytes remaining in the buffer. Each msgpack
62    /// element needs >=1 byte on the wire, so a length prefix can't legitimately exceed
63    /// the remaining bytes — this prevents a malicious count (e.g. 0xFFFFFFFF) from
64    /// forcing a huge pre-allocation before any element is read.
65    pub fn capped_capacity(&self, count: usize) -> usize {
66        count.min(self.len())
67    }
68}
69
70impl<T: DeserializableTraceData> Deref for Buffer<T> {
71    type Target = [u8];
72
73    fn deref(&self) -> &Self::Target {
74        self.0.borrow()
75    }
76}
libdd_trace_utils/msgpack_decoder/decode/buffer.rs

libdd_trace_utils/msgpack_decoder/decode/
buffer.rs
