diff --git a/datadog_lambda/asm.py b/datadog_lambda/asm.py index aab0f1e9..9636760c 100644 --- a/datadog_lambda/asm.py +++ b/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], diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index c7474f65..7abe0fc1 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -9,7 +9,7 @@ from importlib import import_module 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.dsm import set_dsm_context from datadog_lambda.extension import should_use_extension, flush_extension from datadog_lambda.cold_start import ( @@ -239,6 +239,10 @@ def _before(self, event, context): ) if config.data_streams_enabled: set_dsm_context(event, event_source) + + if config.appsec_enabled: + asm_set_context(event_source) + self.span = create_function_execution_span( context=context, function_name=config.function_name,