fix(iast): gevent

Author: avara1986

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

@@ -34,8 +34,8 @@
 
 #include <atomic>
 #include <memory>
-#include <optional>
 #include <mutex>
+#include <optional>
 #include <string>
 #include <vector>
 

ddtrace/appsec/_iast/_taint_tracking/tests/test_taint_engine_context.cpp

@@ -397,3 +397,70 @@ TEST_F(ApplicationContextTest, ClearAllIsIdempotent)
     taint_engine_context->clear_all_request_context_slots();
     ASSERT_EQ(taint_engine_context->get_tainted_object_map_by_ctx_id(id), nullptr);
 }
+
+TEST_F(ApplicationContextTest, ShuttingDownBlocksStartAndLookups)
+{
+    // Enable shutdown: all operations should early-return
+    TaintEngineContext::set_shutting_down(true);
+
+    // start_request_context should refuse to allocate a slot
+    auto idx = taint_engine_context->start_request_context();
+    ASSERT_FALSE(idx.has_value());
+
+    // Lookups should be short-circuited
+    py::str s("x");
+    auto m_obj = taint_engine_context->get_tainted_object_map(s.ptr());
+    ASSERT_EQ(m_obj, nullptr);
+
+    std::vector<PyObject*> objs{ s.ptr() };
+    auto m_list = taint_engine_context->get_tainted_object_map_from_list_of_pyobjects(objs);
+    ASSERT_EQ(m_list, nullptr);
+
+    // By-index lookup should also be blocked
+    auto m_by_id = taint_engine_context->get_tainted_object_map_by_ctx_id(0);
+    ASSERT_EQ(m_by_id, nullptr);
+
+    // Reset shutdown flag for following tests
+    TaintEngineContext::set_shutting_down(false);
+}
+
+TEST_F(ApplicationContextTest, ClearingShuttingDownRestoresOperation)
+{
+    // Block then unblock and verify operations resume
+    TaintEngineContext::set_shutting_down(true);
+    ASSERT_FALSE(taint_engine_context->start_request_context().has_value());
+
+    TaintEngineContext::set_shutting_down(false);
+    auto idx = taint_engine_context->start_request_context();
+    ASSERT_TRUE(idx.has_value());
+    auto m = taint_engine_context->get_tainted_object_map_by_ctx_id(*idx);
+    ASSERT_NE(m, nullptr);
+}
+
+TEST_F(ApplicationContextTest, AtexitHandlerClearsAndResetsContext)
+{
+    // Import the native module to trigger its PYBIND11_MODULE initializer,
+    // which registers an atexit handler that sets shutting_down=true,
+    // clears slots, and resets the global taint_engine_context.
+    auto mod = py::module::import("ddtrace.appsec._iast._taint_tracking._native");
+    auto ctx = mod.attr("context");
+
+    // Before running exit funcs, capacity should be > 0 (module created a context)
+    auto before_size = ctx.attr("debug_context_array_size")().cast<size_t>();
+    ASSERT_GT(before_size, static_cast<size_t>(0));
+
+    // Execute Python atexit callbacks
+    auto atexit = py::module::import("atexit");
+    atexit.attr("_run_exitfuncs")();
+
+    // After atexit, context helpers should reflect cleared/reset state
+    auto after_size = ctx.attr("debug_context_array_size")().cast<size_t>();
+    ASSERT_EQ(after_size, static_cast<size_t>(0));
+
+    // get_tainted_object_map_by_ctx_id returns bool; should be False (nullptr)
+    auto by_id_exists = ctx.attr("get_tainted_object_map_by_ctx_id")(0).cast<bool>();
+    ASSERT_FALSE(by_id_exists);
+
+    // Recreate a fresh context manually to not affect subsequent tests
+    taint_engine_context = std::make_unique<TaintEngineContext>();
+}