1
1
// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2
2
// SPDX-License-Identifier: Apache-2.0
3
3
4
+ /// Read spans from msgpack
4
5
mod span;
5
- mod span_link;
6
6
7
7
use super :: error:: DecodeError ;
8
- use super :: number:: { read_nullable_number_ref, read_number_ref} ;
9
8
use crate :: span_v04:: { Span , SpanSlice } ;
10
- use rmp:: decode:: DecodeStringError ;
11
- use rmp:: { decode, decode:: RmpRead , Marker } ;
12
9
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 ;
17
10
18
11
/// Decodes a Bytes buffer into a vector of `TracerPayloadV04` objects.
19
12
///
@@ -152,181 +145,6 @@ pub fn from_slice(mut data: &[u8]) -> Result<(Vec<Vec<SpanSlice>>, usize), Decod
152
145
) )
153
146
}
154
147
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
-
330
148
#[ cfg( test) ]
331
149
mod tests {
332
150
use super :: * ;
@@ -335,6 +153,7 @@ mod tests {
335
153
use rmp_serde;
336
154
use rmp_serde:: to_vec_named;
337
155
use serde_json:: json;
156
+ use std:: collections:: HashMap ;
338
157
use tinybytes:: BytesString ;
339
158
340
159
fn generate_meta_struct_element ( i : u8 ) -> ( String , Vec < u8 > ) {
@@ -352,6 +171,7 @@ mod tests {
352
171
353
172
( key, rmp_serde:: to_vec_named ( & map) . unwrap ( ) )
354
173
}
174
+
355
175
#[ test]
356
176
fn test_empty_array ( ) {
357
177
let encoded_data = vec ! [ 0x90 ] ;
@@ -824,7 +644,7 @@ mod tests {
824
644
825
645
#[ test]
826
646
#[ cfg_attr( miri, ignore) ]
827
- fn fuzz_from_slice ( ) {
647
+ fn fuzz_from_bytes ( ) {
828
648
check ! ( )
829
649
. with_type :: < (
830
650
String ,
@@ -875,10 +695,11 @@ mod tests {
875
695
start,
876
696
..Default :: default ( )
877
697
} ;
878
- let encoded_data = to_vec_named ( & vec ! [ vec![ span] ] ) . unwrap ( ) ;
698
+ let encoded_data = to_vec_named ( & vec ! [ vec![ span. clone ( ) ] ] ) . unwrap ( ) ;
879
699
let result = from_bytes ( tinybytes:: Bytes :: from ( encoded_data) ) ;
880
700
881
701
assert ! ( result. is_ok( ) ) ;
702
+ assert_eq ! ( result. unwrap( ) . 0 , vec![ vec![ span] ] )
882
703
} ,
883
704
) ;
884
705
}
0 commit comments