Merge branch 'main' into michael.zhao/dsm-lambda

Author: michael-zhao459

datadog_lambda/asm.py

@@ -44,6 +44,16 @@ def _merge_single_and_multi_value_headers(
     return _to_single_value_headers(merged_headers)
 
 
+def asm_set_context(event_source: _EventSource):
+    """Add asm specific items to the ExecutionContext.
+
+    This allows the AppSecSpanProcessor to know information about the event
+    at the moment the span is created and skip it when not relevant.
+    """
+    if event_source.event_type not in _http_event_types:
+        core.set_item("appsec_skip_next_lambda_event", True)
+
+
 def asm_start_request(
     span: Span,
     event: Dict[str, Any],

datadog_lambda/wrapper.py

@@ -10,6 +10,8 @@
 from time import time_ns
 
 from datadog_lambda.asm import asm_start_response, asm_start_request
+from datadog_lambda.asm import asm_set_context, asm_start_response, asm_start_request
+
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -236,6 +238,10 @@ def _before(self, event, context):
                     self.inferred_span = create_inferred_span(
                         event, context, event_source, config.decode_authorizer_context
                     )
+
+                if config.appsec_enabled:
+                    asm_set_context(event_source)
+
                 self.span = create_function_execution_span(
                     context=context,
                     function_name=config.function_name,