fix(appsec): report all tags on the service entry span instead of the local root span

Author: florentinl

ddtrace/appsec/_api_security/api_manager.py

@@ -6,6 +6,7 @@
 from typing import Optional
 
 from ddtrace._trace._limits import MAX_SPAN_META_VALUE_LEN
+from ddtrace.appsec._asm_request_context import get_service_entry_span
 from ddtrace.appsec._constants import API_SECURITY
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
 from ddtrace.appsec._trace_utils import _asm_manual_keep
@@ -133,7 +134,7 @@ def _should_collect_schema(self, env, priority: int) -> Optional[bool]:
     def _schema_callback(self, env):
         if env.span is None or not asm_config._api_security_feature_active:
             return
-        root = env.span._local_root or env.span
+        root = get_service_entry_span(env.span)
         collected = self.BLOCK_COLLECTED if env.blocked else self.COLLECTED
         if not root or any(meta_name in root._meta for _, meta_name, _ in collected):
             return

ddtrace/appsec/_asm_request_context.py

@@ -69,7 +69,7 @@ def report_error_on_span(error: str, message: str) -> None:
     if not span:
         root_span = core.get_root_span()
     else:
-        root_span = span._local_root or span
+        root_span = get_service_entry_span(span)
     if not root_span:
         return
     root_span.set_tag_str(APPSEC.ERROR_TYPE, error)
@@ -118,6 +118,17 @@ def in_asm_context() -> bool:
     return core.get_item(_ASM_CONTEXT) is not None
 
 
+def get_service_entry_span(span: Span) -> Span:
+    """Find the service entry span of a span
+
+    We can't simply go for the span._local_root
+    because it may belong to another service.
+    """
+    while not span._is_top_level and span._parent:
+        span = span._parent
+    return span
+
+
 def is_blocked() -> bool:
     env = _get_asm_context()
     if env is None:
@@ -194,7 +205,7 @@ def flush_waf_triggers(env: ASM_Environment) -> None:
     from ddtrace.appsec._metrics import ddwaf_version
 
     # Make sure we find a root span to attach the triggers to
-    root_span = env.span._local_root or env.span
+    root_span = get_service_entry_span(env.span)
     if env.waf_triggers:
         report_list = get_triggers(root_span)
         if report_list is not None:
@@ -240,7 +251,7 @@ def finalize_asm_env(env: ASM_Environment) -> None:
     flush_waf_triggers(env)
     for function in env.callbacks[_CONTEXT_CALL]:
         function(env)
-    root_span = env.span._local_root or env.span
+    root_span = get_service_entry_span(env.span)
     if root_span:
         if env.waf_info:
             info = env.waf_info()
@@ -255,7 +266,7 @@ def finalize_asm_env(env: ASM_Environment) -> None:
             except Exception:
                 logger.debug("asm_context::finalize_asm_env::exception", extra=log_extra, exc_info=True)
         if asm_config._rc_client_id is not None:
-            root_span._local_root.set_tag(APPSEC.RC_CLIENT_ID, asm_config._rc_client_id)
+            root_span.set_tag(APPSEC.RC_CLIENT_ID, asm_config._rc_client_id)
         waf_adresses = env.waf_addresses
         req_headers = waf_adresses.get(SPAN_DATA_NAMES.REQUEST_HEADERS_NO_COOKIES, {})
         if req_headers:
@@ -661,7 +672,7 @@ def _set_headers(span: Span, headers: Any, kind: str, only_asm_enabled: bool = F
             value = value.decode()
         if key.lower() in (_COLLECTED_REQUEST_HEADERS_ASM_ENABLED if only_asm_enabled else _COLLECTED_REQUEST_HEADERS):
             # since the header value can be a list, use `set_tag()` to ensure it is converted to a string
-            (span._local_root or span).set_tag(_normalize_tag_name(kind, key), value)
+            get_service_entry_span(span).set_tag(_normalize_tag_name(kind, key), value)
 
 
 def asm_listen():

ddtrace/appsec/_exploit_prevention/stack_traces.py

@@ -6,6 +6,7 @@
 from typing import Optional
 
 from ddtrace._trace.span import Span
+from ddtrace.appsec._asm_request_context import get_service_entry_span
 from ddtrace.appsec._constants import STACK_TRACE
 from ddtrace.internal import core
 from ddtrace.settings.asm import config as asm_config
@@ -38,7 +39,7 @@ def report_stack(
 
     if span is None or stack_id is None:
         return False
-    root_span = span._local_root or span
+    root_span = get_service_entry_span(span)
     appsec_traces = root_span.get_struct_tag(STACK_TRACE.TAG) or {}
     current_list = appsec_traces.get(namespace, [])
     total_length = len(current_list)

ddtrace/appsec/_processor.py

@@ -199,7 +199,7 @@ def on_span_start(self, span: Span) -> None:
         peer_ip = _asm_request_context.get_ip()
         headers = _asm_request_context.get_headers()
         headers_case_sensitive = _asm_request_context.get_headers_case_sensitive()
-        root_span = span._local_root or span
+        root_span = _asm_request_context.get_service_entry_span(span)
 
         root_span.set_metric(APPSEC.ENABLED, 1.0)
         root_span.set_tag_str(_RUNTIME_FAMILY, "python")
@@ -293,7 +293,7 @@ def _waf_action(
             waf_results = Binding_error
 
         _asm_request_context.set_waf_info(lambda: self._ddwaf.info)
-        root_span = span._local_root or span
+        root_span = _asm_request_context.get_service_entry_span(span)
         if waf_results.return_code < 0:
             error_tag = APPSEC.RASP_ERROR if rule_type else APPSEC.WAF_ERROR
             previous = root_span.get_tag(error_tag)

tests/appsec/appsec/api_security/test_api_security_manager.py

@@ -39,7 +39,7 @@ def mock_environment(self):
         root_span = MagicMock(spec=Span)
         root_span._meta = {}
         env.span = MagicMock(spec=Span)
-        env.span._local_root = root_span
+        env.span._parent = root_span
         env.span.context.sampling_priority = None
         root_span.context.sampling_priority = None
         env.waf_addresses = {}