chore(iast): context refactor, context-array + contextvar indexing

.gitignore

@@ -138,7 +138,8 @@ ENV/
 .cursor/
 
 # Cascade
-.windsurf/
+# .windsurf/  # If you want to use Windsurf/Cascade, do NOT ignore here.
+# Add '.windsurf/' to .git/info/exclude instead so Cascade can read workflow/config files (gitignore would hide them)
 
 # Claude
 CLAUDE.local.md

ddtrace/appsec/_iast/_iast_request_context_base.py

@@ -1,3 +1,4 @@
+import contextvars
 from typing import Optional
 
 from ddtrace._trace.span import Span
@@ -19,6 +20,8 @@
 
 # Stopgap module for providing ASM context for the blocking features wrapping some contextvars.
 
+IAST_CONTEXT = contextvars.ContextVar("iast_var", default=None)
+
 
 def _set_span_tag_iast_request_tainted(span):
     total_objects_tainted = _num_objects_tainted_in_request()

ddtrace/appsec/_iast/_taint_tracking/CMakeLists.txt

@@ -58,6 +58,7 @@ file(
     GLOB
     SOURCE_FILES
     "*.cpp"
+    "context/*.cpp"
     "aspects/*.cpp"
     "initializer/*.cpp"
     "tainted_ops/*.cpp"
@@ -67,6 +68,7 @@ file(
     GLOB
     HEADER_FILES
     "*.h"
+    "context/*.h"
     "aspects/*.h"
     "initializer/*.h"
     "tainted_ops/*.h"

ddtrace/appsec/_iast/_taint_tracking/__init__.py

@@ -18,9 +18,9 @@
 from ddtrace.appsec._iast._taint_tracking._native.aspects_ospath import _aspect_ospathsplitdrive  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.aspects_ospath import _aspect_ospathsplitext  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.aspects_ospath import _aspect_ospathsplitroot  # noqa: F401
+from ddtrace.appsec._iast._taint_tracking._native.context import get_taint_map  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.initializer import active_map_addreses_size  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.initializer import debug_taint_map  # noqa: F401
-from ddtrace.appsec._iast._taint_tracking._native.initializer import initializer_size  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.initializer import num_objects_tainted  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import OriginType  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import Source  # noqa: F401
@@ -47,6 +47,7 @@
 
 __all__ = [
     "OriginType",
+    "active_map_addreses_size",
     "Source",
     "TagMappingMode",
     "TaintRange",
@@ -64,7 +65,6 @@
     "_aspect_splitlines",
     "_convert_escaped_text_to_tainted_text",
     "_format_aspect",
-    "active_map_addreses_size",
     "are_all_text_all_ranges",
     "as_formatted_evidence",
     "common_replace",
@@ -73,7 +73,7 @@
     "debug_taint_map",
     "get_range_by_hash",
     "get_ranges",
-    "initializer_size",
+    "get_taint_map",
     "is_tainted",
     "new_pyobject_id",
     "num_objects_tainted",
@@ -87,3 +87,4 @@
 ]
 new_pyobject_id = ops.new_pyobject_id
 set_ranges_from_values = ops.set_ranges_from_values
+taint_pyobject = ops.taint_pyobject

ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py

@@ -4,12 +4,14 @@
 
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
+from ddtrace.appsec._iast._iast_request_context_base import IAST_CONTEXT
 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._taint_objects_base import _taint_pyobject_base
+from ddtrace.appsec._iast._taint_tracking._taint_objects_base import _taint_pyobject_base_new
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
@@ -32,6 +34,21 @@ def taint_pyobject(pyobject: Any, source_name: Any, source_value: Any, source_or
     return pyobject
 
 
+def taint_pyobject_new(pyobject: Any, source_name: Any, source_value: Any, source_origin=None) -> Any:
+    try:
+        if (contextid := IAST_CONTEXT.get()) is not None:
+            if source_origin is None:
+                source_origin = OriginType.PARAMETER
+
+            res = _taint_pyobject_base_new(pyobject, source_name, source_value, source_origin, contextid)
+            _set_metric_iast_executed_source(source_origin)
+            increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SOURCE, source_origin)
+            return res
+    except ValueError:
+        iast_propagation_debug_log(f"taint_pyobject error (pyobject type {type(pyobject)})", exc_info=True)
+    return pyobject
+
+
 def copy_ranges_to_string(pyobject: str, ranges: Sequence[TaintRange]) -> str:
     # NB this function uses comment-based type annotation because TaintRange is conditionally imported
     if asm_config.is_iast_request_enabled:

ddtrace/appsec/_iast/_taint_tracking/_taint_objects_base.py

@@ -6,10 +6,12 @@
 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 get_taint_map
 from ddtrace.appsec._iast._taint_tracking import is_tainted
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
 from ddtrace.appsec._iast._taint_tracking import set_ranges
 from ddtrace.appsec._iast._taint_tracking import set_ranges_from_values
+from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 from ddtrace.settings.asm import config as asm_config
 
 
@@ -43,32 +45,77 @@ def _taint_pyobject_base(pyobject: Any, source_name: Any, source_value: Any, sou
         - Returns unmodified object for empty strings
         - Automatically handles bytes/bytearray to str conversion
     """
-    # Early type validation
-    if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):
+    if not isinstance(pyobject, IAST.TAINTEABLE_TYPES) or not pyobject:
         return pyobject
 
-    # Fast path for empty strings
-    if isinstance(pyobject, IAST.TEXT_TYPES) and not pyobject:
+    if isinstance(source_name, (bytes, bytearray)):
+        source_name = source_name.decode("utf-8", errors="ignore")
+    elif isinstance(source_name, OriginType):
+        source_name = origin_to_str(source_name)
+
+    if isinstance(source_value, (bytes, bytearray)):
+        source_value = source_value.decode("utf-8", errors="ignore")
+
+    if source_origin is None:
+        source_origin = OriginType.PARAMETER
+
+    try:
+        pyobject_len = len(pyobject) if isinstance(pyobject, IAST.TEXT_TYPES) else 0
+        return set_ranges_from_values(pyobject, pyobject_len, source_name, source_value, source_origin)
+    except ValueError:
+        iast_propagation_debug_log(f"Tainting object error (pyobject type {type(pyobject)})", exc_info=True)
+        return pyobject
+
+
+def _taint_pyobject_base_new(
+    pyobject: Any, source_name: Any, source_value: Any, source_origin=None, contextid=None
+) -> Any:
+    """Mark a Python object as tainted with information about its origin.
+
+    This function is the base for marking objects as tainted, setting their origin and range.
+    It is optimized for:
+    1. Early validations to avoid unnecessary operations
+    2. Efficient type conversions
+    3. Special case handling (empty objects)
+    4. Robust error handling
+
+    Performance optimizations:
+    - Early return for disabled IAST or non-taintable types
+    - Efficient string length calculation only when needed
+    - Optimized bytes/bytearray to string conversion using decode()
+    - Minimized object allocations and method calls
+
+    Args:
+        pyobject (Any): The object to mark as tainted. Must be a taintable type.
+        source_name (Any): Name of the taint source (e.g., parameter name).
+        source_value (Any): Original value that caused the taint.
+        source_origin (Optional[OriginType]): Origin of the taint. Defaults to PARAMETER.
+
+    Returns:
+        Any: The tainted object if operation was successful, original object if failed.
+
+    Note:
+        - Only applies to taintable types defined in IAST.TAINTEABLE_TYPES
+        - Returns unmodified object for empty strings
+        - Automatically handles bytes/bytearray to str conversion
+    """
+    if not isinstance(pyobject, IAST.TAINTEABLE_TYPES) or not pyobject:
         return pyobject
 
-    # Efficient source_name conversion
     if isinstance(source_name, (bytes, bytearray)):
         source_name = source_name.decode("utf-8", errors="ignore")
     elif isinstance(source_name, OriginType):
         source_name = origin_to_str(source_name)
 
-    # Efficient source_value conversion
     if isinstance(source_value, (bytes, bytearray)):
         source_value = source_value.decode("utf-8", errors="ignore")
 
-    # Default source_origin
     if source_origin is None:
         source_origin = OriginType.PARAMETER
 
     try:
-        # Calculate length only for text types
         pyobject_len = len(pyobject) if isinstance(pyobject, IAST.TEXT_TYPES) else 0
-        return set_ranges_from_values(pyobject, pyobject_len, source_name, source_value, source_origin)
+        return taint_pyobject(pyobject, pyobject_len, source_name, source_value, source_origin, contextid)
     except ValueError:
         iast_propagation_debug_log(f"Tainting object error (pyobject type {type(pyobject)})", exc_info=True)
         return pyobject
@@ -87,8 +134,6 @@ def get_tainted_ranges(pyobject: Any) -> Tuple:
 
 
 def is_pyobject_tainted(pyobject: Any) -> bool:
-    if not asm_config.is_iast_request_enabled:
-        return False
     if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):
         return False
 
@@ -99,6 +144,17 @@ def is_pyobject_tainted(pyobject: Any) -> bool:
     return False
 
 
+def is_pyobject_tainted_new(pyobject: Any) -> bool:
+    if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):
+        return False
+
+    try:
+        return get_taint_map(pyobject)
+    except ValueError as e:
+        iast_propagation_error_log(f"Checking tainted object error: {e}")
+    return False
+
+
 def taint_pyobject_with_ranges(pyobject: Any, ranges: Tuple) -> bool:
     if not asm_config.is_iast_request_enabled:
         return False

ddtrace/appsec/_iast/_taint_tracking/context/_application_context.h

@@ -0,0 +1,13 @@
+#pragma once
+#include <pybind11/pybind11.h>
+
+#include "context/application_context.h"
+
+inline py::module
+pyexport_m_application_context(py::module& m)
+{
+    // ApplicationContext
+    py::module m_context = m.def_submodule("context", "Application Context");
+    pyexport_application_context(m_context);
+    return m_context;
+}