fix(appsec): perform _after extractions on the blocking_response (#656)

Author: florentinl

datadog_lambda/wrapper.py

@@ -284,6 +284,34 @@ def _before(self, event, context):
     def _after(self, event, context):
         try:
             status_code = extract_http_status_code_tag(self.trigger_tags, self.response)
+
+            if self.span:
+                if config.appsec_enabled and not self.blocking_response:
+                    asm_start_response(
+                        self.span,
+                        status_code,
+                        self.event_source,
+                        response=self.response,
+                    )
+                    self.blocking_response = get_asm_blocked_response(self.event_source)
+
+                if self.blocking_response:
+                    status_code = str(self.blocking_response.get("statusCode"))
+                    if config.capture_payload_enabled and self.response:
+                        tag_object(
+                            self.span, "function.blocked_response", self.response
+                        )
+                    self.response = self.blocking_response
+
+                if config.capture_payload_enabled:
+                    tag_object(self.span, "function.request", event)
+                    tag_object(self.span, "function.response", self.response)
+
+                if status_code:
+                    self.span.set_tag("http.status_code", status_code)
+
+                self.span.finish()
+
             if status_code:
                 self.trigger_tags["http.status_code"] = status_code
                 mark_trace_as_error_for_5xx_responses(context, status_code, self.span)
@@ -298,25 +326,6 @@ def _after(self, event, context):
             if should_trace_cold_start:
                 trace_ctx = tracer.current_trace_context()
 
-            if self.span:
-                if config.capture_payload_enabled:
-                    tag_object(self.span, "function.request", event)
-                    tag_object(self.span, "function.response", self.response)
-
-                if status_code:
-                    self.span.set_tag("http.status_code", status_code)
-
-                if config.appsec_enabled and not self.blocking_response:
-                    asm_start_response(
-                        self.span,
-                        status_code,
-                        self.event_source,
-                        response=self.response,
-                    )
-                    self.blocking_response = get_asm_blocked_response(self.event_source)
-
-                self.span.finish()
-
             if self.inferred_span:
                 if status_code:
                     self.inferred_span.set_tag("http.status_code", status_code)

tests/test_wrapper.py

@@ -690,6 +690,9 @@ def setUp(self):
         self.mock_get_asm_blocking_response = patcher.start()
         self.addCleanup(patcher.stop)
 
+        with open("tests/event_samples/api-gateway.json") as f:
+            self.api_gateway_request = json.loads(f.read())
+
         self.fake_blocking_response = {
             "statusCode": "403",
             "headers": {
@@ -706,7 +709,7 @@ def test_blocking_before(self):
 
         lambda_handler = wrapper.datadog_lambda_wrapper(mock_handler)
 
-        response = lambda_handler({}, get_mock_context())
+        response = lambda_handler(self.api_gateway_request, get_mock_context())
         self.assertEqual(response, self.fake_blocking_response)
 
         mock_handler.assert_not_called()
@@ -715,27 +718,31 @@ def test_blocking_before(self):
         self.mock_asm_start_request.assert_called_once()
         self.mock_asm_start_response.assert_not_called()
 
+        assert lambda_handler.span.get_tag("http.status_code") == "403"
+
     def test_blocking_during(self):
         self.mock_get_asm_blocking_response.return_value = None
 
-        @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):
             self.mock_get_asm_blocking_response.return_value = (
                 self.fake_blocking_response
             )
             raise wrapper.BlockingException()
 
-        response = lambda_handler({}, get_mock_context())
+        lambda_handler = wrapper.datadog_lambda_wrapper(lambda_handler)
+
+        response = lambda_handler(self.api_gateway_request, get_mock_context())
         self.assertEqual(response, self.fake_blocking_response)
 
         self.mock_asm_set_context.assert_called_once()
         self.mock_asm_start_request.assert_called_once()
         self.mock_asm_start_response.assert_not_called()
 
+        assert lambda_handler.span.get_tag("http.status_code") == "403"
+
     def test_blocking_after(self):
         self.mock_get_asm_blocking_response.return_value = None
 
-        @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):
             self.mock_get_asm_blocking_response.return_value = (
                 self.fake_blocking_response
@@ -745,13 +752,17 @@ def lambda_handler(event, context):
                 "body": "This should not be returned",
             }
 
-        response = lambda_handler({}, get_mock_context())
+        lambda_handler = wrapper.datadog_lambda_wrapper(lambda_handler)
+
+        response = lambda_handler(self.api_gateway_request, get_mock_context())
         self.assertEqual(response, self.fake_blocking_response)
 
         self.mock_asm_set_context.assert_called_once()
         self.mock_asm_start_request.assert_called_once()
         self.mock_asm_start_response.assert_called_once()
 
+        assert lambda_handler.span.get_tag("http.status_code") == "403"
+
     def test_no_blocking_appsec_disabled(self):
         os.environ["DD_APPSEC_ENABLED"] = "false"
 
@@ -764,14 +775,17 @@ def test_no_blocking_appsec_disabled(self):
             "body": "This should be returned",
         }
 
-        @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):
             return expected_response
 
-        response = lambda_handler({}, get_mock_context())
+        lambda_handler = wrapper.datadog_lambda_wrapper(lambda_handler)
+
+        response = lambda_handler(self.api_gateway_request, get_mock_context())
         self.assertEqual(response, expected_response)
 
         self.mock_get_asm_blocking_response.assert_not_called()
         self.mock_asm_set_context.assert_not_called()
         self.mock_asm_start_request.assert_not_called()
         self.mock_asm_start_response.assert_not_called()
+
+        assert lambda_handler.span.get_tag("http.status_code") == "200"