chore(appsec): report everything on the entry_span

Author: florentinl

ddtrace/appsec/_api_security/api_manager.py

@@ -133,7 +133,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 = env.entry_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

@@ -64,16 +64,12 @@ class WARNING_TAGS(metaclass=Constant_Class):
 GLOBAL_CALLBACKS: Dict[str, List[Callable]] = {_CONTEXT_CALL: []}
 
 
-def report_error_on_span(error: str, message: str) -> None:
-    span = getattr(_get_asm_context(), "span", None) or core.get_span()
-    if not span:
-        root_span = core.get_root_span()
-    else:
-        root_span = span._local_root or span
-    if not root_span:
+def report_error_on_entry_span(error: str, message: str) -> None:
+    entry_span = get_entry_span()
+    if not entry_span:
         return
-    root_span.set_tag_str(APPSEC.ERROR_TYPE, error)
-    root_span.set_tag_str(APPSEC.ERROR_MESSAGE, message)
+    entry_span.set_tag_str(APPSEC.ERROR_TYPE, error)
+    entry_span.set_tag_str(APPSEC.ERROR_MESSAGE, message)
 
 
 class ASM_Environment:
@@ -109,6 +105,10 @@ def __init__(self, span: Optional[Span] = None, rc_products: str = ""):
         self.api_security_reported: int = 0
         self.rc_products: str = rc_products
 
+    @property
+    def entry_span(self):
+        return self.span._service_entry_span
+
 
 def _get_asm_context() -> Optional[ASM_Environment]:
     return core.get_item(_ASM_CONTEXT)
@@ -132,6 +132,14 @@ def get_blocked() -> Dict[str, Any]:
     return env.blocked or {}
 
 
+def get_entry_span() -> Optional[Span]:
+    env = _get_asm_context()
+    if env is None:
+        span = core.get_span()
+        return span._service_entry_span if span else None
+    return env.entry_span
+
+
 def get_framework() -> str:
     env = _get_asm_context()
     if env is None:
@@ -193,42 +201,40 @@ def update_span_metrics(span: Span, name: str, value: Union[float, int]) -> None
 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
     if env.waf_triggers:
-        report_list = get_triggers(root_span)
+        report_list = get_triggers(env.entry_span)
         if report_list is not None:
             report_list.extend(env.waf_triggers)
         else:
             report_list = env.waf_triggers
         if asm_config._use_metastruct_for_triggers:
-            root_span.set_struct_tag(APPSEC.STRUCT, {"triggers": report_list})
+            env.entry_span.set_struct_tag(APPSEC.STRUCT, {"triggers": report_list})
         else:
-            root_span.set_tag(APPSEC.JSON, json.dumps({"triggers": report_list}, separators=(",", ":")))
+            env.entry_span.set_tag(APPSEC.JSON, json.dumps({"triggers": report_list}, separators=(",", ":")))
         env.waf_triggers = []
     telemetry_results: Telemetry_result = env.telemetry
 
-    root_span.set_tag_str(APPSEC.WAF_VERSION, ddwaf_version)
+    env.entry_span.set_tag_str(APPSEC.WAF_VERSION, ddwaf_version)
     if telemetry_results.total_duration:
-        update_span_metrics(root_span, APPSEC.WAF_DURATION, telemetry_results.duration)
+        update_span_metrics(env.entry_span, APPSEC.WAF_DURATION, telemetry_results.duration)
         telemetry_results.duration = 0.0
-        update_span_metrics(root_span, APPSEC.WAF_DURATION_EXT, telemetry_results.total_duration)
+        update_span_metrics(env.entry_span, APPSEC.WAF_DURATION_EXT, telemetry_results.total_duration)
         telemetry_results.total_duration = 0.0
     if telemetry_results.timeout:
-        update_span_metrics(root_span, APPSEC.WAF_TIMEOUTS, telemetry_results.timeout)
+        update_span_metrics(env.entry_span, APPSEC.WAF_TIMEOUTS, telemetry_results.timeout)
     rasp_timeouts = sum(telemetry_results.rasp.timeout.values())
     if rasp_timeouts:
-        update_span_metrics(root_span, APPSEC.RASP_TIMEOUTS, rasp_timeouts)
+        update_span_metrics(env.entry_span, APPSEC.RASP_TIMEOUTS, rasp_timeouts)
     if telemetry_results.rasp.sum_eval:
-        update_span_metrics(root_span, APPSEC.RASP_DURATION, telemetry_results.rasp.duration)
-        update_span_metrics(root_span, APPSEC.RASP_DURATION_EXT, telemetry_results.rasp.total_duration)
-        update_span_metrics(root_span, APPSEC.RASP_RULE_EVAL, telemetry_results.rasp.sum_eval)
+        update_span_metrics(env.entry_span, APPSEC.RASP_DURATION, telemetry_results.rasp.duration)
+        update_span_metrics(env.entry_span, APPSEC.RASP_DURATION_EXT, telemetry_results.rasp.total_duration)
+        update_span_metrics(env.entry_span, APPSEC.RASP_RULE_EVAL, telemetry_results.rasp.sum_eval)
     if telemetry_results.truncation.string_length:
-        root_span.set_metric(APPSEC.TRUNCATION_STRING_LENGTH, max(telemetry_results.truncation.string_length))
+        env.entry_span.set_metric(APPSEC.TRUNCATION_STRING_LENGTH, max(telemetry_results.truncation.string_length))
     if telemetry_results.truncation.container_size:
-        root_span.set_metric(APPSEC.TRUNCATION_CONTAINER_SIZE, max(telemetry_results.truncation.container_size))
+        env.entry_span.set_metric(APPSEC.TRUNCATION_CONTAINER_SIZE, max(telemetry_results.truncation.container_size))
     if telemetry_results.truncation.container_depth:
-        root_span.set_metric(APPSEC.TRUNCATION_CONTAINER_DEPTH, max(telemetry_results.truncation.container_depth))
+        env.entry_span.set_metric(APPSEC.TRUNCATION_CONTAINER_DEPTH, max(telemetry_results.truncation.container_depth))
 
 
 def finalize_asm_env(env: ASM_Environment) -> None:
@@ -240,31 +246,30 @@ 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
-    if root_span:
+    if env.entry_span:
         if env.waf_info:
             info = env.waf_info()
             try:
                 if info.errors:
-                    root_span.set_tag_str(APPSEC.EVENT_RULE_ERRORS, info.errors)
+                    env.entry_span.set_tag_str(APPSEC.EVENT_RULE_ERRORS, info.errors)
                     extra = {"product": "appsec", "more_info": info.errors, "stack_limit": 4}
                     logger.debug("asm_context::finalize_asm_env::waf_errors", extra=extra, stack_info=True)
-                root_span.set_tag_str(APPSEC.EVENT_RULE_VERSION, info.version)
-                root_span.set_metric(APPSEC.EVENT_RULE_LOADED, info.loaded)
-                root_span.set_metric(APPSEC.EVENT_RULE_ERROR_COUNT, info.failed)
+                env.entry_span.set_tag_str(APPSEC.EVENT_RULE_VERSION, info.version)
+                env.entry_span.set_metric(APPSEC.EVENT_RULE_LOADED, info.loaded)
+                env.entry_span.set_metric(APPSEC.EVENT_RULE_ERROR_COUNT, info.failed)
             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)
+            env.entry_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:
-            _set_headers(root_span, req_headers, kind="request")
+            _set_headers(env.entry_span, req_headers, kind="request")
         res_headers = waf_adresses.get(SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES, {})
         if res_headers:
-            _set_headers(root_span, res_headers, kind="response")
+            _set_headers(env.entry_span, res_headers, kind="response")
         if env.rc_products:
-            root_span.set_tag_str(APPSEC.RC_PRODUCTS, env.rc_products)
+            env.entry_span.set_tag_str(APPSEC.RC_PRODUCTS, env.rc_products)
 
     core.discard_local_item(_ASM_CONTEXT)
 
@@ -364,7 +369,7 @@ def call_waf_callback(custom_data: Optional[Dict[str, Any]] = None, **kwargs) ->
         return callback(custom_data, **kwargs)
     else:
         logger.warning(WARNING_TAGS.CALL_WAF_CALLBACK_NOT_SET, extra=log_extra, stack_info=True)
-        report_error_on_span("appsec::instrumentation::diagnostic", WARNING_TAGS.CALL_WAF_CALLBACK_NOT_SET)
+        report_error_on_entry_span("appsec::instrumentation::diagnostic", WARNING_TAGS.CALL_WAF_CALLBACK_NOT_SET)
         return None
 
 
@@ -661,7 +666,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)
+            span.set_tag(_normalize_tag_name(kind, key), value)
 
 
 def asm_listen():

ddtrace/appsec/_exploit_prevention/stack_traces.py

@@ -6,8 +6,8 @@
 from typing import Optional
 
 from ddtrace._trace.span import Span
+from ddtrace.appsec import _asm_request_context
 from ddtrace.appsec._constants import STACK_TRACE
-from ddtrace.internal import core
 from ddtrace.settings.asm import config as asm_config
 
 
@@ -34,12 +34,11 @@ def report_stack(
         return False
 
     if span is None:
-        span = core.get_root_span()
+        span = _asm_request_context.get_entry_span()
 
     if span is None or stack_id is None:
         return False
-    root_span = span._local_root or span
-    appsec_traces = root_span.get_struct_tag(STACK_TRACE.TAG) or {}
+    appsec_traces = span.get_struct_tag(STACK_TRACE.TAG) or {}
     current_list = appsec_traces.get(namespace, [])
     total_length = len(current_list)
 
@@ -77,5 +76,5 @@ def report_stack(
     res["frames"] = frames
     current_list.append(res)
     appsec_traces[namespace] = current_list
-    root_span.set_struct_tag(STACK_TRACE.TAG, appsec_traces)
+    span.set_struct_tag(STACK_TRACE.TAG, appsec_traces)
     return True

ddtrace/appsec/_processor.py

@@ -198,13 +198,15 @@ 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
+        entry_span = _asm_request_context.get_entry_span()
+        if not entry_span:
+            return
 
-        root_span.set_metric(APPSEC.ENABLED, 1.0)
-        root_span.set_tag_str(_RUNTIME_FAMILY, "python")
+        entry_span.set_metric(APPSEC.ENABLED, 1.0)
+        entry_span.set_tag_str(_RUNTIME_FAMILY, "python")
 
         def waf_callable(custom_data=None, **kwargs):
-            return self._waf_action(root_span, ctx, custom_data, **kwargs)
+            return self._waf_action(entry_span, ctx, custom_data, **kwargs)
 
         _asm_request_context.set_waf_callback(waf_callable)
         _asm_request_context.add_context_callback(self.metrics._set_waf_request_metrics)
@@ -225,7 +227,7 @@ def waf_callable(custom_data=None, **kwargs):
 
     def _waf_action(
         self,
-        span: Span,
+        entry_span: Span,
         ctx: "ddwaf.ddwaf_types.ddwaf_context_capsule",
         custom_data: Optional[Dict[str, Any]] = None,
         crop_trace: Optional[str] = None,
@@ -292,18 +294,17 @@ def _waf_action(
             waf_results = Binding_error
 
         _asm_request_context.set_waf_info(lambda: self._ddwaf.info)
-        root_span = span._local_root or 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)
+            previous = entry_span.get_tag(error_tag)
             if previous is None:
-                root_span.set_tag_str(error_tag, str(waf_results.return_code))
+                entry_span.set_tag_str(error_tag, str(waf_results.return_code))
             else:
                 try:
                     int_previous = int(previous)
                 except ValueError:
                     int_previous = -128
-                root_span.set_tag_str(error_tag, str(max(int_previous, waf_results.return_code)))
+                entry_span.set_tag_str(error_tag, str(max(int_previous, waf_results.return_code)))
 
         blocked = {}
         for action, parameters in waf_results.actions.items():
@@ -314,15 +315,17 @@ def _waf_action(
                 blocked[WAF_ACTIONS.TYPE] = "none"
             elif action == WAF_ACTIONS.STACK_ACTION:
                 stack_trace_id = parameters["stack_id"]
-                report_stack("exploit detected", span, crop_trace, stack_id=stack_trace_id, namespace=STACK_TRACE.RASP)
+                report_stack(
+                    "exploit detected", entry_span, crop_trace, stack_id=stack_trace_id, namespace=STACK_TRACE.RASP
+                )
                 for rule in waf_results.data:
                     rule[EXPLOIT_PREVENTION.STACK_TRACE_ID] = stack_trace_id
 
         # Trace tagging
         for key, value in waf_results.meta_tags.items():
-            root_span.set_tag_str(key, value)
+            entry_span.set_tag_str(key, value)
         for key, value in waf_results.metrics.items():
-            root_span.set_metric(key, value)
+            entry_span.set_metric(key, value)
 
         if waf_results.data:
             log.debug("[DDAS-011-00] ASM In-App WAF returned: %s. Timeout %s", waf_results.data, waf_results.timeout)
@@ -345,24 +348,24 @@ def _waf_action(
         if waf_results.data:
             _asm_request_context.store_waf_results_data(waf_results.data)
             if blocked:
-                span.set_tag(APPSEC.BLOCKED, "true")
+                entry_span.set_tag(APPSEC.BLOCKED, "true")
 
             # Partial DDAS-011-00
-            span.set_tag_str(APPSEC.EVENT, "true")
+            entry_span.set_tag_str(APPSEC.EVENT, "true")
 
             remote_ip = _asm_request_context.get_waf_address(SPAN_DATA_NAMES.REQUEST_HTTP_IP)
             if remote_ip:
                 # Note that if the ip collection is disabled by the env var
                 # DD_TRACE_CLIENT_IP_HEADER_DISABLED actor.ip won't be sent
-                span.set_tag_str("actor.ip", remote_ip)
+                entry_span.set_tag_str("actor.ip", remote_ip)
 
             # Right now, we overwrite any value that could be already there. We need to reconsider when ASM/AppSec's
             # specs are updated.
-            if span.get_tag(_ORIGIN_KEY) is None:
-                span.set_tag_str(_ORIGIN_KEY, APPSEC.ORIGIN_VALUE)
+            if entry_span.get_tag(_ORIGIN_KEY) is None:
+                entry_span.set_tag_str(_ORIGIN_KEY, APPSEC.ORIGIN_VALUE)
 
         if waf_results.keep and allowed:
-            _asm_manual_keep(span)
+            _asm_manual_keep(entry_span)
 
         return waf_results