feat(appsec): enable api security

Author: florentinl

datadog_lambda/asm.py

@@ -185,13 +185,19 @@ def asm_start_response(
             "content-type": "application/json",
         }
 
+    if isinstance(response, dict) and "statusCode" in response:
+        body = response.get("body")
+    else:
+        body = response
+
     core.dispatch(
         # The matching listener is registered in ddtrace.appsec._handlers
         "aws_lambda.start_response",
         (
             span,
             status_code,
             response_headers,
+            body,
         ),
     )
 

datadog_lambda/wrapper.py

@@ -318,6 +318,9 @@ def _after(self, event, context):
                 if status_code:
                     self.inferred_span.set_tag("http.status_code", status_code)
 
+                if self.trigger_tags and (route := self.trigger_tags.get("http.route")):
+                    self.inferred_span.set_tag("http.route", route)
+
                 if config.service:
                     self.inferred_span.set_tag("peer.service", config.service)
 

tests/test_asm.py

@@ -106,6 +106,7 @@
         },
         "200",
         {"Content-Type": "text/html"},
+        None,
         True,
     ),
     (
@@ -123,6 +124,7 @@
             "Content-Type": "text/plain",
             "X-Error": "Not Found",
         },
+        None,
         True,
     ),
     (
@@ -140,6 +142,7 @@
             "Location": "/user/123",
             "Content-Type": "application/json",
         },
+        None,
         True,
     ),
     (
@@ -158,6 +161,7 @@
             "Content-Type": "application/json",
             "X-Custom-Header": "test-value",
         },
+        '{"message": "success"}',
         True,
     ),
     (
@@ -169,6 +173,7 @@
         },
         "200",
         {"Content-Type": "application/json"},
+        None,
         True,
     ),
     (
@@ -180,6 +185,7 @@
         },
         "200",
         {"Content-Type": "text/plain"},
+        None,
         True,
     ),
     (
@@ -188,6 +194,7 @@
         {"statusCode": 200},
         "200",
         {},
+        None,
         False,  # Should not dispatch for non-HTTP events
     ),
     (
@@ -196,6 +203,7 @@
         "Hello, World!",
         "200",
         {"content-type": "application/json"},
+        "Hello, World!",
         True,
     ),
     (
@@ -204,6 +212,7 @@
         {"message": "Hello, World!"},
         "200",
         {"content-type": "application/json"},
+        {"message": "Hello, World!"},
         True,
     ),
 ]
@@ -326,7 +335,7 @@ def test_asm_start_request_parametrized(
 
 
 @pytest.mark.parametrize(
-    "name,event_file,response,status_code,expected_headers,should_dispatch",
+    "name,event_file,response,status_code,expected_headers,expected_body,should_dispatch",
     ASM_START_RESPONSE_TEST_CASES,
 )
 @patch("datadog_lambda.asm.core")
@@ -337,6 +346,7 @@ def test_asm_start_response_parametrized(
     response,
     status_code,
     expected_headers,
+    expected_body,
     should_dispatch,
 ):
     """Test ASM start response for various HTTP event types using parametrization"""
@@ -362,11 +372,12 @@ def test_asm_start_response_parametrized(
 
         # Extract the dispatched arguments
         dispatch_args = call_args[0][1]
-        span, response_status_code, response_headers = dispatch_args
+        span, response_status_code, response_headers, response_body = dispatch_args
 
         assert span == mock_span
         assert response_status_code == status_code
         assert response_headers == expected_headers
+        assert response_body == expected_body
     else:
         # Verify core.dispatch was not called for non-HTTP events
         mock_core.dispatch.assert_not_called()