chore(appsec): report everything on the entry_span

Author: florentinl

.github/codeql-config.yml

@@ -1,3 +1,4 @@
 name: "CodeQL config"
 paths-ignore:
-  - 'tests/appsec/iast_packages/packages/**'
+  - "tests/appsec/iast_packages/packages/**"
+  - "tests/appsec/contrib_appsec/**"

ddtrace/appsec/_api_security/api_manager.py

@@ -131,7 +131,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:
@@ -93,6 +89,7 @@ def __init__(self, span: Optional[Span] = None, rc_products: str = ""):
             logger.warning(WARNING_TAGS.ASM_ENV_NO_SPAN, extra=log_extra, stack_info=True)
             raise TypeError("ASM_Environment requires a span")
         self.span: Span = context_span
+        self.entry_span: Span = self.span._service_entry_span
         if self.span.name.endswith(".request"):
             self.framework = self.span.name[:-8]
         else:
@@ -132,6 +129,17 @@ 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()
+        if span:
+            return span._service_entry_span
+        else:
+            return core.get_root_span()
+    return env.entry_span
+
+
 def get_framework() -> str:
     env = _get_asm_context()
     if env is None:
@@ -193,42 +201,41 @@ 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
+    entry_span = env.entry_span
     if env.waf_triggers:
-        report_list = get_triggers(root_span)
+        report_list = get_triggers(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})
+            entry_span.set_struct_tag(APPSEC.STRUCT, {"triggers": report_list})
         else:
-            root_span.set_tag(APPSEC.JSON, json.dumps({"triggers": report_list}, separators=(",", ":")))
+            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)
+    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(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(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(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(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(entry_span, APPSEC.RASP_DURATION, telemetry_results.rasp.duration)
+        update_span_metrics(entry_span, APPSEC.RASP_DURATION_EXT, telemetry_results.rasp.total_duration)
+        update_span_metrics(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))
+        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))
+        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))
+        entry_span.set_metric(APPSEC.TRUNCATION_CONTAINER_DEPTH, max(telemetry_results.truncation.container_depth))
 
 
 def finalize_asm_env(env: ASM_Environment) -> None:
@@ -240,31 +247,31 @@ 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:
+    entry_span = env.entry_span
+    if 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)
+                    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)
+                entry_span.set_tag_str(APPSEC.EVENT_RULE_VERSION, info.version)
+                entry_span.set_metric(APPSEC.EVENT_RULE_LOADED, info.loaded)
+                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)
+            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(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(entry_span, res_headers, kind="response")
         if env.rc_products:
-            root_span.set_tag_str(APPSEC.RC_PRODUCTS, env.rc_products)
+            entry_span.set_tag_str(APPSEC.RC_PRODUCTS, env.rc_products)
 
     core.discard_local_item(_ASM_CONTEXT)
 
@@ -364,7 +371,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 +668,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,6 +6,7 @@
 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
@@ -33,13 +34,15 @@ def report_stack(
         # iast stack trace with iast disabled
         return False
 
-    if span is None:
+    if namespace == STACK_TRACE.IAST and asm_config._iast_use_root_span:
         span = core.get_root_span()
 
+    if span is None:
+        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 +80,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/_iast/__init__.py

@@ -27,6 +27,7 @@ def wrapped_function(wrapped, instance, args, kwargs):
     )
     return wrapped(*args, **kwargs)
 """
+
 import os
 import sys
 import types
@@ -107,7 +108,6 @@ def _iast_pytest_activation():
     if _iast_propagation_enabled:
         return
     os.environ["DD_IAST_ENABLED"] = os.environ.get("DD_IAST_ENABLED") or "1"
-    os.environ["_DD_IAST_USE_ROOT_SPAN"] = os.environ.get("_DD_IAST_USE_ROOT_SPAN") or "true"
     os.environ["DD_IAST_REQUEST_SAMPLING"] = os.environ.get("DD_IAST_REQUEST_SAMPLING") or "100.0"
     os.environ["_DD_APPSEC_DEDUPLICATION_ENABLED"] = os.environ.get("_DD_APPSEC_DEDUPLICATION_ENABLED") or "false"
     os.environ["DD_IAST_VULNERABILITIES_PER_REQUEST"] = os.environ.get("DD_IAST_VULNERABILITIES_PER_REQUEST") or "1000"

ddtrace/appsec/_iast/_iast_request_context.py

@@ -77,7 +77,7 @@ def _create_and_attach_iast_report_to_span(
 
 def _iast_end_request(ctx=None, span=None, *args, **kwargs):
     try:
-        move_to_root = base._move_iast_data_to_root_span()
+        move_to_root = asm_config._iast_use_root_span
         if move_to_root:
             req_span = core.get_root_span()
         else:

ddtrace/appsec/_iast/_iast_request_context_base.py

@@ -1,4 +1,3 @@
-import os
 from typing import Optional
 
 from ddtrace._trace.span import Span
@@ -13,7 +12,6 @@
 from ddtrace.appsec._iast.sampling.vulnerability_detection import update_global_vulnerability_limit
 from ddtrace.internal import core
 from ddtrace.internal.logger import get_logger
-from ddtrace.internal.utils.formats import asbool
 from ddtrace.settings.asm import config as asm_config
 
 
@@ -80,10 +78,6 @@ def set_iast_request_endpoint(method, route) -> None:
             log.debug("iast::propagation::context::Trying to set IAST request endpoint but no context is present")
 
 
-def _move_iast_data_to_root_span():
-    return asbool(os.getenv("_DD_IAST_USE_ROOT_SPAN"))
-
-
 def _iast_start_request(span=None, *args, **kwargs):
     try:
         if asm_config._iast_enabled:

ddtrace/appsec/_processor.py

@@ -211,13 +211,12 @@ 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.set_metric(APPSEC.ENABLED, 1.0)
-        root_span.set_tag_str(_RUNTIME_FAMILY, "python")
+        entry_span = span._service_entry_span
+        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)
@@ -238,7 +237,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,
@@ -304,18 +303,17 @@ def _waf_action(
             log.debug("appsec::processor::waf::run", exc_info=True)
             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():
@@ -326,15 +324,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)
@@ -357,24 +357,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