python_lambda: test lambda inferred span compliance with RFC-1081

.github/workflows/run-end-to-end.yml

@@ -457,6 +457,9 @@ jobs:
     - name: Run APPSEC_LAMBDA_RASP scenario
       if: always() && steps.build.outcome == 'success' && contains(inputs.scenarios, '"APPSEC_LAMBDA_RASP"')
       run: ./run.sh APPSEC_LAMBDA_RASP
+    - name: Run APPSEC_LAMBDA_INFERRED_SPANS scenario
+      if: always() && steps.build.outcome == 'success' && contains(inputs.scenarios, '"APPSEC_LAMBDA_INFERRED_SPANS"')
+      run: ./run.sh APPSEC_LAMBDA_INFERRED_SPANS
     - name: Run EXTERNAL_PROCESSING scenario
       if: always() && steps.build.outcome == 'success' && contains(inputs.scenarios, '"EXTERNAL_PROCESSING"')
       run: ./run.sh EXTERNAL_PROCESSING

manifests/python_lambda.yml

@@ -209,6 +209,7 @@ manifest:
     - weblog_declaration:
         "*": v8.113.0
         alb-multi: v8.114.0.dev
+  tests/appsec/test_inferred_spans.py::Test_Inferred_Span_Tags: missing_feature
   tests/appsec/test_only_python.py::Test_ImportError: v7.112.0
   tests/appsec/test_reports.py::Test_ExtraTagsFromRule: v7.112.0
   tests/appsec/test_reports.py::Test_Info: v7.112.0

tests/appsec/test_inferred_spans.py

@@ -0,0 +1,48 @@
+import json
+from typing import Any
+
+from utils import weblog, interfaces, scenarios, features
+
+
+INFERRED_SPAN_NAMES = {"aws.apigateway", "aws.httpapi"}
+
+
+@scenarios.appsec_lambda_inferred_spans
+@features.appsec_api_gateway_inferred_span_discovery
+class Test_Inferred_Span_Tags:
+    """Tests for endpoint discovery & correlation from lambda inferred spans"""
+
+    def setup_lambda_inferred_span(self) -> None:
+        self.r = weblog.get("/waf/?message=<script>alert()</script>")
+
+    def test_lambda_inferred_span(self) -> None:
+        for _, _, span, appsec_data in interfaces.library.get_appsec_events(self.r):
+            if span.get("name") == "aws.lambda":
+                lambda_span_appsec_data = appsec_data
+
+        assert lambda_span_appsec_data, "Expected non empty appsec data on aws.lambda span"
+
+        def validate_inferred_span(span: dict[str, Any]) -> bool:
+            if span.get("name") not in INFERRED_SPAN_NAMES:
+                return False
+
+            assert span.get("type") == "web", "Lambda inferred spans must be of type web"
+            assert "operation_name" not in span.get("meta", {}), "operation_name should be removed"
+
+            metrics = span.get("metrics", {})
+            appsec_enabled = metrics.get("_dd.appsec.enabled")
+            assert appsec_enabled is not None, "Lambda inferred spans must report _dd.appsec.enabled"
+            assert float(appsec_enabled) == 1.0
+
+            inferred_span_payload = span.get("meta", {}).get("_dd.appsec.json", {}) or span.get("meta_struct", {}).get(
+                "appsec", {}
+            )
+            assert inferred_span_payload, "Lambda inferred spans must include the appsec payload"
+            inferred_payload = (
+                json.loads(inferred_span_payload) if isinstance(inferred_span_payload, str) else inferred_span_payload
+            )
+            assert inferred_payload == lambda_span_appsec_data, "AppSec Data must match the service-entry span"
+
+            return True
+
+        interfaces.library.validate_one_span(self.r, validator=validate_inferred_span, full_trace=True)

utils/_context/_scenarios/__init__.py

@@ -1146,6 +1146,12 @@ class _Scenarios:
         scenario_groups=[scenario_groups.appsec, scenario_groups.appsec_lambda],
     )
     appsec_lambda_rasp = AppSecLambdaRaspScenario("APPSEC_LAMBDA_RASP")
+    appsec_lambda_inferred_spans = LambdaScenario(
+        "APPSEC_LAMBDA_INFERRED_SPANS",
+        doc="Lambda scenario with managed services tracing enabled",
+        scenario_groups=[scenario_groups.appsec, scenario_groups.appsec_lambda],
+        trace_managed_services=True,
+    )
 
     otel_collector = OtelCollectorScenario("OTEL_COLLECTOR")
     otel_collector_e2e = OtelCollectorScenario("OTEL_COLLECTOR_E2E", mocked_backend=False)

utils/_context/_scenarios/aws_lambda.py

@@ -27,6 +27,7 @@ def __init__(
         scenario_groups: list[ScenarioGroup] | None = None,
         weblog_env: dict[str, str | None] | None = None,
         weblog_volumes: dict[str, dict[str, str]] | None = None,
+        trace_managed_services: bool = False,
     ):
         scenario_groups = [
             all_scenario_groups.tracer_release,
@@ -37,8 +38,7 @@ def __init__(
         super().__init__(name, github_workflow=github_workflow, doc=doc, scenario_groups=scenario_groups)
 
         self.lambda_weblog = LambdaWeblogContainer(
-            environment=weblog_env or {},
-            volumes=weblog_volumes or {},
+            environment=weblog_env or {}, volumes=weblog_volumes or {}, trace_managed_services=trace_managed_services
         )
 
         self.lambda_proxy_container = LambdaProxyContainer(

utils/_context/containers.py

@@ -1077,13 +1077,14 @@ def __init__(
         *,
         environment: dict[str, str | None] | None = None,
         volumes: dict | None = None,
+        trace_managed_services: bool = False,
     ):
         environment = (environment or {}) | {
             "DD_HOSTNAME": "test",
             "DD_SITE": os.environ.get("DD_SITE", "datad0g.com"),
             "DD_API_KEY": os.environ.get("DD_API_KEY", _FAKE_DD_API_KEY),
             "DD_SERVERLESS_FLUSH_STRATEGY": "periodically,100",
-            "DD_TRACE_MANAGED_SERVICES": "false",
+            "DD_TRACE_MANAGED_SERVICES": str(trace_managed_services).lower(),
         }
 
         volumes = volumes or {}

utils/_features.py

@@ -2687,5 +2687,13 @@ def llm_observability_anthropic_messages(test_object):
         """
         return _mark_test_object(test_object, feature_id=524, owner=_Owner.ml_observability)
 
+    @staticmethod
+    def appsec_api_gateway_inferred_span_discovery(test_object):
+        """Support API Gateway Inferred span discovery and correlation in the App & API Protection  API Catalog
+
+        https://feature-parity.us1.prod.dog/#/?feature=526
+        """
+        return _mark_test_object(test_object, feature_id=526, owner=_Owner.asm)
+
 
 features = _Features()

utils/build/docker/lambda_proxy/main.py

@@ -39,6 +39,8 @@ def invoke_lambda_function_api_gateway_rest():
         stage_name="Prod",
     )
 
+    converted_event["requestContext"]["identity"]["userAgent"] = request.user_agent.string
+
     response = post(
         RIE_URL,
         json=converted_event,
@@ -64,6 +66,7 @@ def invoke_lambda_function_api_gateway_http():
     path = PathConverter.convert_path_to_api_gateway(request.path)
     route_key = LocalApigwService._v2_route_key(request.method, path, is_default_route=False)
     converted_event = construct_v2_event_http(request, PORT, binary_types=BINARY_TYPES, route_key=route_key)
+    converted_event["requestContext"]["http"]["userAgent"] = request.user_agent.string
 
     response = post(
         RIE_URL,

utils/scripts/ci_orchestrators/workflow_data.py

@@ -488,6 +488,7 @@ def _is_supported(library: str, weblog: str, scenario: str, _ci_environment: str
         "APPSEC_LAMBDA_BLOCKING",
         "APPSEC_LAMBDA_API_SECURITY",
         "APPSEC_LAMBDA_RASP",
+        "APPSEC_LAMBDA_INFERRED_SPANS",
     )
     if is_lambda_library != is_lambda_scenario:
         return False