Skip to content

Commit 55b16bf

Browse files
Move number, string and map decoding to separate modules
1 parent 84f8853 commit 55b16bf

File tree

9 files changed

+271
-241
lines changed

9 files changed

+271
-241
lines changed
+1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
22
// SPDX-License-Identifier: Apache-2.0
33

4+
/// Decoding logic for V04 encoded trace payload
45
pub mod v04;

trace-utils/src/msgpack_decoder/v04/decoder/mod.rs

+6-185
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,12 @@
11
// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
22
// SPDX-License-Identifier: Apache-2.0
33

4+
/// Read spans from msgpack
45
mod span;
5-
mod span_link;
66

77
use super::error::DecodeError;
8-
use super::number::{read_nullable_number_ref, read_number_ref};
98
use crate::span_v04::{Span, SpanSlice};
10-
use rmp::decode::DecodeStringError;
11-
use rmp::{decode, decode::RmpRead, Marker};
129
use span::decode_span;
13-
use std::{collections::HashMap, f64};
14-
15-
// https://docs.rs/rmp/latest/rmp/enum.Marker.html#variant.Null (0xc0 == 192)
16-
const NULL_MARKER: &u8 = &0xc0;
1710

1811
/// Decodes a Bytes buffer into a vector of `TracerPayloadV04` objects.
1912
///
@@ -152,181 +145,6 @@ pub fn from_slice(mut data: &[u8]) -> Result<(Vec<Vec<SpanSlice>>, usize), Decod
152145
))
153146
}
154147

155-
#[inline]
156-
fn read_string_nomut(buf: &[u8]) -> Result<(&str, &[u8]), DecodeError> {
157-
decode::read_str_from_slice(buf).map_err(|e| match e {
158-
DecodeStringError::InvalidMarkerRead(e) => DecodeError::InvalidFormat(e.to_string()),
159-
DecodeStringError::InvalidDataRead(e) => DecodeError::InvalidConversion(e.to_string()),
160-
DecodeStringError::TypeMismatch(marker) => {
161-
DecodeError::InvalidType(format!("Type mismatch at marker {:?}", marker))
162-
}
163-
DecodeStringError::InvalidUtf8(_, e) => DecodeError::Utf8Error(e.to_string()),
164-
_ => DecodeError::IOError,
165-
})
166-
}
167-
168-
#[inline]
169-
fn read_string<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
170-
read_string_nomut(buf).map(|(str, newbuf)| {
171-
*buf = newbuf;
172-
str
173-
})
174-
}
175-
176-
#[inline]
177-
fn read_nullable_string<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
178-
if is_null_marker(buf) {
179-
Ok("")
180-
} else {
181-
read_string(buf)
182-
}
183-
}
184-
185-
#[inline]
186-
fn read_str_map_to_str<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
187-
let len = decode::read_map_len(buf)
188-
.map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
189-
190-
let mut map = HashMap::with_capacity(len.try_into().expect("Unable to cast map len to usize"));
191-
for _ in 0..len {
192-
let key = read_string(buf)?;
193-
let value = read_string(buf)?;
194-
map.insert(key, value);
195-
}
196-
Ok(map)
197-
}
198-
199-
#[inline]
200-
fn read_nullable_str_map_to_str<'a>(
201-
buf: &mut &'a [u8],
202-
) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
203-
if is_null_marker(buf) {
204-
return Ok(HashMap::default());
205-
}
206-
207-
read_str_map_to_str(buf)
208-
}
209-
210-
#[inline]
211-
fn read_metrics<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, f64>, DecodeError> {
212-
if is_null_marker(buf) {
213-
return Ok(HashMap::default());
214-
}
215-
216-
fn read_metric_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, f64), DecodeError> {
217-
let key = read_string(buf)?;
218-
let v = read_number_ref(buf)?;
219-
220-
Ok((key, v))
221-
}
222-
223-
let len = read_map_len(buf)?;
224-
225-
read_map(len, buf, read_metric_pair)
226-
}
227-
228-
#[inline]
229-
fn read_meta_struct<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, Vec<u8>>, DecodeError> {
230-
if is_null_marker(buf) {
231-
return Ok(HashMap::default());
232-
}
233-
234-
fn read_meta_struct_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, Vec<u8>), DecodeError> {
235-
let key = read_string(buf)?;
236-
let array_len = decode::read_array_len(buf).map_err(|_| {
237-
DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned())
238-
})?;
239-
240-
let mut v = Vec::with_capacity(array_len as usize);
241-
242-
for _ in 0..array_len {
243-
let value = read_number_ref(buf)?;
244-
v.push(value);
245-
}
246-
Ok((key, v))
247-
}
248-
249-
let len = read_map_len(buf)?;
250-
read_map(len, buf, read_meta_struct_pair)
251-
}
252-
253-
/// Reads a map from the buffer and returns it as a `HashMap`.
254-
///
255-
/// This function is generic over the key and value types of the map, and it uses a provided
256-
/// function to read key-value pairs from the buffer.
257-
///
258-
/// # Arguments
259-
///
260-
/// * `len` - The number of key-value pairs to read from the buffer.
261-
/// * `buf` - A reference to the slice containing the encoded map data.
262-
/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a
263-
/// `Result<(K, V), DecodeError>`.
264-
///
265-
/// # Returns
266-
///
267-
/// * `Ok(HashMap<K, V>)` - A `HashMap` containing the decoded key-value pairs if successful.
268-
/// * `Err(DecodeError)` - An error if the decoding process fails.
269-
///
270-
/// # Errors
271-
///
272-
/// This function will return an error if:
273-
/// - The `read_pair` function returns an error while reading a key-value pair.
274-
///
275-
/// # Type Parameters
276-
///
277-
/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`.
278-
/// * `V` - The type of the values in the map.
279-
/// * `F` - The type of the function used to read key-value pairs from the buffer.
280-
#[inline]
281-
fn read_map<'a, K, V, F>(
282-
len: usize,
283-
buf: &mut &'a [u8],
284-
read_pair: F,
285-
) -> Result<HashMap<K, V>, DecodeError>
286-
where
287-
K: std::hash::Hash + Eq,
288-
F: Fn(&mut &'a [u8]) -> Result<(K, V), DecodeError>,
289-
{
290-
let mut map = HashMap::with_capacity(len);
291-
for _ in 0..len {
292-
let (k, v) = read_pair(buf)?;
293-
map.insert(k, v);
294-
}
295-
Ok(map)
296-
}
297-
298-
#[inline]
299-
fn read_map_len(buf: &mut &[u8]) -> Result<usize, DecodeError> {
300-
match decode::read_marker(buf)
301-
.map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))?
302-
{
303-
Marker::FixMap(len) => Ok(len as usize),
304-
Marker::Map16 => buf
305-
.read_data_u16()
306-
.map_err(|_| DecodeError::IOError)
307-
.map(|len| len as usize),
308-
Marker::Map32 => buf
309-
.read_data_u32()
310-
.map_err(|_| DecodeError::IOError)
311-
.map(|len| len as usize),
312-
_ => Err(DecodeError::InvalidType(
313-
"Unable to read map from buffer".to_owned(),
314-
)),
315-
}
316-
}
317-
318-
/// When you want to "peek" if the next value is a null marker, and only advance the buffer if it is
319-
/// null. If it is not null, you can continue to decode as expected.
320-
#[inline]
321-
fn is_null_marker(buf: &mut &[u8]) -> bool {
322-
if buf.first() == Some(NULL_MARKER) {
323-
*buf = &buf[1..];
324-
true
325-
} else {
326-
false
327-
}
328-
}
329-
330148
#[cfg(test)]
331149
mod tests {
332150
use super::*;
@@ -335,6 +153,7 @@ mod tests {
335153
use rmp_serde;
336154
use rmp_serde::to_vec_named;
337155
use serde_json::json;
156+
use std::collections::HashMap;
338157
use tinybytes::BytesString;
339158

340159
fn generate_meta_struct_element(i: u8) -> (String, Vec<u8>) {
@@ -352,6 +171,7 @@ mod tests {
352171

353172
(key, rmp_serde::to_vec_named(&map).unwrap())
354173
}
174+
355175
#[test]
356176
fn test_empty_array() {
357177
let encoded_data = vec![0x90];
@@ -824,7 +644,7 @@ mod tests {
824644

825645
#[test]
826646
#[cfg_attr(miri, ignore)]
827-
fn fuzz_from_slice() {
647+
fn fuzz_from_bytes() {
828648
check!()
829649
.with_type::<(
830650
String,
@@ -875,10 +695,11 @@ mod tests {
875695
start,
876696
..Default::default()
877697
};
878-
let encoded_data = to_vec_named(&vec![vec![span]]).unwrap();
698+
let encoded_data = to_vec_named(&vec![vec![span.clone()]]).unwrap();
879699
let result = from_bytes(tinybytes::Bytes::from(encoded_data));
880700

881701
assert!(result.is_ok());
702+
assert_eq!(result.unwrap().0, vec![vec![span]])
882703
},
883704
);
884705
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
use super::is_null_marker;
5+
use super::number::read_number;
6+
use super::string::read_string;
7+
use crate::msgpack_decoder::v04::error::DecodeError;
8+
use rmp::{decode, decode::RmpRead, Marker};
9+
use std::collections::HashMap;
10+
11+
/// Read a map of string to string from `buf`.
12+
#[inline]
13+
pub fn read_str_map_to_str<'a>(
14+
buf: &mut &'a [u8],
15+
) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
16+
let len = decode::read_map_len(buf)
17+
.map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
18+
19+
let mut map = HashMap::with_capacity(len.try_into().expect("Unable to cast map len to usize"));
20+
for _ in 0..len {
21+
let key = read_string(buf)?;
22+
let value = read_string(buf)?;
23+
map.insert(key, value);
24+
}
25+
Ok(map)
26+
}
27+
28+
/// Read a nullable map of string to string from `buf`.
29+
#[inline]
30+
pub fn read_nullable_str_map_to_str<'a>(
31+
buf: &mut &'a [u8],
32+
) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
33+
if is_null_marker(buf) {
34+
return Ok(HashMap::default());
35+
}
36+
37+
read_str_map_to_str(buf)
38+
}
39+
40+
/// Read a map of string to f64 from `buf`.
41+
#[inline]
42+
pub fn read_metrics<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, f64>, DecodeError> {
43+
if is_null_marker(buf) {
44+
return Ok(HashMap::default());
45+
}
46+
47+
fn read_metric_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, f64), DecodeError> {
48+
let key = read_string(buf)?;
49+
let v = read_number(buf)?;
50+
51+
Ok((key, v))
52+
}
53+
54+
let len = read_map_len(buf)?;
55+
56+
read_map(len, buf, read_metric_pair)
57+
}
58+
59+
/// Read a map of string to u8 array from `buf`.
60+
///
61+
/// The struct can't be a u8 slice since it is encoded as a msgpack array and not as a raw bytes
62+
/// buffer.
63+
#[inline]
64+
pub fn read_meta_struct<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, Vec<u8>>, DecodeError> {
65+
if is_null_marker(buf) {
66+
return Ok(HashMap::default());
67+
}
68+
69+
fn read_meta_struct_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, Vec<u8>), DecodeError> {
70+
let key = read_string(buf)?;
71+
let array_len = decode::read_array_len(buf).map_err(|_| {
72+
DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned())
73+
})?;
74+
75+
let mut v = Vec::with_capacity(array_len as usize);
76+
77+
for _ in 0..array_len {
78+
let value = read_number(buf)?;
79+
v.push(value);
80+
}
81+
Ok((key, v))
82+
}
83+
84+
let len = read_map_len(buf)?;
85+
read_map(len, buf, read_meta_struct_pair)
86+
}
87+
88+
/// Reads a map from the buffer and returns it as a `HashMap`.
89+
///
90+
/// This function is generic over the key and value types of the map, and it uses a provided
91+
/// function to read key-value pairs from the buffer.
92+
///
93+
/// # Arguments
94+
///
95+
/// * `len` - The number of key-value pairs to read from the buffer.
96+
/// * `buf` - A reference to the slice containing the encoded map data.
97+
/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a
98+
/// `Result<(K, V), DecodeError>`.
99+
///
100+
/// # Returns
101+
///
102+
/// * `Ok(HashMap<K, V>)` - A `HashMap` containing the decoded key-value pairs if successful.
103+
/// * `Err(DecodeError)` - An error if the decoding process fails.
104+
///
105+
/// # Errors
106+
///
107+
/// This function will return an error if:
108+
/// - The `read_pair` function returns an error while reading a key-value pair.
109+
///
110+
/// # Type Parameters
111+
///
112+
/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`.
113+
/// * `V` - The type of the values in the map.
114+
/// * `F` - The type of the function used to read key-value pairs from the buffer.
115+
#[inline]
116+
fn read_map<'a, K, V, F>(
117+
len: usize,
118+
buf: &mut &'a [u8],
119+
read_pair: F,
120+
) -> Result<HashMap<K, V>, DecodeError>
121+
where
122+
K: std::hash::Hash + Eq,
123+
F: Fn(&mut &'a [u8]) -> Result<(K, V), DecodeError>,
124+
{
125+
let mut map = HashMap::with_capacity(len);
126+
for _ in 0..len {
127+
let (k, v) = read_pair(buf)?;
128+
map.insert(k, v);
129+
}
130+
Ok(map)
131+
}
132+
133+
/// Read the length of a msgpack map.
134+
#[inline]
135+
fn read_map_len(buf: &mut &[u8]) -> Result<usize, DecodeError> {
136+
match decode::read_marker(buf)
137+
.map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))?
138+
{
139+
Marker::FixMap(len) => Ok(len as usize),
140+
Marker::Map16 => buf
141+
.read_data_u16()
142+
.map_err(|_| DecodeError::IOError)
143+
.map(|len| len as usize),
144+
Marker::Map32 => buf
145+
.read_data_u32()
146+
.map_err(|_| DecodeError::IOError)
147+
.map(|len| len as usize),
148+
_ => Err(DecodeError::InvalidType(
149+
"Unable to read map from buffer".to_owned(),
150+
)),
151+
}
152+
}

0 commit comments

Comments
 (0)