chore(iast): context refactor III

.gitlab/benchmarks/bp-runner.microbenchmarks.fail-on-breach.yml

@@ -722,15 +722,15 @@ experiments:
               - max_rss_usage < 39.00 MB
           - name: iastpropagation-propagation_enabled
             thresholds:
-              - execution_time < 0.16 ms
+              - execution_time < 0.19 ms
               - max_rss_usage < 39.00 MB
           - name: iastpropagation-propagation_enabled_100
             thresholds:
-              - execution_time < 1.80 ms
+              - execution_time < 1.90 ms
               - max_rss_usage < 39.00 MB
           - name: iastpropagation-propagation_enabled_1000
             thresholds:
-              - execution_time < 30.55 ms
+              - execution_time < 35.55 ms
               - max_rss_usage < 39.00 MB
 
           # otelsdkspan

benchmarks/appsec_iast_aspects/scenario.py

@@ -5,21 +5,29 @@
 
 
 try:
-    # 3.15+
+    # >= 3.15
     from ddtrace.appsec._iast._iast_request_context_base import _iast_finish_request as end_iast_context
     from ddtrace.appsec._iast._iast_request_context_base import _iast_start_request as iast_start_request
-    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
 except ImportError:
     try:
-        # 3.6+
+        # >= 3.6; < 3.15
         from ddtrace.appsec._iast._iast_request_context_base import end_iast_context
-        from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
         from ddtrace.appsec._iast._iast_request_context_base import start_iast_context as iast_start_request
     except ImportError:
-        # Pre 3.6
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import end_iast_context
         from ddtrace.appsec._iast._iast_request_context import iast_start_request
+
+try:
+    # >= 3.6; < 3.15
+    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
+except ImportError:
+    try:
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import set_iast_request_enabled
+    except ImportError:
+        # >= 3.15
+        set_iast_request_enabled = lambda x: None  # noqa: E731
 
 
 def _start_iast_context_and_oce():

benchmarks/appsec_iast_aspects_ospath/scenario.py

@@ -3,21 +3,29 @@
 
 
 try:
-    # 3.15+
+    # >= 3.15
     from ddtrace.appsec._iast._iast_request_context_base import _iast_finish_request as end_iast_context
     from ddtrace.appsec._iast._iast_request_context_base import _iast_start_request as iast_start_request
-    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
 except ImportError:
     try:
-        # 3.6+
+        # >= 3.6; < 3.15
         from ddtrace.appsec._iast._iast_request_context_base import end_iast_context
-        from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
         from ddtrace.appsec._iast._iast_request_context_base import start_iast_context as iast_start_request
     except ImportError:
-        # Pre 3.6
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import end_iast_context
         from ddtrace.appsec._iast._iast_request_context import iast_start_request
+
+try:
+    # >= 3.6; < 3.15
+    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
+except ImportError:
+    try:
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import set_iast_request_enabled
+    except ImportError:
+        # >= 3.15
+        set_iast_request_enabled = lambda x: None  # noqa: E731
 
 
 def _start_iast_context_and_oce():

benchmarks/appsec_iast_aspects_re_module/scenario.py

@@ -3,21 +3,29 @@
 
 
 try:
-    # 3.15+
+    # >= 3.15
     from ddtrace.appsec._iast._iast_request_context_base import _iast_finish_request as end_iast_context
     from ddtrace.appsec._iast._iast_request_context_base import _iast_start_request as iast_start_request
-    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
 except ImportError:
     try:
-        # 3.6+
+        # >= 3.6; < 3.15
         from ddtrace.appsec._iast._iast_request_context_base import end_iast_context
-        from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
         from ddtrace.appsec._iast._iast_request_context_base import start_iast_context as iast_start_request
     except ImportError:
-        # Pre 3.6
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import end_iast_context
         from ddtrace.appsec._iast._iast_request_context import iast_start_request
+
+try:
+    # >= 3.6; < 3.15
+    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
+except ImportError:
+    try:
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import set_iast_request_enabled
+    except ImportError:
+        # >= 3.15
+        set_iast_request_enabled = lambda x: None  # noqa: E731
 
 
 def _start_iast_context_and_oce():

benchmarks/appsec_iast_aspects_split/scenario.py

@@ -3,21 +3,29 @@
 
 
 try:
-    # 3.15+
+    # >= 3.15
     from ddtrace.appsec._iast._iast_request_context_base import _iast_finish_request as end_iast_context
     from ddtrace.appsec._iast._iast_request_context_base import _iast_start_request as iast_start_request
-    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
 except ImportError:
     try:
-        # 3.6+
+        # >= 3.6; < 3.15
         from ddtrace.appsec._iast._iast_request_context_base import end_iast_context
-        from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
         from ddtrace.appsec._iast._iast_request_context_base import start_iast_context as iast_start_request
     except ImportError:
-        # Pre 3.6
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import end_iast_context
         from ddtrace.appsec._iast._iast_request_context import iast_start_request
+
+try:
+    # >= 3.6; < 3.15
+    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
+except ImportError:
+    try:
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import set_iast_request_enabled
+    except ImportError:
+        # >= 3.15
+        set_iast_request_enabled = lambda x: None  # noqa: E731
 
 
 def _start_iast_context_and_oce():

benchmarks/appsec_iast_propagation/scenario.py

@@ -2,21 +2,29 @@
 
 
 try:
-    # 3.15+
+    # >= 3.15
     from ddtrace.appsec._iast._iast_request_context_base import _iast_finish_request as end_iast_context
     from ddtrace.appsec._iast._iast_request_context_base import _iast_start_request as iast_start_request
-    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
 except ImportError:
     try:
-        # 3.6+
+        # >= 3.6; < 3.15
         from ddtrace.appsec._iast._iast_request_context_base import end_iast_context
-        from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
         from ddtrace.appsec._iast._iast_request_context_base import start_iast_context as iast_start_request
     except ImportError:
-        # Pre 3.6
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import end_iast_context
         from ddtrace.appsec._iast._iast_request_context import iast_start_request
+
+try:
+    # >= 3.6; < 3.15
+    from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_enabled
+except ImportError:
+    try:
+        # < 3.6
         from ddtrace.appsec._iast._iast_request_context import set_iast_request_enabled
+    except ImportError:
+        # >= 3.15
+        set_iast_request_enabled = lambda x: None  # noqa: E731
 
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject

ddtrace/appsec/_iast/_iast_env.py

@@ -21,7 +21,6 @@ class IASTEnvironment:
     def __init__(self, span: Optional[Span] = None):
         self.span = span or core.get_span()
 
-        self.request_enabled: bool = False
         self.iast_reporter: Optional["IastSpanReporter"] = None
         self.iast_span_metrics: Dict[str, int] = {}
         self.iast_stack_trace_reported: bool = False

ddtrace/appsec/_iast/_iast_request_context.py

@@ -9,7 +9,6 @@
 from ddtrace.appsec._iast._iast_env import _get_iast_env
 import ddtrace.appsec._iast._iast_request_context_base as base
 from ddtrace.appsec._iast._metrics import _set_metric_iast_request_tainted
-from ddtrace.appsec._iast._overhead_control_engine import oce
 from ddtrace.appsec._iast._span_metrics import _set_span_tag_iast_executed_sink
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
@@ -62,14 +61,11 @@ def _create_and_attach_iast_report_to_span(
     base._set_span_tag_iast_request_tainted(req_span)
     _set_span_tag_iast_executed_sink(req_span)
 
-    base.set_iast_request_enabled(False)
     base._iast_finish_request(req_span)
 
     if req_span.get_tag(_ORIGIN_KEY) is None:
         req_span.set_tag_str(_ORIGIN_KEY, APPSEC.ORIGIN_VALUE)
 
-    oce.release_request()
-
 
 def _iast_end_request(ctx=None, span=None, *args, **kwargs):
     try:
@@ -84,14 +80,14 @@ def _iast_end_request(ctx=None, span=None, *args, **kwargs):
         if req_span is None:
             log.debug("iast::propagation::context::Error finishing IAST context. There isn't a SPAN")
             return
+
         if asm_config._iast_enabled:
             existing_data = req_span.get_tag(IAST.JSON) or req_span.get_struct_tag(IAST.STRUCT)
             if existing_data is None:
                 if req_span.get_metric(IAST.ENABLED) is None:
                     if not base.is_iast_request_enabled():
                         req_span.set_metric(IAST.ENABLED, 0.0)
                         base._iast_finish_request(req_span)
-                        oce.release_request()
                         return
 
                     req_span.set_metric(IAST.ENABLED, 1.0)

ddtrace/appsec/_iast/_iast_request_context_base.py

@@ -1,15 +1,14 @@
 import contextvars
 from typing import Optional
 
-from ddtrace._trace.span import Span
 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 import num_objects_tainted
-from ddtrace.appsec._iast._taint_tracking._context import create_context as create_propagation_context
-from ddtrace.appsec._iast._taint_tracking._context import reset_context as reset_propagation_context
+from ddtrace.appsec._iast._taint_tracking._context import debug_num_tainted_objects
+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.internal import core
 from ddtrace.internal.logger import get_logger
@@ -20,7 +19,7 @@
 
 # Stopgap module for providing ASM context for the blocking features wrapping some contextvars.
 
-IAST_CONTEXT = contextvars.ContextVar("iast_var", default=None)
+IAST_CONTEXT: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("iast_var", default=None)
 
 
 def _set_span_tag_iast_request_tainted(span):
@@ -30,10 +29,6 @@ def _set_span_tag_iast_request_tainted(span):
         span.set_tag(IAST_SPAN_TAGS.TELEMETRY_REQUEST_TAINTED, total_objects_tainted)
 
 
-def finalize_iast_env(env: IASTEnvironment) -> None:
-    core.discard_item(IAST.REQUEST_CONTEXT_KEY)
-
-
 def get_iast_stacktrace_reported() -> bool:
     env = _get_iast_env()
     if env:
@@ -47,14 +42,6 @@ def set_iast_stacktrace_reported(reported: bool) -> None:
         env.iast_stack_trace_reported = reported
 
 
-def set_iast_request_enabled(request_enabled) -> None:
-    env = _get_iast_env()
-    if env:
-        env.request_enabled = request_enabled
-    else:
-        log.debug("iast::propagation::context::Trying to set IAST reporter but no context is present")
-
-
 def set_iast_request_endpoint(method, route) -> None:
     if asm_config._iast_enabled:
         env = _get_iast_env()
@@ -67,36 +54,72 @@ 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 _iast_start_request(span=None, *args, **kwargs):
-    try:
-        if asm_config._iast_enabled:
-            create_propagation_context()
-            core.set_item(IAST.REQUEST_CONTEXT_KEY, IASTEnvironment(span))
-            request_iast_enabled = False
-            if oce.acquire_request(span):
-                request_iast_enabled = True
-            set_iast_request_enabled(request_iast_enabled)
-    except Exception:
-        log.debug("iast::propagation::context::Error starting IAST context", exc_info=True)
+def _iast_start_request(span=None) -> Optional[int]:
+    """Initialize the IAST request context for the current execution.
+
+    This function acquires the IAST request budget via the Overhead Control Engine,
+    creates a new native taint context, and stores its identifier in a ContextVar so
+    subsequent IAST operations can locate the request-local taint map. If a context
+    is already active, the existing identifier is reused.
+
+    The provided span, when present, is attached to the IAST environment for later
+    enrichment and end-of-request processing.
+    """
+    context_id = None
+
+    if asm_config._iast_enabled:
+        if oce.acquire_request(span):
+            if not is_iast_request_enabled():
+                context_id = start_request_context()
+                IAST_CONTEXT.set(context_id)
+                if context_id is not None:
+                    core.set_item(IAST.REQUEST_CONTEXT_KEY, IASTEnvironment(span))
+        elif (context_id := _get_iast_context_id()) is not None:
+            finish_request_context(context_id)
+            IAST_CONTEXT.set(None)
+    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: Optional["Span"] = None):
+def _iast_finish_request(span=None, shoud_update_global_vulnerability_limit: bool = True) -> bool:
+    """Finalize the IAST request context and optionally update global limits.
+
+    This function discards the per-request IAST environment, optionally updates the
+    global vulnerability optimization data, and releases the native taint context
+    associated with the active request.
+    """
     env = _get_iast_env()
     if env is not None and env.span is span:
-        update_global_vulnerability_limit(env)
-        finalize_iast_env(env)
-    reset_propagation_context()
+        if shoud_update_global_vulnerability_limit:
+            update_global_vulnerability_limit(env)
+        core.discard_item(IAST.REQUEST_CONTEXT_KEY)
+
+    context_id = _get_iast_context_id()
+    print(f"context_id: {context_id}")
+    if context_id is not None:
+        finish_request_context(context_id)
+        print(f"finish_request_context: {context_id}")
+        IAST_CONTEXT.set(None)
+        return True
+
+    return False
 
 
 def is_iast_request_enabled() -> bool:
     """Check whether IAST is currently operating within an active request context."""
-    return asm_config.is_iast_request_enabled
+    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.
 
-def _num_objects_tainted_in_request():
-    return num_objects_tainted()
+    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