fix(iast): wrong memory address in subprocess in mcp servers

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

@@ -591,8 +591,8 @@ experiments:
               - max_rss_usage < 41.50 MB
           - name: iastaspects-slice_aspect
             thresholds:
-              - execution_time < 0.01 ms
-              - max_rss_usage < 41.50 MB
+              - execution_time < 0.02 ms
+              - max_rss_usage < 110.50 MB
           - name: iastaspects-slice_noaspect
             thresholds:
               - execution_time < 0.01 ms

ddtrace/appsec/_iast/__init__.py

@@ -51,33 +51,37 @@ def wrapped_function(wrapped, instance, args, kwargs):
 
 def _disable_iast_after_fork():
     """
-    Conditionally disable IAST in forked child processes to prevent segmentation faults.
+    Handle IAST state after fork to prevent segmentation faults.
 
-    This fork handler differentiates between two types of forks:
+    This fork handler works in conjunction with the C++ pthread_atfork handler
+    (registered in native.cpp) to provide complete fork-safety:
 
-    1. **Early forks (web framework workers)**: Gunicorn, uvicorn, Django, Flask workers
-       fork BEFORE IAST initializes any state. These are safe - IAST remains enabled.
+    **C++ pthread_atfork handler (automatic, runs first):**
+    - Clears all taint maps (removes stale PyObject pointers from parent)
+    - Resets taint_engine_context (recreates context array with fresh state)
+    - Resets initializer (recreates memory pools)
+    - Happens automatically for ALL forks before this Python handler runs
 
-    2. **Late forks (multiprocessing)**: multiprocessing.Process forks AFTER IAST has
-       initialized state. These inherit corrupted native extension state and must have
-       IAST disabled to prevent segmentation faults.
+    **Python fork handler (this function, runs second):**
+    - Detects fork type (early vs late)
+    - Manages Python-level IAST_CONTEXT
+    - Conditionally disables IAST for late forks (multiprocessing)
 
-    Detection logic:
-    - If IAST has active request contexts when fork occurs → Late fork → Disable IAST
-    - If IAST has no active state → Early fork (worker) → Keep IAST enabled
+    Fork types:
+    1. **Early forks (web framework workers)**: Gunicorn, uvicorn, etc. fork BEFORE
+       IAST has active request contexts. Native state was reset by pthread_atfork.
+       IAST remains enabled and works correctly in workers.
 
-    This is critical for multiprocessing compatibility while maintaining IAST coverage
-    in web framework workers. The native extension state (taint maps, context slots,
-    object pools, shared_ptr references) cannot be safely used across fork boundaries
-    when it exists, but is safe to initialize fresh in clean workers.
+    2. **Late forks (multiprocessing)**: fork AFTER IAST has active contexts.
+       Native state was reset by pthread_atfork, but we disable IAST in these
+       processes for multiprocessing compatibility and to avoid confusion.
 
-    For late forks, the child process:
-    - Clears all C++ taint maps and context slots
-    - Resets the Python-level IAST_CONTEXT
-    - Disables IAST by setting asm_config._iast_enabled = False
+    Detection logic:
+    - If is_iast_request_enabled() → Late fork → Disable IAST
+    - If not is_iast_request_enabled() → Early fork → Keep IAST enabled
 
-    This prevents segmentation faults in multiprocessing while allowing IAST to work
-    in web framework workers.
+    This prevents segmentation faults in ALL fork scenarios while maintaining
+    IAST coverage in web framework workers.
     """
     if not asm_config._iast_enabled:
         return
@@ -86,33 +90,45 @@ def _disable_iast_after_fork():
         # Import locally to avoid issues if the module hasn't been loaded yet
         from ddtrace.appsec._iast._iast_request_context_base import IAST_CONTEXT
         from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled
-        from ddtrace.appsec._iast._taint_tracking._context import clear_all_request_context_slots
+
+        # Note: The C++ pthread_atfork handler (in native.cpp) automatically resets
+        # native state in actual fork scenarios. It clears the inherited state and
+        # creates fresh instances without touching the invalid PyObject pointers.
+        # We don't need to (and shouldn't) call clear_all_request_context_slots()
+        # here because the C++ handler has already done the necessary cleanup.
 
         # In pytest mode, always disable IAST in child processes to avoid segfaults
         # when tests create multiprocesses (e.g., for testing fork behavior)
         if _iast_in_pytest_mode:
             log.debug("IAST fork handler: Pytest mode detected, disabling IAST in child process")
-            clear_all_request_context_slots()
             IAST_CONTEXT.set(None)
             asm_config._iast_enabled = False
             return
 
         if not is_iast_request_enabled():
             # No active context - this is an early fork (web framework worker)
-            # IAST can be safely initialized fresh in this child process
-            log.debug("IAST fork handler: No active context, keeping IAST enabled (web worker fork)")
+            # The C++ pthread_atfork handler has already reset native state with fresh instances.
+            # IAST can continue working correctly in this child process.
+            log.debug(
+                "IAST fork handler: No active context (early fork/web worker). "
+                "Native state auto-reset by pthread_atfork. IAST remains enabled."
+            )
+            # Clear Python-side context just in case
+            IAST_CONTEXT.set(None)
             return
 
         # Active context exists - this is a late fork (multiprocessing)
-        # Native state is corrupted, must disable IAST
-        log.debug("IAST fork handler: Active context detected, disabling IAST (multiprocessing fork)")
+        # The C++ pthread_atfork handler has already reset native state.
+        # We disable IAST in these processes for consistency.
+        log.debug(
+            "IAST fork handler: Active context (late fork/multiprocessing). "
+            "Native state auto-reset by pthread_atfork. Disabling IAST in child."
+        )
 
-        # Clear C++ side: all taint maps and context slots
-        clear_all_request_context_slots()
         # Clear Python side: reset the context ID
         IAST_CONTEXT.set(None)
 
-        # Disable IAST to prevent segmentation faults
+        # Disable IAST for multiprocessing compatibility
         asm_config._iast_enabled = False
 
     except Exception as e:
@@ -182,6 +198,7 @@ def enable_iast_propagation():
     if asm_config._iast_propagation_enabled:
         from ddtrace.appsec._iast._ast.ast_patching import _should_iast_patch
         from ddtrace.appsec._iast._loader import _exec_iast_patched_module
+        from ddtrace.appsec._iast._taint_tracking import initialize_native_state
 
         global _iast_propagation_enabled
         if _iast_propagation_enabled:
@@ -190,6 +207,7 @@ def enable_iast_propagation():
         log.debug("iast::instrumentation::starting IAST")
         ModuleWatchdog.register_pre_exec_module_hook(_should_iast_patch, _exec_iast_patched_module)
         _iast_propagation_enabled = True
+        initialize_native_state()
         _register_fork_handler()
 
 

ddtrace/appsec/_iast/_taint_tracking/CMakeLists.txt

@@ -63,6 +63,7 @@ file(
     GLOB
     SOURCE_FILES
     "*.cpp"
+    "api/*.cpp"
     "context/*.cpp"
     "aspects/*.cpp"
     "initializer/*.cpp"
@@ -73,6 +74,7 @@ file(
     GLOB
     HEADER_FILES
     "*.h"
+    "api/*.h"
     "context/*.h"
     "aspects/*.h"
     "initializer/*.h"

ddtrace/appsec/_iast/_taint_tracking/__init__.py

@@ -1,4 +1,6 @@
+from ddtrace.appsec._iast._taint_tracking._native import initialize_native_state  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native import ops  # noqa: F401
+from ddtrace.appsec._iast._taint_tracking._native import reset_native_state  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.aspect_format import _format_aspect  # noqa: F401
 from ddtrace.appsec._iast._taint_tracking._native.aspect_helpers import _convert_escaped_text_to_tainted_text
 
@@ -65,6 +67,8 @@
     "copy_ranges_from_strings",
     "get_range_by_hash",
     "get_ranges",
+    "reset_native_state",
+    "initialize_native_state",
     "is_tainted",
     "new_pyobject_id",
     "origin_to_str",

ddtrace/appsec/_iast/_taint_tracking/api/safe_context.cpp

@@ -0,0 +1,35 @@
+#include "api/safe_context.h"
+
+// ============================================================================
+// Safe wrapper functions for global pointers
+// ============================================================================
+// These functions centralize null checks for taint_engine_context and
+// initializer to prevent segmentation faults when called before
+// initialize_native_state(). All aspect functions should use these wrappers
+// instead of accessing the globals directly.
+TaintedObjectMapTypePtr
+safe_get_tainted_object_map(PyObject* tainted_object)
+{
+    if (!taint_engine_context) {
+        return nullptr;
+    }
+    return taint_engine_context->get_tainted_object_map(tainted_object);
+}
+
+TaintedObjectMapTypePtr
+safe_get_tainted_object_map_from_list_of_pyobjects(const std::vector<PyObject*>& objects)
+{
+    if (!taint_engine_context) {
+        return nullptr;
+    }
+    return taint_engine_context->get_tainted_object_map_from_list_of_pyobjects(objects);
+}
+
+TaintedObjectMapTypePtr
+safe_get_tainted_object_map_by_ctx_id(size_t ctx_id)
+{
+    if (!taint_engine_context) {
+        return nullptr;
+    }
+    return taint_engine_context->get_tainted_object_map_by_ctx_id(ctx_id);
+}

ddtrace/appsec/_iast/_taint_tracking/api/safe_context.h

@@ -0,0 +1,30 @@
+#pragma once
+#include "context/taint_engine_context.h"
+
+// Safe wrapper functions that check for null before accessing global pointers
+// These functions centralize null checks to prevent segmentation faults when
+// called before initialize_native_state()
+
+/**
+ * @brief Safely get tainted object map for a PyObject.
+ * @param tainted_object The Python object to look up
+ * @return TaintedObjectMapTypePtr or nullptr if not initialized or not found
+ */
+TaintedObjectMapTypePtr
+safe_get_tainted_object_map(PyObject* tainted_object);
+
+/**
+ * @brief Safely get tainted object map by context ID.
+ * @param ctx_id The context ID to look up
+ * @return TaintedObjectMapTypePtr or nullptr if not initialized or not found
+ */
+TaintedObjectMapTypePtr
+safe_get_tainted_object_map_by_ctx_id(size_t ctx_id);
+
+/**
+ * @brief Safely get tainted object map from a list of PyObjects.
+ * @param objects Vector of PyObject pointers to search
+ * @return TaintedObjectMapTypePtr or nullptr if not initialized or not found
+ */
+TaintedObjectMapTypePtr
+safe_get_tainted_object_map_from_list_of_pyobjects(const std::vector<PyObject*>& objects);

ddtrace/appsec/_iast/_taint_tracking/api/safe_initializer.cpp

@@ -0,0 +1,46 @@
+#include "api/safe_initializer.h"
+
+TaintRangePtr
+safe_allocate_taint_range(RANGE_START start, RANGE_LENGTH length, const Source& source, const SecureMarks secure_marks)
+{
+    if (!initializer) {
+        return nullptr;
+    }
+    return initializer->allocate_taint_range(start, length, source, secure_marks);
+}
+
+TaintedObjectPtr
+safe_allocate_tainted_object_copy(const TaintedObjectPtr& from)
+{
+    if (!initializer) {
+        return nullptr;
+    }
+    return initializer->allocate_tainted_object_copy(from);
+}
+
+TaintedObjectPtr
+safe_allocate_tainted_object()
+{
+    if (!initializer) {
+        return nullptr;
+    }
+    return initializer->allocate_tainted_object();
+}
+
+TaintedObjectPtr
+safe_allocate_ranges_into_taint_object(TaintRangeRefs ranges)
+{
+    if (!initializer) {
+        return nullptr;
+    }
+    return initializer->allocate_ranges_into_taint_object(ranges);
+}
+
+TaintedObjectPtr
+safe_allocate_ranges_into_taint_object_copy(const TaintRangeRefs& ranges)
+{
+    if (!initializer) {
+        return nullptr;
+    }
+    return initializer->allocate_ranges_into_taint_object_copy(ranges);
+}

ddtrace/appsec/_iast/_taint_tracking/api/safe_initializer.h

@@ -0,0 +1,49 @@
+#pragma once
+
+#include "initializer/initializer.h"
+#include <Python.h>
+
+/**
+ * @brief Safely allocate a taint range.
+ * @param start The start position of the taint range
+ * @param length The length of the taint range
+ * @param source The source of the taint
+ * @param secure_marks Optional secure marks
+ * @return TaintRangePtr or nullptr if not initialized
+ */
+TaintRangePtr
+safe_allocate_taint_range(RANGE_START start,
+                          RANGE_LENGTH length,
+                          const Source& source,
+                          const SecureMarks secure_marks = 0);
+
+/**
+ * @brief Safely allocate a copy of a tainted object.
+ * @param from The tainted object to copy
+ * @return TaintedObjectPtr or nullptr if not initialized
+ */
+TaintedObjectPtr
+safe_allocate_tainted_object_copy(const TaintedObjectPtr& from);
+
+/**
+ * @brief Safely allocate a tainted object.
+ * @return TaintedObjectPtr or nullptr if not initialized
+ */
+TaintedObjectPtr
+safe_allocate_tainted_object();
+
+/**
+ * @brief Safely allocate ranges into a tainted object.
+ * @param ranges The ranges to allocate
+ * @return TaintedObjectPtr or nullptr if not initialized
+ */
+TaintedObjectPtr
+safe_allocate_ranges_into_taint_object(TaintRangeRefs ranges);
+
+/**
+ * @brief Safely allocate a copy of ranges into a tainted object.
+ * @param ranges The ranges to copy and allocate
+ * @return TaintedObjectPtr or nullptr if not initialized
+ */
+TaintedObjectPtr
+safe_allocate_ranges_into_taint_object_copy(const TaintRangeRefs& ranges);

ddtrace/appsec/_iast/_taint_tracking/api/utils.h

@@ -0,0 +1,4 @@
+#define CHECK_IAST_INITIALIZED_OR_RETURN(fallback_result)                                                              \
+    if (!taint_engine_context || !initializer) {                                                                       \
+        return fallback_result;                                                                                        \
+    }

ddtrace/appsec/_iast/_taint_tracking/aspects/aspect_extend.cpp

@@ -29,7 +29,7 @@ api_extend_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
         return nullptr;
     }
 
-    auto ctx_map = taint_engine_context->get_tainted_object_map_from_list_of_pyobjects({ candidate_text, to_add });
+    auto ctx_map = safe_get_tainted_object_map_from_list_of_pyobjects({ candidate_text, to_add });
     if (not ctx_map or ctx_map->empty()) {
         auto method_name = PyUnicode_FromString("extend");
         PyObject_CallMethodObjArgs(candidate_text, method_name, to_add, nullptr);
@@ -40,7 +40,7 @@ api_extend_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
         Py_DecRef(method_name);
     } else {
         const auto& to_candidate = get_tainted_object(candidate_text, ctx_map);
-        auto to_result = initializer->allocate_tainted_object_copy(to_candidate);
+        auto to_result = safe_allocate_tainted_object_copy(to_candidate);
         const auto& to_toadd = get_tainted_object(to_add, ctx_map);
 
         // Ensure no returns are done before this method call

ddtrace/appsec/_iast/_taint_tracking/aspects/aspect_extend.h

@@ -1,7 +1,7 @@
 #pragma once
 
-#include "initializer/initializer.h"
-#include <Python.h>
+#include "api/safe_context.h"
+#include "api/safe_initializer.h"
 
 PyObject*
 api_extend_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs);

ddtrace/appsec/_iast/_taint_tracking/aspects/aspect_format.cpp

@@ -18,16 +18,17 @@ api_format_aspect(StrType& candidate_text,
                   const py::args& args,
                   const py::kwargs& kwargs)
 {
-    auto tx_map = taint_engine_context->get_tainted_object_map_from_list_of_pyobjects(
-      { candidate_text.ptr(), parameter_list.ptr() });
+    auto return_candidate_result = [&]() -> StrType { return py::getattr(candidate_text, "format")(*args, **kwargs); };
+
+    auto tx_map = safe_get_tainted_object_map_from_list_of_pyobjects({ candidate_text.ptr(), parameter_list.ptr() });
 
     if (not tx_map or tx_map->empty()) {
-        return py::getattr(candidate_text, "format")(*args, **kwargs);
+        return return_candidate_result();
     }
 
     auto [ranges_orig, candidate_text_ranges] = are_all_text_all_ranges(candidate_text.ptr(), parameter_list, tx_map);
     if (ranges_orig.empty() and candidate_text_ranges.empty()) {
-        return py::getattr(candidate_text, "format")(*args, **kwargs);
+        return return_candidate_result();
     }
 
     auto new_template =

ddtrace/appsec/_iast/_taint_tracking/aspects/aspect_format.h

@@ -1,6 +1,7 @@
 #pragma once
+#include "api/safe_context.h"
+#include "api/safe_initializer.h"
 #include "helpers.h"
-#include "initializer/initializer.h"
 
 template<class StrType>
 StrType