feat(appsec): enable request blocking

Author: florentinl

datadog_lambda/asm.py

@@ -4,6 +4,7 @@
 
 from ddtrace.contrib.internal.trace_utils import _get_request_header_client_ip
 from ddtrace.internal import core
+from ddtrace.internal.utils import get_blocked, set_blocked
 from ddtrace.trace import Span
 
 from datadog_lambda.trigger import (
@@ -182,3 +183,17 @@ def asm_start_response(
             response_headers,
         ),
     )
+
+
+def get_asm_blocked_response(
+    event_source: _EventSource,
+) -> Optional[Dict[str, Any]]:
+    """Get the blocked response for the given event source."""
+    if event_source.event_type not in _http_event_types:
+        return None
+
+    blocked = get_blocked()
+    if blocked:
+        set_blocked(blocked)
+        return blocked.get("type", "auto")
+    return None

datadog_lambda/wrapper.py

@@ -9,7 +9,12 @@
 from importlib import import_module
 from time import time_ns
 
-from datadog_lambda.asm import asm_set_context, asm_start_response, asm_start_request
+from datadog_lambda.asm import (
+    asm_set_context,
+    asm_start_response,
+    asm_start_request,
+    get_asm_blocked_response,
+)
 from datadog_lambda.dsm import set_dsm_context
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
@@ -159,8 +164,16 @@ def __init__(self, func):
     def __call__(self, event, context, **kwargs):
         """Executes when the wrapped function gets called"""
         self._before(event, context)
+        if config.appsec_enabled:
+            blocking_response = get_asm_blocked_response(self.event_source)
+            if blocking_response:
+                return blocking_response
         try:
             self.response = self.func(event, context, **kwargs)
+            if config.appsec_enabled:
+                blocking_response = get_asm_blocked_response(self.event_source)
+                if blocking_response:
+                    return blocking_response
             return self.response
         except Exception:
             from datadog_lambda.metric import submit_errors_metric