diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index 2d864a6dab..d8ff841a35 100644 --- a/.github/workflows/fuzz.yml +++ b/.github/workflows/fuzz.yml @@ -7,7 +7,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - directory: [alloc, profiling, ddcommon-ffi] + directory: [alloc, profiling, ddcommon-ffi, trace-utils] env: CARGO_TERM_COLOR: always steps: diff --git a/Cargo.lock b/Cargo.lock index 93f5cd9303..7ca6f5cef4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1644,6 +1644,8 @@ name = "datadog-trace-utils" version = "12.0.0" dependencies = [ "anyhow", + "bolero", + "bolero-generator", "bytes", "cargo_metadata", "criterion", diff --git a/trace-utils/Cargo.toml b/trace-utils/Cargo.toml index 51b1f0a733..5e4c34fcdb 100644 --- a/trace-utils/Cargo.toml +++ b/trace-utils/Cargo.toml @@ -39,6 +39,8 @@ testcontainers = { version = "0.17.0", optional = true } cargo_metadata = { version = "0.18.1", optional = true } [dev-dependencies] +bolero = "0.10.1" +bolero-generator = "0.10.2" criterion = "0.5.1" httpmock = { version = "0.7.0"} serde_json = "1.0" diff --git a/trace-utils/src/lib.rs b/trace-utils/src/lib.rs index 70d845a08c..fffe4f8c19 100644 --- a/trace-utils/src/lib.rs +++ b/trace-utils/src/lib.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 pub mod config_utils; +pub mod msgpack_decoder; pub mod send_data; pub mod stats_utils; #[cfg(any(test, feature = "test-utils"))] diff --git a/trace-utils/src/msgpack_decoder/mod.rs b/trace-utils/src/msgpack_decoder/mod.rs new file mode 100644 index 0000000000..f4e980a0ae --- /dev/null +++ b/trace-utils/src/msgpack_decoder/mod.rs @@ -0,0 +1,4 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +pub mod v04; diff --git a/trace-utils/src/msgpack_decoder/v04/decoder/mod.rs b/trace-utils/src/msgpack_decoder/v04/decoder/mod.rs new file mode 100644 index 0000000000..76a7c7a413 --- /dev/null +++ b/trace-utils/src/msgpack_decoder/v04/decoder/mod.rs @@ -0,0 +1,541 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +mod span; +mod span_link; + +use self::span::decode_span; +use super::error::DecodeError; +use super::number::read_number; +use crate::tracer_payload::TracerPayloadV04; +use datadog_trace_protobuf::pb::Span; +use rmp::decode::DecodeStringError; +use rmp::{ + decode, + decode::{read_array_len, RmpRead}, + Marker, +}; +use std::{collections::HashMap, f64}; + +/// Decodes a slice of bytes into a vector of `TracerPayloadV04` objects. +/// +/// +/// +/// # Arguments +/// +/// * `data` - A mutable reference to a slice of bytes containing the encoded data.Bytes are +/// expected to be encoded msgpack data containing a list of a list of v04 spans. +/// +/// # Returns +/// +/// * `Ok(Vec)` - A vector of decoded `TracerPayloadV04` objects if successful. +/// * `Err(DecodeError)` - An error if the decoding process fails. +/// +/// # Errors +/// +/// This function will return an error if: +/// - The array length for trace count or span count cannot be read. +/// - Any span cannot be decoded. +/// +/// # Examples +/// +/// ``` +/// use datadog_trace_protobuf::pb::Span; +/// use datadog_trace_utils::msgpack_decoder::v04::decoder::from_slice; +/// use rmp_serde::to_vec_named; +/// +/// let span = Span { +/// name: "test-span".to_owned(), +/// ..Default::default() +/// }; +/// let encoded_data = to_vec_named(&vec![vec![span]]).unwrap(); +/// let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); +/// +/// assert_eq!(1, decoded_traces.len()); +/// assert_eq!(1, decoded_traces[0].len()); +/// let decoded_span = &decoded_traces[0][0]; +/// assert_eq!("test-span", decoded_span.name); +/// ``` +pub fn from_slice(data: &mut &[u8]) -> Result, DecodeError> { + let trace_count = read_array_len(data).map_err(|_| { + DecodeError::InvalidFormat("Unable to read array len for trace count".to_owned()) + })?; + + let mut traces: Vec = Default::default(); + + for _ in 0..trace_count { + let span_count = match read_array_len(data) { + Ok(count) => count, + Err(_) => { + return Err(DecodeError::InvalidFormat( + "Unable to read array len for span count".to_owned(), + )) + } + }; + + let mut trace: Vec = Default::default(); + + for _ in 0..span_count { + let span = decode_span(data)?; + trace.push(span); + } + traces.push(trace); + } + + Ok(traces) +} + +#[inline] +fn read_string_ref(buf: &[u8]) -> Result<(&str, &[u8]), DecodeError> { + decode::read_str_from_slice(buf).map_err(|e| match e { + DecodeStringError::InvalidMarkerRead(e) => DecodeError::InvalidFormat(e.to_string()), + DecodeStringError::InvalidDataRead(e) => DecodeError::InvalidConversion(e.to_string()), + DecodeStringError::TypeMismatch(marker) => { + DecodeError::InvalidType(format!("Type mismatch at marker {:?}", marker)) + } + DecodeStringError::InvalidUtf8(_, e) => DecodeError::Utf8Error(e.to_string()), + _ => DecodeError::IOError, + }) +} + +#[inline] +fn read_string(buf: &mut &[u8]) -> Result { + let (str_ref, remaining_buf) = read_string_ref(buf)?; + *buf = remaining_buf; + Ok(str_ref.to_string()) +} + +#[inline] +fn read_str_pair(buf: &mut &[u8]) -> Result<(String, String), DecodeError> { + let k = read_string(buf)?; + let v = read_string(buf)?; + + Ok((k, v)) +} + +#[inline] +fn read_metric_pair(buf: &mut &[u8]) -> Result<(String, f64), DecodeError> { + let k = read_string(buf)?; + let v = read_number(buf)?.try_into()?; + + Ok((k, v)) +} + +fn read_map_strs(buf: &mut &[u8]) -> Result, DecodeError> { + let len = read_map_len(buf)?; + read_map(len, buf, read_str_pair) +} + +fn read_metrics(buf: &mut &[u8]) -> Result, DecodeError> { + let len = read_map_len(buf)?; + read_map(len, buf, read_metric_pair) +} + +fn read_meta_struct(buf: &mut &[u8]) -> Result>, DecodeError> { + fn read_meta_struct_pair(buf: &mut &[u8]) -> Result<(String, Vec), DecodeError> { + let k = read_string(buf)?; + let mut v = vec![]; + let array_len = decode::read_array_len(buf).map_err(|_| { + DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned()) + })?; + for _ in 0..array_len { + let value = read_number(buf)?.try_into()?; + v.push(value); + } + Ok((k, v)) + } + + let len = read_map_len(buf)?; + read_map(len, buf, read_meta_struct_pair) +} + +/// Reads a map from the buffer and returns it as a `HashMap`. +/// +/// This function is generic over the key and value types of the map, and it uses a provided +/// function to read key-value pairs from the buffer. +/// +/// # Arguments +/// +/// * `len` - The number of key-value pairs to read from the buffer. +/// * `buf` - A mutable reference to the buffer containing the encoded map data. +/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a +/// `Result<(K, V), DecodeError>`. +/// +/// # Returns +/// +/// * `Ok(HashMap)` - A `HashMap` containing the decoded key-value pairs if successful. +/// * `Err(DecodeError)` - An error if the decoding process fails. +/// +/// # Errors +/// +/// This function will return an error if: +/// - The `read_pair` function returns an error while reading a key-value pair. +/// +/// # Type Parameters +/// +/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`. +/// * `V` - The type of the values in the map. +/// * `F` - The type of the function used to read key-value pairs from the buffer. +fn read_map( + len: usize, + buf: &mut &[u8], + read_pair: F, +) -> Result, DecodeError> +where + K: std::hash::Hash + Eq, + F: Fn(&mut &[u8]) -> Result<(K, V), DecodeError>, +{ + let mut map = HashMap::new(); + for _ in 0..len { + let (k, v) = read_pair(buf)?; + map.insert(k, v); + } + Ok(map) +} + +fn read_map_len(buf: &mut &[u8]) -> Result { + match decode::read_marker(buf) + .map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))? + { + Marker::FixMap(len) => Ok(len as usize), + Marker::Map16 => buf + .read_data_u16() + .map_err(|_| DecodeError::IOError) + .map(|len| len as usize), + Marker::Map32 => buf + .read_data_u32() + .map_err(|_| DecodeError::IOError) + .map(|len| len as usize), + _ => Err(DecodeError::InvalidType( + "Unable to read map from buffer".to_owned(), + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datadog_trace_protobuf::pb::SpanLink; + + #[test] + fn decoder_read_string_success() { + let expected_string = "test-service-name"; + let span = Span { + name: expected_string.to_owned(), + ..Default::default() + }; + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_string, decoded_span.name); + } + + #[test] + fn test_decoder_meta_struct_fixed_map_success() { + let expected_meta_struct = HashMap::from([ + ("key1".to_string(), vec![1, 2, 3]), + ("key2".to_string(), vec![4, 5, 6]), + ]); + let span = Span { + meta_struct: expected_meta_struct.clone(), + ..Default::default() + }; + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_meta_struct, decoded_span.meta_struct); + } + + #[test] + fn test_decoder_meta_struct_map_16_success() { + let expected_meta_struct: HashMap> = (0..20) + .map(|i| (format!("key {}", i), vec![1 + i, 2 + i, 3 + i])) + .collect(); + + let span = Span { + meta_struct: expected_meta_struct.clone(), + ..Default::default() + }; + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_meta_struct, decoded_span.meta_struct); + } + + #[test] + fn test_decoder_meta_fixed_map_success() { + let expected_meta = HashMap::from([ + ("key1".to_string(), "value1".to_string()), + ("key2".to_string(), "value2".to_string()), + ]); + let span = Span { + meta: expected_meta.clone(), + ..Default::default() + }; + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_meta, decoded_span.meta); + } + + #[test] + fn test_decoder_meta_map_16_success() { + let expected_meta: HashMap = (0..20) + .map(|i| (format!("key {}", i), format!("value {}", i))) + .collect(); + + let span = Span { + meta: expected_meta.clone(), + ..Default::default() + }; + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_meta, decoded_span.meta); + } + + #[test] + fn test_decoder_metrics_fixed_map_success() { + let mut span = Span::default(); + let expected_metrics = + HashMap::from([("metric1".to_string(), 1.23), ("metric2".to_string(), 4.56)]); + span.metrics = expected_metrics.clone(); + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_metrics, decoded_span.metrics); + } + + #[test] + fn test_decoder_metrics_map16_success() { + let mut span = Span::default(); + let expected_metrics: HashMap = (0..20) + .map(|i| (format!("metric{}", i), i as f64)) + .collect(); + + span.metrics = expected_metrics.clone(); + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = from_slice(&mut encoded_data.as_slice()).expect("Decoding failed"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_metrics, decoded_span.metrics); + } + + #[test] + fn test_decoder_span_link_success() { + let expected_span_links = vec![SpanLink { + trace_id: 1, + trace_id_high: 0, + span_id: 1, + attributes: HashMap::from([ + ("attr1".to_string(), "test_value".to_string()), + ("attr2".to_string(), "test_value".to_string()), + ]), + tracestate: "state_test".to_string(), + flags: 0b101, + }]; + + let span = Span { + span_links: expected_span_links.clone(), + ..Default::default() + }; + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + let decoded_traces = + from_slice(&mut encoded_data.as_slice()).expect("unable to decode span"); + + assert_eq!(1, decoded_traces.len()); + assert_eq!(1, decoded_traces[0].len()); + let decoded_span = &decoded_traces[0][0]; + assert_eq!(expected_span_links, decoded_span.span_links); + } + + #[test] + fn test_decoder_read_string_wrong_format() { + let span = Span { + service: "my_service".to_owned(), + ..Default::default() + }; + let mut encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + // This changes the map size from 11 to 12 to trigger an InvalidMarkerRead error. + encoded_data[2] = 0x8c; + + let result = from_slice(&mut encoded_data.as_slice()); + assert_eq!( + Err(DecodeError::InvalidFormat( + "Expected at least bytes 1, but only got 0 (pos 0)".to_owned() + )), + result + ); + } + + #[test] + fn test_decoder_read_string_utf8_error() { + let invalid_seq = vec![0, 159, 146, 150]; + let invalid_str = unsafe { String::from_utf8_unchecked(invalid_seq) }; + let span = Span { + name: invalid_str.to_owned(), + ..Default::default() + }; + let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + + let result = from_slice(&mut encoded_data.as_slice()); + assert_eq!( + Err(DecodeError::Utf8Error( + "invalid utf-8 sequence of 1 bytes from index 1".to_owned() + )), + result + ); + } + + #[test] + fn test_decoder_invalid_marker_for_trace_count_read() { + let span = Span::default(); + let mut encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + // This changes the entire payload to a map with 12 keys in order to trigger an error when + // reading the array len of traces + encoded_data[0] = 0x8c; + + let result = from_slice(&mut encoded_data.as_ref()); + assert_eq!( + Err(DecodeError::InvalidFormat( + "Unable to read array len for trace count".to_string() + )), + result + ); + } + + #[test] + fn test_decoder_invalid_marker_for_span_count_read() { + let span = Span::default(); + let mut encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + // This changes the entire payload to a map with 12 keys in order to trigger an error when + // reading the array len of spans + encoded_data[1] = 0x8c; + + let result = from_slice(&mut encoded_data.as_ref()); + assert_eq!( + Err(DecodeError::InvalidFormat( + "Unable to read array len for span count".to_owned() + )), + result + ); + } + + #[test] + fn test_decoder_read_string_invalid_data_read() { + let span = Span { + name: "test-span".to_owned(), + ..Default::default() + }; + let mut encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + // This changes the marker for the empty metrics map to a str8 marker + encoded_data[104] = 0xD9; + + let result = from_slice(&mut encoded_data.as_slice()); + assert_eq!( + Err(DecodeError::InvalidConversion( + "Expected at least bytes 1, but only got 0 (pos 1)".to_owned() + )), + result + ); + } + + #[test] + fn test_decoder_read_string_type_mismatch() { + let span = Span::default(); + let mut encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap(); + // Modify the encoded data to cause a type mismatch by changing the marker for the `name` + // field to an integer marker + encoded_data[3] = 0x01; + + let result = from_slice(&mut encoded_data.as_slice()); + assert_eq!( + Err(DecodeError::InvalidType( + "Type mismatch at marker FixPos(1)".to_owned() + )), + result + ); + } + + use bolero::check; + use datadog_trace_protobuf::pb::Span; + use rmp_serde::to_vec_named; + + #[test] + fn fuzz_from_slice() { + check!() + .with_type::<( + String, + String, + String, + String, + String, + String, + String, + String, + u64, + u64, + u64, + i64, + )>() + .cloned() + .for_each( + |( + name, + service, + resource, + span_type, + meta_key, + meta_value, + metric_key, + metric_value, + trace_id, + span_id, + parent_id, + start, + )| { + let span = Span { + name, + service, + resource, + r#type: span_type, + meta: HashMap::from([(meta_key, meta_value)]), + metrics: HashMap::from([( + metric_key, + metric_value.parse::().unwrap_or_default(), + )]), + trace_id, + span_id, + parent_id, + start, + ..Default::default() + }; + let encoded_data = to_vec_named(&vec![vec![span]]).unwrap(); + let result = from_slice(&mut encoded_data.as_slice()); + println!("result: {:?}", result); + + assert!(result.is_ok()); + }, + ); + } +} diff --git a/trace-utils/src/msgpack_decoder/v04/decoder/span.rs b/trace-utils/src/msgpack_decoder/v04/decoder/span.rs new file mode 100644 index 0000000000..4bd9c3d61a --- /dev/null +++ b/trace-utils/src/msgpack_decoder/v04/decoder/span.rs @@ -0,0 +1,153 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +use super::{ + read_map_strs, read_meta_struct, read_metrics, read_string, read_string_ref, + span_link::read_span_links, +}; +use crate::msgpack_decoder::v04::error::DecodeError; +use crate::msgpack_decoder::v04::number::read_number; +use datadog_trace_protobuf::pb::Span; +use std::str::FromStr; + +/// Decodes a slice of bytes into a `Span` object. +/// +/// # Arguments +/// +/// * `buf` - A mutable reference to a slice of bytes containing the encoded data. +/// +/// # Returns +/// +/// * `Ok(Span)` - A decoded `Span` object if successful. +/// * `Err(DecodeError)` - An error if the decoding process fails. +/// +/// # Errors +/// +/// This function will return an error if: +/// - The map length cannot be read. +/// - Any key or value cannot be decoded. +#[inline] +pub(crate) fn decode_span(buf: &mut &[u8]) -> Result { + let mut span = Span::default(); + let span_size = rmp::decode::read_map_len(buf).map_err(|_| { + DecodeError::InvalidFormat("Unable to get map len for span size".to_owned()) + })?; + + for _ in 0..span_size { + fill_span(&mut span, buf)?; + } + Ok(span) +} + +#[derive(Debug, PartialEq)] +enum SpanKey { + Service, + Name, + Resource, + TraceId, + SpanId, + ParentId, + Start, + Duration, + Error, + Meta, + Metrics, + Type, + MetaStruct, + SpanLinks, +} + +impl FromStr for SpanKey { + type Err = DecodeError; + + fn from_str(s: &str) -> Result { + match s { + "service" => Ok(SpanKey::Service), + "name" => Ok(SpanKey::Name), + "resource" => Ok(SpanKey::Resource), + "trace_id" => Ok(SpanKey::TraceId), + "span_id" => Ok(SpanKey::SpanId), + "parent_id" => Ok(SpanKey::ParentId), + "start" => Ok(SpanKey::Start), + "duration" => Ok(SpanKey::Duration), + "error" => Ok(SpanKey::Error), + "meta" => Ok(SpanKey::Meta), + "metrics" => Ok(SpanKey::Metrics), + "type" => Ok(SpanKey::Type), + "meta_struct" => Ok(SpanKey::MetaStruct), + "span_links" => Ok(SpanKey::SpanLinks), + _ => Err(DecodeError::InvalidFormat( + format!("Invalid span key: {}", s).to_owned(), + )), + } + } +} + +fn fill_span(span: &mut Span, buf: &mut &[u8]) -> Result<(), DecodeError> { + let (key, value) = read_string_ref(buf)?; + let key = key.parse::()?; + + *buf = value; + + match key { + SpanKey::Service => { + let value = read_string(buf)?; + span.service = value; + } + SpanKey::Name => { + let value = read_string(buf)?; + span.name = value; + } + SpanKey::Resource => { + let value = read_string(buf)?; + span.resource = value; + } + SpanKey::TraceId => span.trace_id = read_number(buf)?.try_into()?, + SpanKey::SpanId => span.span_id = read_number(buf)?.try_into()?, + SpanKey::ParentId => span.parent_id = read_number(buf)?.try_into()?, + SpanKey::Start => span.start = read_number(buf)?.try_into()?, + SpanKey::Duration => span.duration = read_number(buf)?.try_into()?, + SpanKey::Error => span.error = read_number(buf)?.try_into()?, + SpanKey::Meta => span.meta = read_map_strs(buf)?, + SpanKey::Metrics => span.metrics = read_metrics(buf)?, + SpanKey::Type => { + let value = read_string(buf)?; + span.r#type = value; + } + SpanKey::MetaStruct => span.meta_struct = read_meta_struct(buf)?, + SpanKey::SpanLinks => span.span_links = read_span_links(buf)?, + } + Ok(()) +} +#[cfg(test)] +mod tests { + use super::SpanKey; + use crate::msgpack_decoder::v04::error::DecodeError; + use std::str::FromStr; + + #[test] + fn test_span_key_from_str() { + assert_eq!(SpanKey::from_str("service").unwrap(), SpanKey::Service); + assert_eq!(SpanKey::from_str("name").unwrap(), SpanKey::Name); + assert_eq!(SpanKey::from_str("resource").unwrap(), SpanKey::Resource); + assert_eq!(SpanKey::from_str("trace_id").unwrap(), SpanKey::TraceId); + assert_eq!(SpanKey::from_str("span_id").unwrap(), SpanKey::SpanId); + assert_eq!(SpanKey::from_str("parent_id").unwrap(), SpanKey::ParentId); + assert_eq!(SpanKey::from_str("start").unwrap(), SpanKey::Start); + assert_eq!(SpanKey::from_str("duration").unwrap(), SpanKey::Duration); + assert_eq!(SpanKey::from_str("error").unwrap(), SpanKey::Error); + assert_eq!(SpanKey::from_str("meta").unwrap(), SpanKey::Meta); + assert_eq!(SpanKey::from_str("metrics").unwrap(), SpanKey::Metrics); + assert_eq!(SpanKey::from_str("type").unwrap(), SpanKey::Type); + assert_eq!( + SpanKey::from_str("meta_struct").unwrap(), + SpanKey::MetaStruct + ); + assert_eq!(SpanKey::from_str("span_links").unwrap(), SpanKey::SpanLinks); + + assert!(matches!( + SpanKey::from_str("invalid_key"), + Err(DecodeError::InvalidFormat(_)) + )); + } +} diff --git a/trace-utils/src/msgpack_decoder/v04/decoder/span_link.rs b/trace-utils/src/msgpack_decoder/v04/decoder/span_link.rs new file mode 100644 index 0000000000..07a48d7233 --- /dev/null +++ b/trace-utils/src/msgpack_decoder/v04/decoder/span_link.rs @@ -0,0 +1,134 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +use crate::msgpack_decoder::v04::decoder::{read_map_strs, read_string, read_string_ref}; +use crate::msgpack_decoder::v04::error::DecodeError; +use crate::msgpack_decoder::v04::number::read_number; +use datadog_trace_protobuf::pb::SpanLink; +use rmp::Marker; +use std::str::FromStr; + +/// Reads a slice of bytes and decodes it into a vector of `SpanLink` objects. +/// +/// # Arguments +/// +/// * `buf` - A mutable reference to a slice of bytes containing the encoded data. +/// +/// # Returns +/// +/// * `Ok(Vec)` - A vector of decoded `SpanLink` objects if successful. +/// * `Err(DecodeError)` - An error if the decoding process fails. +/// +/// # Errors +/// +/// This function will return an error if: +/// - The marker for the array length cannot be read. +/// - Any `SpanLink` cannot be decoded. +/// ``` +pub(crate) fn read_span_links(buf: &mut &[u8]) -> Result, DecodeError> { + match rmp::decode::read_marker(buf).map_err(|_| { + DecodeError::InvalidFormat("Unable to read marker for span links".to_owned()) + })? { + Marker::FixArray(len) => { + let mut vec: Vec = Vec::new(); + for _ in 0..len { + vec.push(decode_span_link(buf)?); + } + Ok(vec) + } + _ => Err(DecodeError::InvalidType( + "Unable to read span link from buffer".to_owned(), + )), + } +} +#[derive(Debug, PartialEq)] +enum SpanLinkKey { + TraceId, + TraceIdHigh, + SpanId, + Attributes, + Tracestate, + Flags, +} +impl FromStr for SpanLinkKey { + type Err = DecodeError; + + fn from_str(s: &str) -> Result { + match s { + "trace_id" => Ok(SpanLinkKey::TraceId), + "trace_id_high" => Ok(SpanLinkKey::TraceIdHigh), + "span_id" => Ok(SpanLinkKey::SpanId), + "attributes" => Ok(SpanLinkKey::Attributes), + "tracestate" => Ok(SpanLinkKey::Tracestate), + "flags" => Ok(SpanLinkKey::Flags), + _ => Err(DecodeError::InvalidFormat( + format!("Invalid span link key: {}", s).to_owned(), + )), + } + } +} + +fn decode_span_link(buf: &mut &[u8]) -> Result { + let mut span = SpanLink::default(); + let span_size = rmp::decode::read_map_len(buf) + .map_err(|_| DecodeError::InvalidType("Unable to get map len for span size".to_owned()))?; + + for _ in 0..span_size { + let (key, value) = read_string_ref(buf)?; + *buf = value; + let key = key.parse::()?; + + match key { + SpanLinkKey::TraceId => span.trace_id = read_number(buf)?.try_into()?, + SpanLinkKey::TraceIdHigh => span.trace_id_high = read_number(buf)?.try_into()?, + SpanLinkKey::SpanId => span.span_id = read_number(buf)?.try_into()?, + SpanLinkKey::Attributes => span.attributes = read_map_strs(buf)?, + SpanLinkKey::Tracestate => { + let value = read_string(buf)?; + span.tracestate = value; + } + SpanLinkKey::Flags => span.flags = read_number(buf)?.try_into()?, + } + } + + Ok(span) +} + +#[cfg(test)] +mod tests { + use super::SpanLinkKey; + use crate::msgpack_decoder::v04::error::DecodeError; + use std::str::FromStr; + + #[test] + fn test_span_link_key_from_str() { + // Valid cases + assert_eq!( + SpanLinkKey::from_str("trace_id").unwrap(), + SpanLinkKey::TraceId + ); + assert_eq!( + SpanLinkKey::from_str("trace_id_high").unwrap(), + SpanLinkKey::TraceIdHigh + ); + assert_eq!( + SpanLinkKey::from_str("span_id").unwrap(), + SpanLinkKey::SpanId + ); + assert_eq!( + SpanLinkKey::from_str("attributes").unwrap(), + SpanLinkKey::Attributes + ); + assert_eq!( + SpanLinkKey::from_str("tracestate").unwrap(), + SpanLinkKey::Tracestate + ); + assert_eq!(SpanLinkKey::from_str("flags").unwrap(), SpanLinkKey::Flags); + + // Invalid case + assert!(matches!( + SpanLinkKey::from_str("invalid_key"), + Err(DecodeError::InvalidFormat(_)) + )); + } +} diff --git a/trace-utils/src/msgpack_decoder/v04/error.rs b/trace-utils/src/msgpack_decoder/v04/error.rs new file mode 100644 index 0000000000..ff74819f57 --- /dev/null +++ b/trace-utils/src/msgpack_decoder/v04/error.rs @@ -0,0 +1,23 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#[derive(Debug, PartialEq)] +pub enum DecodeError { + InvalidConversion(String), + InvalidType(String), + InvalidFormat(String), + IOError, + Utf8Error(String), +} + +impl std::fmt::Display for DecodeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DecodeError::InvalidConversion(msg) => write!(f, "Failed to convert value: {}", msg), + DecodeError::IOError => write!(f, "Failed to read from buffer"), + DecodeError::InvalidType(msg) => write!(f, "Invalid type encountered: {}", msg), + DecodeError::InvalidFormat(msg) => write!(f, "Invalid format: {}", msg), + DecodeError::Utf8Error(msg) => write!(f, "Failed to read utf8 value: {}", msg), + } + } +} diff --git a/trace-utils/src/msgpack_decoder/v04/mod.rs b/trace-utils/src/msgpack_decoder/v04/mod.rs new file mode 100644 index 0000000000..5b789e7c18 --- /dev/null +++ b/trace-utils/src/msgpack_decoder/v04/mod.rs @@ -0,0 +1,6 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +pub mod decoder; +pub mod error; +pub mod number; diff --git a/trace-utils/src/msgpack_decoder/v04/number.rs b/trace-utils/src/msgpack_decoder/v04/number.rs new file mode 100644 index 0000000000..9a4f39cfc8 --- /dev/null +++ b/trace-utils/src/msgpack_decoder/v04/number.rs @@ -0,0 +1,550 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +use super::error::DecodeError; +use rmp::{decode::RmpRead, Marker}; +use std::fmt; + +#[derive(Debug, PartialEq)] +pub enum Number { + Unsigned(u64), + Signed(i64), + Float(f64), +} + +impl fmt::Display for Number { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Number::Signed(val) => write!(f, "{}", val), + Number::Unsigned(val) => write!(f, "{}", val), + Number::Float(val) => write!(f, "{}", val), + } + } +} + +impl Number { + pub fn bounded_int_conversion( + self, + lower_bound: T, + upper_bound: Option, + ) -> Result + where + T: TryInto + TryInto + TryInto + Copy + fmt::Display, + i64: TryInto, + u64: TryInto, + >::Error: fmt::Debug, + >::Error: fmt::Debug, + >::Error: fmt::Debug, + >::Error: fmt::Debug, + { + match self { + Number::Signed(val) => { + let upper_bound_check = if let Some(upper_bound) = upper_bound { + val <= upper_bound.try_into().unwrap() + } else { + true + }; + if val >= lower_bound.try_into().unwrap() && upper_bound_check { + val.try_into() + .map_err(|e| DecodeError::InvalidConversion(format!("{:?}", e))) + } else { + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + val + ))) + } + } + Number::Unsigned(val) => { + let upper_bound_check = if let Some(upper_bound) = upper_bound { + val <= upper_bound.try_into().unwrap() + } else { + true + }; + + if upper_bound_check { + val.try_into() + .map_err(|e| DecodeError::InvalidConversion(format!("{:?}", e))) + } else { + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + val + ))) + } + } + _ => Err(DecodeError::InvalidConversion( + "Cannot convert float to int".to_owned(), + )), + } + } +} + +impl TryFrom for i8 { + type Error = DecodeError; + fn try_from(value: Number) -> Result { + value.bounded_int_conversion(i8::MIN, Some(i8::MAX)) + } +} + +impl TryFrom for i32 { + type Error = DecodeError; + fn try_from(value: Number) -> Result { + value.bounded_int_conversion(i32::MIN, Some(i32::MAX)) + } +} + +impl TryFrom for i64 { + type Error = DecodeError; + fn try_from(value: Number) -> Result { + value.bounded_int_conversion(i64::MIN, Some(i64::MAX)) + } +} + +impl TryFrom for u8 { + type Error = DecodeError; + fn try_from(value: Number) -> Result { + value.bounded_int_conversion(u8::MIN, Some(u8::MAX)) + } +} + +impl TryFrom for u32 { + type Error = DecodeError; + fn try_from(value: Number) -> Result { + value.bounded_int_conversion(u32::MIN, Some(u32::MAX)) + } +} + +impl TryFrom for u64 { + type Error = DecodeError; + fn try_from(value: Number) -> Result { + value.bounded_int_conversion(u64::MIN, None) + } +} + +impl TryFrom for f64 { + type Error = DecodeError; + fn try_from(value: Number) -> Result { + match value { + Number::Unsigned(val) => { + if val <= f64::MAX as u64 { + Ok(val as f64) + } else { + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + val + ))) + } + } + Number::Signed(val) => { + if val >= f64::MIN as i64 && val <= f64::MAX as i64 { + Ok(val as f64) + } else { + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + val + ))) + } + } + Number::Float(val) => Ok(val), + } + } +} + +pub fn read_number(buf: &mut &[u8]) -> Result { + match rmp::decode::read_marker(buf) + .map_err(|_| DecodeError::InvalidFormat("Unable to read marker for number".to_owned()))? + { + Marker::FixPos(val) => Ok(Number::Unsigned(val as u64)), + Marker::FixNeg(val) => Ok(Number::Signed(val as i64)), + Marker::U8 => Ok(Number::Unsigned( + buf.read_data_u8().map_err(|_| DecodeError::IOError)? as u64, + )), + Marker::U16 => Ok(Number::Unsigned( + buf.read_data_u16().map_err(|_| DecodeError::IOError)? as u64, + )), + Marker::U32 => Ok(Number::Unsigned( + buf.read_data_u32().map_err(|_| DecodeError::IOError)? as u64, + )), + Marker::U64 => Ok(Number::Unsigned( + buf.read_data_u64().map_err(|_| DecodeError::IOError)?, + )), + Marker::I8 => Ok(Number::Signed( + buf.read_data_i8().map_err(|_| DecodeError::IOError)? as i64, + )), + Marker::I16 => Ok(Number::Signed( + buf.read_data_i16().map_err(|_| DecodeError::IOError)? as i64, + )), + Marker::I32 => Ok(Number::Signed( + buf.read_data_i32().map_err(|_| DecodeError::IOError)? as i64, + )), + Marker::I64 => Ok(Number::Signed( + buf.read_data_i64().map_err(|_| DecodeError::IOError)?, + )), + Marker::F32 => Ok(Number::Float( + buf.read_data_f32().map_err(|_| DecodeError::IOError)? as f64, + )), + Marker::F64 => Ok(Number::Float( + buf.read_data_f64().map_err(|_| DecodeError::IOError)?, + )), + _ => Err(DecodeError::InvalidType("Invalid number type".to_owned())), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::f64; + + #[test] + fn test_i64_conversions() { + let valid_max = i64::MAX; + let valid_unsigned_number = Number::Unsigned(valid_max as u64); + let zero_unsigned = Number::Unsigned(0u64); + let zero_signed = Number::Unsigned(0u64); + let valid_signed_number = Number::Signed(valid_max); + let invalid_float_number = Number::Float(4.14); + let invalid_unsigned = u64::MAX; + let invalid_unsigned_number = Number::Unsigned(invalid_unsigned); + + assert_eq!( + valid_max, + TryInto::::try_into(valid_unsigned_number).unwrap() + ); + assert_eq!( + valid_max, + TryInto::::try_into(valid_signed_number).unwrap() + ); + assert_eq!(0, TryInto::::try_into(zero_signed).unwrap()); + assert_eq!(0, TryInto::::try_into(zero_unsigned).unwrap()); + assert_eq!( + Err(DecodeError::InvalidConversion( + "Cannot convert float to int".to_owned() + )), + TryInto::::try_into(invalid_float_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_unsigned + ))), + TryInto::::try_into(invalid_unsigned_number) + ); + } + #[test] + fn test_i32_conversions() { + let valid_signed_upper = i32::MAX; + let valid_unsigned_number = Number::Unsigned(valid_signed_upper as u64); + let zero_unsigned = Number::Unsigned(0u64); + let zero_signed = Number::Unsigned(0u64); + let valid_signed_number_upper = Number::Signed(valid_signed_upper as i64); + let valid_signed_lower = i32::MIN; + let valid_signed_number_lower = Number::Signed(valid_signed_lower as i64); + let invalid_float_number = Number::Float(4.14); + let invalid_unsigned = u64::MAX; + let invalid_unsigned_number = Number::Unsigned(invalid_unsigned); + let invalid_signed_upper = i32::MAX as i64 + 1; + let invalid_signed_number_upper = Number::Signed(invalid_signed_upper); + let invalid_signed_lower = i32::MIN as i64 - 1; + let invalid_signed_number_lower = Number::Signed(invalid_signed_lower); + + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_unsigned_number).unwrap() + ); + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_signed_number_upper).unwrap() + ); + assert_eq!( + valid_signed_lower, + TryInto::::try_into(valid_signed_number_lower).unwrap() + ); + assert_eq!(0, TryInto::::try_into(zero_signed).unwrap()); + assert_eq!(0, TryInto::::try_into(zero_unsigned).unwrap()); + assert_eq!( + Err(DecodeError::InvalidConversion( + "Cannot convert float to int".to_owned() + )), + TryInto::::try_into(invalid_float_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_unsigned + ))), + TryInto::::try_into(invalid_unsigned_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_upper + ))), + TryInto::::try_into(invalid_signed_number_upper) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_lower + ))), + TryInto::::try_into(invalid_signed_number_lower) + ); + } + + #[test] + fn test_i8_conversions() { + let valid_signed_upper = i8::MAX; + let valid_unsigned_number = Number::Unsigned(valid_signed_upper as u64); + let zero_unsigned = Number::Unsigned(0u64); + let zero_signed = Number::Unsigned(0u64); + let valid_signed_number_upper = Number::Signed(valid_signed_upper as i64); + let valid_signed_lower = i8::MIN; + let valid_signed_number_lower = Number::Signed(valid_signed_lower as i64); + let invalid_float_number = Number::Float(4.14); + let invalid_unsigned = u8::MAX; + let invalid_unsigned_number = Number::Unsigned(invalid_unsigned as u64); + let invalid_signed_upper = i8::MAX as i64 + 1; + let invalid_signed_number_upper = Number::Signed(invalid_signed_upper); + let invalid_signed_lower = i8::MIN as i64 - 1; + let invalid_signed_number_lower = Number::Signed(invalid_signed_lower); + + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_unsigned_number).unwrap() + ); + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_signed_number_upper).unwrap() + ); + assert_eq!( + valid_signed_lower, + TryInto::::try_into(valid_signed_number_lower).unwrap() + ); + assert_eq!(0, TryInto::::try_into(zero_signed).unwrap()); + assert_eq!(0, TryInto::::try_into(zero_unsigned).unwrap()); + assert_eq!( + Err(DecodeError::InvalidConversion( + "Cannot convert float to int".to_owned() + )), + TryInto::::try_into(invalid_float_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_unsigned + ))), + TryInto::::try_into(invalid_unsigned_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_upper + ))), + TryInto::::try_into(invalid_signed_number_upper) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_lower + ))), + TryInto::::try_into(invalid_signed_number_lower) + ); + } + + #[test] + fn test_u8_conversions() { + let valid_signed_upper = u8::MAX; + let valid_unsigned_number = Number::Unsigned(valid_signed_upper as u64); + let zero_unsigned = Number::Unsigned(0u64); + let zero_signed = Number::Unsigned(0u64); + let valid_signed_number_upper = Number::Signed(valid_signed_upper as i64); + let valid_signed_lower = u8::MIN; + let valid_signed_number_lower = Number::Signed(valid_signed_lower as i64); + let invalid_float_number = Number::Float(4.14); + let invalid_unsigned = (u8::MAX as u64) + 1; + let invalid_unsigned_number = Number::Unsigned(invalid_unsigned); + let invalid_signed_upper = i32::MAX as i64 + 1; + let invalid_signed_number_upper = Number::Signed(invalid_signed_upper); + let invalid_signed_lower = i8::MIN as i64; + let invalid_signed_number_lower = Number::Signed(invalid_signed_lower); + + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_unsigned_number).unwrap() + ); + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_signed_number_upper).unwrap() + ); + assert_eq!( + valid_signed_lower, + TryInto::::try_into(valid_signed_number_lower).unwrap() + ); + assert_eq!(0, TryInto::::try_into(zero_signed).unwrap()); + assert_eq!(0, TryInto::::try_into(zero_unsigned).unwrap()); + assert_eq!( + Err(DecodeError::InvalidConversion( + "Cannot convert float to int".to_owned() + )), + TryInto::::try_into(invalid_float_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_unsigned + ))), + TryInto::::try_into(invalid_unsigned_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_upper + ))), + TryInto::::try_into(invalid_signed_number_upper) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_lower + ))), + TryInto::::try_into(invalid_signed_number_lower) + ); + } + + #[test] + fn test_u32_conversions() { + let valid_signed_upper = u32::MAX; + let valid_unsigned_number = Number::Unsigned(valid_signed_upper as u64); + let zero_unsigned = Number::Unsigned(0u64); + let zero_signed = Number::Unsigned(0u64); + let valid_signed_number_upper = Number::Signed(valid_signed_upper as i64); + let valid_signed_lower = u32::MIN; + let valid_signed_number_lower = Number::Signed(valid_signed_lower as i64); + let invalid_float_number = Number::Float(4.14); + let invalid_unsigned = (u32::MAX as u64) + 1; + let invalid_unsigned_number = Number::Unsigned(invalid_unsigned); + let invalid_signed_upper = i64::MAX; + let invalid_signed_number_upper = Number::Signed(invalid_signed_upper); + let invalid_signed_lower = i8::MIN as i64; + let invalid_signed_number_lower = Number::Signed(invalid_signed_lower); + + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_unsigned_number).unwrap() + ); + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_signed_number_upper).unwrap() + ); + assert_eq!( + valid_signed_lower, + TryInto::::try_into(valid_signed_number_lower).unwrap() + ); + assert_eq!(0, TryInto::::try_into(zero_signed).unwrap()); + assert_eq!(0, TryInto::::try_into(zero_unsigned).unwrap()); + assert_eq!( + Err(DecodeError::InvalidConversion( + "Cannot convert float to int".to_owned() + )), + TryInto::::try_into(invalid_float_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_unsigned + ))), + TryInto::::try_into(invalid_unsigned_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_upper + ))), + TryInto::::try_into(invalid_signed_number_upper) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_lower + ))), + TryInto::::try_into(invalid_signed_number_lower) + ); + } + + #[test] + fn test_u64_conversions() { + let valid_unsigned = u64::MAX; + let valid_unsigned_number = Number::Unsigned(valid_unsigned); + let zero_unsigned = Number::Unsigned(0u64); + let zero_signed = Number::Unsigned(0u64); + let valid_signed_upper = i64::MAX as u64; + let valid_signed_number_upper = Number::Signed(valid_signed_upper as i64); + let valid_signed_lower = u32::MIN as u64; + let valid_signed_number_lower = Number::Signed(valid_signed_lower as i64); + let invalid_float_number = Number::Float(4.14); + let invalid_signed_lower = i8::MIN as i64; + let invalid_signed_number_lower = Number::Signed(invalid_signed_lower); + + assert_eq!( + valid_unsigned, + TryInto::::try_into(valid_unsigned_number).unwrap() + ); + assert_eq!( + valid_signed_upper, + TryInto::::try_into(valid_signed_number_upper).unwrap() + ); + assert_eq!( + valid_signed_lower, + TryInto::::try_into(valid_signed_number_lower).unwrap() + ); + assert_eq!(0, TryInto::::try_into(zero_signed).unwrap()); + assert_eq!(0, TryInto::::try_into(zero_unsigned).unwrap()); + assert_eq!( + Err(DecodeError::InvalidConversion( + "Cannot convert float to int".to_owned() + )), + TryInto::::try_into(invalid_float_number) + ); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_signed_lower + ))), + TryInto::::try_into(invalid_signed_number_lower) + ); + } + + #[test] + fn test_f64_conversions() { + let valid_signed_upper = i64::MAX; + let valid_unsigned_upper = u64::MAX; + let valid_signed_number_upper = Number::Signed(valid_signed_upper); + let valid_signed_lower = i64::MIN; + let valid_signed_number_lower = Number::Signed(valid_signed_lower); + let valid_unsigned_number = Number::Unsigned(valid_unsigned_upper); + let zero_unsigned = Number::Unsigned(0u64); + let zero_signed = Number::Unsigned(0u64); + let invalid_unsigned = u64::MAX; + let invalid_unsigned_number = Number::Unsigned(invalid_unsigned); + + assert_eq!( + valid_unsigned_upper as f64, + TryInto::::try_into(valid_unsigned_number).unwrap() + ); + assert_eq!( + valid_signed_upper as f64, + TryInto::::try_into(valid_signed_number_upper).unwrap() + ); + assert_eq!( + valid_signed_lower as f64, + TryInto::::try_into(valid_signed_number_lower).unwrap() + ); + assert_eq!(0f64, TryInto::::try_into(zero_signed).unwrap()); + assert_eq!(0f64, TryInto::::try_into(zero_unsigned).unwrap()); + assert_eq!( + Err(DecodeError::InvalidConversion(format!( + "{} is out of bounds for conversion", + invalid_unsigned + ))), + TryInto::::try_into(invalid_unsigned_number) + ); + } +} diff --git a/trace-utils/src/tracer_payload.rs b/trace-utils/src/tracer_payload.rs index 5e2be722ef..8ae0a963fc 100644 --- a/trace-utils/src/tracer_payload.rs +++ b/trace-utils/src/tracer_payload.rs @@ -1,11 +1,14 @@ // Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 -use crate::trace_utils::{cmp_send_data_payloads, collect_trace_chunks, TracerHeaderTags}; +use crate::{ + msgpack_decoder, + trace_utils::{cmp_send_data_payloads, collect_trace_chunks, TracerHeaderTags}, +}; use datadog_trace_protobuf::pb; use std::cmp::Ordering; -type TracerPayloadV04 = Vec; +pub type TracerPayloadV04 = Vec; #[derive(Debug, Clone)] /// Enumerates the different encoding types. @@ -241,12 +244,18 @@ impl<'a, T: TraceChunkProcessor + 'a> TryInto fn try_into(self) -> Result { match self.encoding_type { TraceEncoding::V04 => { - let traces: Vec> = match rmp_serde::from_slice(self.data) { - Ok(res) => res, - Err(e) => { - anyhow::bail!("Error deserializing trace from request body: {e}") - } - }; + // msgpack::decoder::from_slice requires a mutable ref to self.data, so we need to + // create a local copy of the ref to make the ref mutable + let mut data_slice: &[u8] = self.data; + let data: &mut &[u8] = &mut data_slice; + + let traces: Vec> = + match msgpack_decoder::v04::decoder::from_slice(data) { + Ok(res) => res, + Err(e) => { + anyhow::bail!("Error deserializing trace from request body: {e}") + } + }; if traces.is_empty() { anyhow::bail!("No traces deserialized from the request body."); @@ -294,6 +303,15 @@ mod tests { }]) } + fn create_trace() -> Vec { + vec![ + // create a root span with metrics + create_test_span(1234, 12341, 0, 1, true), + create_test_span(1234, 12342, 12341, 1, false), + create_test_span(1234, 12343, 12342, 1, false), + ] + } + #[test] fn test_append_traces_v07() { let mut trace = create_dummy_collection_v07(); @@ -357,6 +375,7 @@ mod tests { "error": 0, "meta": {}, "metrics": {}, + "type": "serverless", }]); let expected_serialized_span_data1 = vec![pb::Span { @@ -372,7 +391,7 @@ mod tests { meta: HashMap::new(), metrics: HashMap::new(), meta_struct: HashMap::new(), - r#type: "".to_string(), + r#type: "serverless".to_string(), span_links: vec![], }]; @@ -388,6 +407,7 @@ mod tests { "error": 1, "meta": {}, "metrics": {}, + "type": "", }]); let expected_serialized_span_data2 = vec![pb::Span { @@ -433,4 +453,31 @@ mod tests { panic!("Invalid collection type returned for try_into"); } } + + #[test] + fn test_try_into_meta_metrics_success() { + let dummy_trace = create_trace(); + let expected = vec![dummy_trace.clone()]; + let payload = rmp_serde::to_vec_named(&expected).unwrap(); + let tracer_header_tags = &TracerHeaderTags::default(); + + let result: anyhow::Result = TracerPayloadParams::new( + &payload, + tracer_header_tags, + &mut DefaultTraceChunkProcessor, + false, + TraceEncoding::V04, + ) + .try_into(); + + assert!(result.is_ok()); + + let collection = result.unwrap(); + assert_eq!(1, collection.size()); + if let TracerPayloadCollection::V04(traces) = collection { + assert_eq!(dummy_trace, traces[0]); + } else { + panic!("Invalid collection type returned for try_into"); + } + } }