refactor(appsec): remove ASM and IAST import cycle

Author: florentinl

ddtrace/appsec/_asm_request_context.py

@@ -16,6 +16,7 @@
 from ddtrace.appsec._constants import APPSEC
 from ddtrace.appsec._constants import EXPLOIT_PREVENTION
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
+from ddtrace.appsec._iast_request_state import iast_suppress_context
 from ddtrace.appsec._metrics import UNKNOWN_VERSION
 from ddtrace.appsec._metrics import report_waf_run_error
 from ddtrace.appsec._metrics import report_waf_truncation
@@ -788,7 +789,5 @@ def asm_listen() -> None:
 
 def iast_disabled_taint_sources() -> "contextlib.AbstractContextManager[None]":
     if asm_config._iast_enabled:
-        from ddtrace.appsec._iast._iast_request_context_base import iast_suppress_context
-
         return iast_suppress_context()
     return contextlib.nullcontext()

ddtrace/appsec/_exploit_prevention/stack_traces.py

@@ -6,7 +6,6 @@
 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 span_bus
 from ddtrace.internal.settings.asm import config as asm_config
@@ -17,41 +16,39 @@
 ) + (os.path.sep + "wrapt" + os.path.sep,)
 
 
+def _get_stack_span() -> Optional[Span]:
+    span = span_bus.get_span()
+    if span is not None:
+        return span._service_entry_span
+    return span_bus.get_root_span()
+
+
 def report_stack(
     message: Optional[str] = None,
     span: Optional[Span] = None,
     crop_stack: Optional[str] = None,
     stack_id: Optional[str] = None,
     namespace: str = STACK_TRACE.RASP,
 ) -> bool:
-    """
-    Report a stack trace to the current span.
-    This is used to report stack traces for exploit prevention.
-    Return the stack id for the reported stack trace to link it in triggers.
-    """
+    """Report an AppSec stack trace to a span."""
     if not asm_config._ep_stack_trace_enabled:
-        # stack trace report disabled
         return False
     if namespace == STACK_TRACE.RASP and not (asm_config._asm_enabled and asm_config._ep_enabled):
-        # exploit prevention stack trace with ep disabled
         return False
-    if namespace == STACK_TRACE.IAST and not (asm_config._iast_enabled):
-        # iast stack trace with iast disabled
+    if namespace == STACK_TRACE.IAST and not asm_config._iast_enabled:
         return False
 
     if namespace == STACK_TRACE.IAST and asm_config._iast_use_root_span:
         span = span_bus.get_root_span()
-
-    if span is None:
-        span = _asm_request_context.get_entry_span()
+    elif span is None:
+        span = _get_stack_span()
 
     if span is None or stack_id is None:
         return False
     appsec_traces = span._get_struct_tag(STACK_TRACE.TAG) or {}
     current_list = appsec_traces.get(namespace, [])
     total_length = len(current_list)
 
-    # Do not report more than the maximum number of stack traces
     if asm_config._ep_max_stack_traces and total_length >= asm_config._ep_max_stack_traces:
         return False
 
@@ -62,7 +59,6 @@ def report_stack(
             if frame.frame.f_code.co_name == crop_stack:
                 crop_index = i + 1
                 break
-        # Strip any remaining ddtrace/wrapt frames from the top of the stack
         while crop_index < len(stack) and any(d in stack[crop_index].filename for d in _INTERNAL_FRAMES):
             crop_index += 1
     if crop_index:

ddtrace/appsec/_iast/_iast_request_context_base.py

@@ -1,48 +1,25 @@
-import contextlib
-import contextvars
 from typing import Optional
 
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast._iast_env import IASTEnvironment
 from ddtrace.appsec._iast._iast_env import _get_iast_env
 from ddtrace.appsec._iast._overhead_control_engine import oce
-from ddtrace.appsec._iast._taint_tracking._context import debug_num_tainted_objects
+from ddtrace.appsec._iast._request_taint_state import _num_objects_tainted_in_request
 from ddtrace.appsec._iast._taint_tracking._context import finish_request_context
 from ddtrace.appsec._iast._taint_tracking._context import start_request_context
 from ddtrace.appsec._iast.sampling.vulnerability_detection import update_global_vulnerability_limit
+from ddtrace.appsec._iast_request_state import IAST_CONTEXT
+from ddtrace.appsec._iast_request_state import _get_iast_context_id
+from ddtrace.appsec._iast_request_state import iast_suppress_context as iast_suppress_context  # noqa: F401
+from ddtrace.appsec._iast_request_state import is_iast_request_enabled as is_iast_request_enabled
 from ddtrace.internal import core
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.settings.asm import config as asm_config
 
 
 log = get_logger(__name__)
 
-# Stopgap module for providing ASM context for the blocking features wrapping some contextvars.
-
-IAST_CONTEXT: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("iast_var", default=None)
-
-# Keep source suppression separate from IAST_CONTEXT. Clearing the
-# request context id disables request-scoped taint queries and propagation and
-# can send no-context queries through unsafe native fallback paths.
-_IAST_TAINT_SOURCES_SUPPRESSED: contextvars.ContextVar[bool] = contextvars.ContextVar(
-    "iast_taint_sources_suppressed", default=False
-)
-
-
-@contextlib.contextmanager
-def iast_suppress_context():
-    """Temporarily disable IAST taint *source* generation for the current context."""
-    token = _IAST_TAINT_SOURCES_SUPPRESSED.set(True)
-    try:
-        yield
-    finally:
-        _IAST_TAINT_SOURCES_SUPPRESSED.reset(token)
-
-
-def _is_iast_taint_source_enabled() -> bool:
-    return not _IAST_TAINT_SOURCES_SUPPRESSED.get()
-
 
 def _set_span_tag_iast_request_tainted(span):
     total_objects_tainted = _num_objects_tainted_in_request()
@@ -102,11 +79,6 @@ def _iast_start_request(span=None) -> Optional[int]:
     return context_id
 
 
-def _get_iast_context_id() -> Optional[int]:
-    """Retrieve the current IAST context identifier from the ContextVar."""
-    return IAST_CONTEXT.get()
-
-
 def _iast_finish_request(span=None, shoud_update_global_vulnerability_limit: bool = True) -> bool:
     """Finalize the IAST request context and optionally update global limits.
 
@@ -129,22 +101,6 @@ def _iast_finish_request(span=None, shoud_update_global_vulnerability_limit: boo
     return False
 
 
-def is_iast_request_enabled() -> bool:
-    """Check whether IAST is currently operating within an active request context."""
-    return _get_iast_context_id() is not None
-
-
-def _num_objects_tainted_in_request() -> int:
-    """Get the count of tainted objects tracked in the active IAST request context.
-
-    Useful for span metrics and internal telemetry.
-    """
-    context_id = _get_iast_context_id()
-    if context_id is not None:
-        return debug_num_tainted_objects(context_id)
-    return 0
-
-
 def get_hash_object_tracking_len():
     env = _get_iast_env()
     if env:

ddtrace/appsec/_iast/_metrics.py

@@ -5,7 +5,7 @@
 from ddtrace.appsec._constants import TELEMETRY_INFORMATION_VERBOSITY
 from ddtrace.appsec._constants import TELEMETRY_MANDATORY_VERBOSITY
 from ddtrace.appsec._deduplications import deduplication
-from ddtrace.appsec._iast._iast_request_context_base import _num_objects_tainted_in_request
+from ddtrace.appsec._iast._request_taint_state import _num_objects_tainted_in_request
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
 from ddtrace.appsec._iast._utils import _is_iast_debug_enabled

ddtrace/appsec/_iast/_request_taint_state.py

@@ -0,0 +1,11 @@
+from ddtrace.appsec._iast._taint_tracking._context import debug_num_tainted_objects
+from ddtrace.appsec._iast_request_state import _get_iast_context_id
+
+
+def _num_objects_tainted_in_request() -> int:
+    """Return the number of objects tainted in the current request."""
+    context_id = _get_iast_context_id()
+    if context_id is not None:
+        num_tainted: int = debug_num_tainted_objects(context_id)
+        return num_tainted
+    return 0

ddtrace/appsec/_iast/_span_metrics.py

@@ -1,7 +1,7 @@
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast._iast_env import _get_iast_env
-from ddtrace.appsec._iast._iast_request_context_base import _num_objects_tainted_in_request
 from ddtrace.appsec._iast._metrics import _metric_key_as_snake_case
+from ddtrace.appsec._iast._request_taint_state import _num_objects_tainted_in_request
 
 
 def _set_span_tag_iast_request_tainted(span):

ddtrace/appsec/_iast/_taint_tracking/__init__.py

@@ -43,35 +43,24 @@
 from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import shift_taint_ranges  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import str_to_origin  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import taint_range as TaintRange  # noqa: F401
+from ddtrace.appsec._iast_request_state import _get_iast_context_id
 from ddtrace.internal.logger import get_logger
 
 
 log = get_logger(__name__)
 
-# AIDEV-NOTE: _get_iast_context_id is imported lazily — a top-level import here
-# circularly bootstraps via _iast_request_context_base -> _taint_tracking._context
-# -> _taint_tracking/__init__.py. The cached module-global avoids the per-call
-# import dance on this hot path.
-_CACHE_GET_IAST_CONTEXT_ID = None
+# AIDEV-NOTE: Resolve the active native context through the lightweight request-state
+# module. Keeping lifecycle imports out of taint tracking avoids circular initialization
+# and makes this hot path a direct ContextVar lookup.
 
 
 def _current_iast_context_id() -> Optional[int]:
-    global _CACHE_GET_IAST_CONTEXT_ID
-    if _CACHE_GET_IAST_CONTEXT_ID is None:
-        from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id
-
-        _CACHE_GET_IAST_CONTEXT_ID = _get_iast_context_id
-    return _CACHE_GET_IAST_CONTEXT_ID()
+    return _get_iast_context_id()
 
 
 def get_ranges(string_input: Any, context_id: Optional[int] = None) -> Any:
     if context_id is None:
-        global _CACHE_GET_IAST_CONTEXT_ID
-        if _CACHE_GET_IAST_CONTEXT_ID is None:
-            from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id
-
-            _CACHE_GET_IAST_CONTEXT_ID = _get_iast_context_id
-        context_id = _CACHE_GET_IAST_CONTEXT_ID()
+        context_id = _get_iast_context_id()
     if context_id is None:
         return []
     return _native_get_ranges(string_input, context_id)

ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py

@@ -4,15 +4,15 @@
 
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
-from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id
-from ddtrace.appsec._iast._iast_request_context_base import _is_iast_taint_source_enabled
 from ddtrace.appsec._iast._logs import iast_propagation_debug_log
 from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_source
 from ddtrace.appsec._iast._span_metrics import increment_iast_span_metric
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import TaintRange
 from ddtrace.appsec._iast._taint_tracking import set_ranges
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import _taint_pyobject_base
+from ddtrace.appsec._iast_request_state import _get_iast_context_id
+from ddtrace.appsec._iast_request_state import _is_iast_taint_source_enabled
 from ddtrace.internal.logger import get_logger
 
 

ddtrace/appsec/_iast/_taint_tracking/_taint_objects_base.py

@@ -1,14 +1,14 @@
 from typing import Any
 
 from ddtrace.appsec._constants import IAST
-from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id
 from ddtrace.appsec._iast._logs import iast_propagation_debug_log
 from ddtrace.appsec._iast._logs import iast_propagation_error_log
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import get_ranges
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
 from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 from ddtrace.appsec._iast._taint_tracking._context import is_in_taint_map
+from ddtrace.appsec._iast_request_state import _get_iast_context_id
 
 
 def _taint_pyobject_base(pyobject: Any, source_name: Any, source_value: Any, source_origin=None, contextid=None) -> Any:

ddtrace/appsec/_iast/_taint_tracking/aspects.py

@@ -16,7 +16,6 @@
 from typing import Union
 
 from ddtrace.appsec._constants import IAST
-from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled
 from ddtrace.appsec._iast._logs import iast_propagation_error_log
 from ddtrace.appsec._iast._taint_tracking import TagMappingMode
 from ddtrace.appsec._iast._taint_tracking import TaintRange
@@ -50,6 +49,7 @@
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import get_tainted_ranges
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import is_pyobject_tainted
 from ddtrace.appsec._iast._taint_utils import taint_structure
+from ddtrace.appsec._iast_request_state import is_iast_request_enabled
 
 
 TEXT_TYPES = Union[str, bytes, bytearray]