Restrict the logic for trace detection in logs (#519)

* restrict the logic for trace detection in logs

* lint

* lint

* typo

Author: maxday

aws/logs_monitoring/lambda_function.py

@@ -175,9 +175,19 @@ def extract_trace_payload(event):
     try:
         message = event["message"]
         obj = json.loads(event["message"])
-        if not "traces" in obj or not isinstance(obj["traces"], list):
-            return None
-        return {"message": message, "tags": event[DD_CUSTOM_TAGS]}
+
+        obj_has_traces = "traces" in obj
+        traces_is_a_list = isinstance(obj["traces"], list)
+        # check that the log is not containing a trace array unrelated to Datadog
+        trace_id_found = (
+            len(obj["traces"]) > 0
+            and len(obj["traces"][0]) > 0
+            and obj["traces"][0][0]["trace_id"] is not None
+        )
+
+        if obj_has_traces and traces_is_a_list and trace_id_found:
+            return {"message": message, "tags": event[DD_CUSTOM_TAGS]}
+        return None
     except Exception:
         return None
 

aws/logs_monitoring/tests/test_lambda_function.py

@@ -29,6 +29,7 @@
     extract_host_from_cloudtrails,
     extract_host_from_guardduty,
     extract_host_from_route53,
+    extract_trace_payload,
     enrich,
     transform,
     split,
@@ -198,5 +199,34 @@ def test_datadog_forwarder(self, mock_get_s3_cache, mock_forward_metrics):
         del os.environ["DD_FETCH_LAMBDA_TAGS"]
 
 
+class TestLambdaFunctionExtractTracePayload(unittest.TestCase):
+    def test_extract_trace_payload_none_no_trace(self):
+        message_json = """{
+            "key": "value"
+        }"""
+        self.assertEqual(extract_trace_payload({"message": message_json}), None)
+
+    def test_extract_trace_payload_none_exception(self):
+        message_json = """{
+            "invalid_json"
+        }"""
+        self.assertEqual(extract_trace_payload({"message": message_json}), None)
+
+    def test_extract_trace_payload_unrelated_datadog_trace(self):
+        message_json = """{"traces":["I am a trace"]}"""
+        self.assertEqual(extract_trace_payload({"message": message_json}), None)
+
+    def test_extract_trace_payload_valid_trace(self):
+        message_json = """{"traces":[[{"trace_id":1234}]]}"""
+        tags_json = """["key0:value", "key1:value1"]"""
+        item = {
+            "message": '{"traces":[[{"trace_id":1234}]]}',
+            "tags": '["key0:value", "key1:value1"]',
+        }
+        self.assertEqual(
+            extract_trace_payload({"message": message_json, "ddtags": tags_json}), item
+        )
+
+
 if __name__ == "__main__":
     unittest.main()