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

.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/_trace/span.py

@@ -128,6 +128,7 @@ class Span(object):
         "_context",
         "_parent_context",
         "_local_root_value",
+        "_service_entry_span_value",
         "_parent",
         "_ignored_exceptions",
         "_on_finish_callbacks",
@@ -221,6 +222,7 @@ def __init__(
         self._parent: Optional["Span"] = None
         self._ignored_exceptions: Optional[List[Type[Exception]]] = None
         self._local_root_value: Optional["Span"] = None  # None means this is the root span.
+        self._service_entry_span_value: Optional["Span"] = None  # None means this is the service entry span.
         self._store: Optional[Dict[str, Any]] = None
 
     def _update_tags_from_context(self) -> None:
@@ -710,21 +712,28 @@ def context(self) -> Context:
 
     @property
     def _local_root(self) -> "Span":
-        if self._local_root_value is None:
-            return self
-        return self._local_root_value
+        return self._local_root_value or self
 
     @_local_root.setter
     def _local_root(self, value: "Span") -> None:
-        if value is not self:
-            self._local_root_value = value
-        else:
-            self._local_root_value = None
+        self._local_root_value = value if value is not self else None
 
     @_local_root.deleter
     def _local_root(self) -> None:
         del self._local_root_value
 
+    @property
+    def _service_entry_span(self) -> "Span":
+        return self._service_entry_span_value or self
+
+    @_service_entry_span.setter
+    def _service_entry_span(self, span: "Span") -> None:
+        self._service_entry_span_value = None if span is self else span
+
+    @_service_entry_span.deleter
+    def _service_entry_span(self) -> None:
+        del self._service_entry_span_value
+
     def link_span(self, context: Context, attributes: Optional[Dict[str, Any]] = None) -> None:
         """Defines a causal relationship between two spans"""
         if not context.trace_id or not context.span_id:
@@ -859,7 +868,8 @@ def __repr__(self) -> str:
             f"metrics={self._metrics}, "
             f"links={self._links}, "
             f"events={self._events}, "
-            f"context={self._context})"
+            f"context={self._context}, "
+            f"service_entry_span_name={self._service_entry_span.name})"
         )
 
     def __str__(self) -> str:

ddtrace/_trace/tracer.py

@@ -530,6 +530,8 @@ def _start_span(
             if parent:
                 span._parent = parent
                 span._local_root = parent._local_root
+                if span._parent.service == service:
+                    span._service_entry_span = parent._service_entry_span
 
             for k, v in _get_metas_to_propagate(context):
                 # We do not want to propagate AppSec propagation headers

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: