diff --git a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs b/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs index eb81944eb..c8f3eb70f 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs +++ b/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs @@ -1,5 +1,6 @@ use crate::config::aws::get_aws_partition_by_region; use crate::lifecycle::invocation::{ + base64_to_string, processor::MS_TO_NS, triggers::{ DATADOG_CARRIER_KEY, FUNCTION_TRIGGER_EVENT_SOURCE_TAG, ServiceNameResolver, Trigger, @@ -177,10 +178,24 @@ impl Trigger for SqsRecord { fn get_carrier(&self) -> HashMap { let carrier = HashMap::new(); - if let Some(ma) = self.message_attributes.get(DATADOG_CARRIER_KEY) - && let Some(string_value) = &ma.string_value - { - return serde_json::from_str(string_value).unwrap_or_default(); + if let Some(ma) = self.message_attributes.get(DATADOG_CARRIER_KEY) { + // dataType may carry a custom label suffix, e.g. "String.foo" or "Binary.foo". + if ma.data_type.starts_with("String") { + if let Some(string_value) = &ma.string_value { + return serde_json::from_str(string_value).unwrap_or_default(); + } + } else if ma.data_type.starts_with("Binary") { + if let Some(binary_value) = &ma.binary_value + && let Ok(carrier) = base64_to_string(binary_value) + { + return serde_json::from_str(&carrier).unwrap_or_default(); + } + } else { + debug!( + "Unsupported dataType in SQS message attribute: {}", + ma.data_type + ); + } } // Check for SNS event sent through SQS @@ -419,6 +434,50 @@ mod tests { assert_eq!(carrier, expected); } + #[test] + fn test_get_carrier_binary() { + let json = read_json_file("sqs_event_binary.json"); + let payload = serde_json::from_str(&json).expect("Failed to deserialize into Value"); + let event = SqsRecord::new(payload).expect("Failed to deserialize SqsRecord"); + let carrier = event.get_carrier(); + + let expected = HashMap::from([ + ( + "x-datadog-trace-id".to_string(), + "1111111111111111111".to_string(), + ), + ( + "x-datadog-parent-id".to_string(), + "2222222222222222222".to_string(), + ), + ("x-datadog-sampling-priority".to_string(), "1".to_string()), + ]); + + assert_eq!(carrier, expected); + } + + #[test] + fn test_get_carrier_binary_custom_label() { + let json = read_json_file("sqs_event_custom_label.json"); + let payload = serde_json::from_str(&json).expect("Failed to deserialize into Value"); + let event = SqsRecord::new(payload).expect("Failed to deserialize SqsRecord"); + let carrier = event.get_carrier(); + + let expected = HashMap::from([ + ( + "x-datadog-trace-id".to_string(), + "3333333333333333333".to_string(), + ), + ( + "x-datadog-parent-id".to_string(), + "4444444444444444444".to_string(), + ), + ("x-datadog-sampling-priority".to_string(), "1".to_string()), + ]); + + assert_eq!(carrier, expected); + } + #[test] fn test_get_carrier_from_sns() { let json = read_json_file("sns_sqs_event.json"); diff --git a/bottlecap/tests/payloads/sqs_event_binary.json b/bottlecap/tests/payloads/sqs_event_binary.json new file mode 100644 index 000000000..5e8dc27d3 --- /dev/null +++ b/bottlecap/tests/payloads/sqs_event_binary.json @@ -0,0 +1,27 @@ +{ + "Records": [ + { + "messageId": "19dd0b57-b21e-4ac1-bd88-01bbb068cb78", + "receiptHandle": "MessageReceiptHandle", + "body": "Hello from SQS!", + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1523232000000", + "SenderId": "123456789012", + "ApproximateFirstReceiveTimestamp": "1523232000001" + }, + "messageAttributes": { + "_datadog": { + "binaryValue": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiIxMTExMTExMTExMTExMTExMTExIiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjIyMjIyMjIyMjIyMjIyMjIyMjIiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0=", + "stringListValues": [], + "binaryListValues": [], + "dataType": "Binary" + } + }, + "md5OfBody": "{{{md5_of_body}}}", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:MyQueue", + "awsRegion": "us-east-1" + } + ] +} diff --git a/bottlecap/tests/payloads/sqs_event_custom_label.json b/bottlecap/tests/payloads/sqs_event_custom_label.json new file mode 100644 index 000000000..0501724a3 --- /dev/null +++ b/bottlecap/tests/payloads/sqs_event_custom_label.json @@ -0,0 +1,27 @@ +{ + "Records": [ + { + "messageId": "19dd0b57-b21e-4ac1-bd88-01bbb068cb78", + "receiptHandle": "MessageReceiptHandle", + "body": "Hello from SQS!", + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1523232000000", + "SenderId": "123456789012", + "ApproximateFirstReceiveTimestamp": "1523232000001" + }, + "messageAttributes": { + "_datadog": { + "binaryValue": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiIzMzMzMzMzMzMzMzMzMzMzMzMzIiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjQ0NDQ0NDQ0NDQ0NDQ0NDQ0NDQiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0=", + "stringListValues": [], + "binaryListValues": [], + "dataType": "Binary.custom" + } + }, + "md5OfBody": "{{{md5_of_body}}}", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:MyQueue", + "awsRegion": "us-east-1" + } + ] +}