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

IAST-enabled applications using Gunicorn/Uvicorn workers were
experiencing segmentation faults (~33% crash rate on MCP streaming
requests) due to memory corruption when processes fork.

- C++ global singletons (`taint_engine_context`, `initializer`)
initialized at module load
- Taint maps storing PyObject pointers by memory address
- Child processes after fork inherited stale pointers from parent
process memory
- Accessing these stale pointers → use-after-free → SIGSEGV crash

- Implemented `pthread_atfork` handler that automatically resets C++
global state in child processes after every fork:
- Added comprehensive null-check wrappers around all native functions to
prevent crashes when native state is
- Fixed test regression issues where context slots weren't being freed:

**[AddressSanitizer
(ASAN)](https://github.com/google/sanitizers/wiki/AddressSanitizer)** is
a fast memory error detector that catches use-after-free, buffer
overflows, and other memory corruption bugs at runtime.

**1. Runtime Environment (No Recompilation Required)**

The simplest way to test is using LD_PRELOAD with the system's libasan:

```bash
ASAN_LIB=$(gcc -print-file-name=libasan.so)

LD_PRELOAD=$ASAN_LIB \
ASAN_OPTIONS="detect_leaks=0:symbolize=1:abort_on_error=0" \
python3 -m pytest tests/appsec/iast/test_fork_handler_regression.py -v
```
**ASAN_OPTIONS explained:**
- `detect_leaks=0` - Disable leak detection (Python has many false
positives)
- `symbolize=1` - Show human-readable stack traces
- `abort_on_error=0` - Continue after first error (collect all errors)

**2. Build with ASAN (Optional, for deeper analysis)**

For more thorough testing, compile the native extension with ASAN:

```bash
export CFLAGS="-fsanitize=address -fno-omit-frame-pointer -g"
export CXXFLAGS="-fsanitize=address -fno-omit-frame-pointer -g"
export LDFLAGS="-fsanitize=address"

pip install --no-build-isolation --force-reinstall -e .

ASAN_OPTIONS="detect_leaks=0:symbolize=1:abort_on_error=0" \
python3 -m pytest tests/appsec/iast/test_fork_handler_regression.py -v
```

This script demonstrates the fork safety fix and can be used to verify
ASAN finds no errors:

```python
"""Minimal fork safety reproduction test for ASAN verification."""
import os
from ddtrace.appsec._iast._taint_tracking import OriginType
from ddtrace.appsec._iast._taint_tracking._native import initialize_native_state
from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
from ddtrace.appsec._iast._taint_tracking._context import (
    start_request_context,
    debug_context_array_free_slots_number,
    debug_num_tainted_objects
)

def main():
    print(f"[Parent PID {os.getpid()}] Initializing IAST...")

    # Initialize native state
    initialize_native_state()

    # Create context and tainted objects in parent
    ctx_id = start_request_context()
    print(f"[Parent] Context created: {ctx_id}")

    # Create some tainted objects (populates native maps)
    for i in range(10):
        taint_pyobject(f"data_{i}", "source", f"value_{i}", OriginType.PARAMETER)

    tainted_count = debug_num_tainted_objects(ctx_id)
    print(f"[Parent] Tainted objects: {tainted_count}")

    # Fork
    pid = os.fork()

    if pid == 0:
        # Child process
        print(f"[Child PID {os.getpid()}] Started after fork")

        # Verify pthread_atfork reset worked
        free_slots = debug_context_array_free_slots_number()
        print(f"[Child] Free slots: {free_slots}")

        if free_slots > 0:
            print("[Child] ✅ Context slots were reset (pthread_atfork worked!)")
            os._exit(0)  # Success
        else:
            print("[Child] ❌ Context slots NOT reset")
            os._exit(1)  # Failure
    else:
        # Parent waits for child
        _, status = os.waitpid(pid, 0)
        exit_code = os.WEXITSTATUS(status)

        if exit_code == 0:
            print(f"[Parent] ✅ Child exited cleanly - Fork safety verified!")
            return 0
        else:
            print(f"[Parent] ❌ Child failed with exit code {exit_code}")
            return 1

if __name__ == "__main__":
    exit(main())
```

**Run with ASAN:**
```bash
LD_PRELOAD=$(gcc -print-file-name=libasan.so) \
ASAN_OPTIONS="detect_leaks=0:symbolize=1:abort_on_error=0" \
python3 test_fork_asan.py
```

**Expected output (success):**
```
[Parent PID 12345] Initializing IAST...
[Parent] Context created: 0
[Parent] Tainted objects: 10
[Child PID 12346] Started after fork
[Child] Free slots: 2
[Child] ✅ Context slots were reset (pthread_atfork worked!)
[Parent] ✅ Child exited cleanly - Fork safety verified!
```

**What ASAN would report WITHOUT this fix:**
```
==12346==ERROR: AddressSanitizer: SEGV on unknown address 0x000000000040
==12346==The signal is caused by a READ memory access.
    #0 0x7f... in get_tainted_object_map_by_ctx_id
    #1 0x7f... in debug_context_array_free_slots_number
```

(cherry picked from commit d1b4fd8)

Author: avara1986

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

@@ -59,6 +59,7 @@ file(
     GLOB
     SOURCE_FILES
     "*.cpp"
+    "api/*.cpp"
     "context/*.cpp"
     "aspects/*.cpp"
     "initializer/*.cpp"
@@ -69,6 +70,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,7 +67,8 @@
     "copy_ranges_from_strings",
     "get_range_by_hash",
     "get_ranges",
-    "is_in_taint_map",
+    "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);