Backport fix from 10706 to 2.10

Signed-off-by: Juanjo Alvarez <[email protected]>
DataDog · Sep 20, 2024 · 2d295dd · 2d295dd
1 parent f3fdff9
commit 2d295dd
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 51 deletions.
diff --git a/ddtrace/appsec/_iast/_taint_tracking/_native.cpp b/ddtrace/appsec/_iast/_taint_tracking/_native.cpp
@@ -2,7 +2,8 @@
  * IAST Native Module
  * This C++ module contains IAST propagation features:
  * - Taint tracking: propagation of a tainted variable.
- * - Aspects: common string operations are replaced by functions that propagate the taint variables.
+ * - Aspects: common string operations are replaced by functions that propagate
+ * the taint variables.
  * - Taint ranges: Information related to tainted values.
  */
 #include <memory>
@@ -19,74 +20,85 @@
 #include "TaintTracking/_taint_tracking.h"
 #include "TaintedOps/TaintedOps.h"
 
-#define PY_MODULE_NAME_ASPECTS                                                                                         \
-    PY_MODULE_NAME "."                                                                                                 \
-                   "aspects"
+#define PY_MODULE_NAME_ASPECTS                                                 \
+  PY_MODULE_NAME "."                                                           \
+                 "aspects"
 
 using namespace pybind11::literals;
 namespace py = pybind11;
 
 static PyMethodDef AspectsMethods[] = {
-    { "add_aspect", ((PyCFunction)api_add_aspect), METH_FASTCALL, "aspect add" },
-    { "extend_aspect", ((PyCFunction)api_extend_aspect), METH_FASTCALL, "aspect extend" },
-    { "index_aspect", ((PyCFunction)api_index_aspect), METH_FASTCALL, "aspect index" },
-    { "join_aspect", ((PyCFunction)api_join_aspect), METH_FASTCALL, "aspect join" },
-    { "slice_aspect", ((PyCFunction)api_slice_aspect), METH_FASTCALL, "aspect slice" },
-    { nullptr, nullptr, 0, nullptr }
-};
+    {"add_aspect", ((PyCFunction)api_add_aspect), METH_FASTCALL, "aspect add"},
+    {"extend_aspect", ((PyCFunction)api_extend_aspect), METH_FASTCALL,
+     "aspect extend"},
+    {"index_aspect", ((PyCFunction)api_index_aspect), METH_FASTCALL,
+     "aspect index"},
+    {"join_aspect", ((PyCFunction)api_join_aspect), METH_FASTCALL,
+     "aspect join"},
+    {"slice_aspect", ((PyCFunction)api_slice_aspect), METH_FASTCALL,
+     "aspect slice"},
+    {nullptr, nullptr, 0, nullptr}};
 
-static struct PyModuleDef aspects = { PyModuleDef_HEAD_INIT,
-                                      .m_name = PY_MODULE_NAME_ASPECTS,
-                                      .m_doc = "Taint tracking Aspects",
-                                      .m_size = -1,
-                                      .m_methods = AspectsMethods };
+static struct PyModuleDef aspects = {PyModuleDef_HEAD_INIT,
+                                     .m_name = PY_MODULE_NAME_ASPECTS,
+                                     .m_doc = "Taint tracking Aspects",
+                                     .m_size = -1, .m_methods = AspectsMethods};
 
 static PyMethodDef OpsMethods[] = {
-    { "new_pyobject_id", (PyCFunction)api_new_pyobject_id, METH_FASTCALL, "new pyobject id" },
-    { "set_ranges_from_values", ((PyCFunction)api_set_ranges_from_values), METH_FASTCALL, "set_ranges_from_values" },
-    { nullptr, nullptr, 0, nullptr }
-};
+    {"new_pyobject_id", (PyCFunction)api_new_pyobject_id, METH_FASTCALL,
+     "new pyobject id"},
+    {"set_ranges_from_values", ((PyCFunction)api_set_ranges_from_values),
+     METH_FASTCALL, "set_ranges_from_values"},
+    {nullptr, nullptr, 0, nullptr}};
 
-static struct PyModuleDef ops = { PyModuleDef_HEAD_INIT,
-                                  .m_name = PY_MODULE_NAME_ASPECTS,
-                                  .m_doc = "Taint tracking operations",
-                                  .m_size = -1,
-                                  .m_methods = OpsMethods };
+static struct PyModuleDef ops = {PyModuleDef_HEAD_INIT,
+                                 .m_name = PY_MODULE_NAME_ASPECTS,
+                                 .m_doc = "Taint tracking operations",
+                                 .m_size = -1, .m_methods = OpsMethods};
 
 /**
  * This function initializes the native module.
  */
-PYBIND11_MODULE(_native, m)
-{
-    const char* env_iast_enabled = std::getenv("DD_IAST_ENABLED");
-    if (env_iast_enabled == nullptr) {
-        py::module::import("logging").attr("warning")("IAST not enabled but native module is being loaded");
-    } else {
-        std::string iast_enabled = std::string(env_iast_enabled);
-        std::transform(iast_enabled.begin(), iast_enabled.end(), iast_enabled.begin(), [](unsigned char c) {
-            return std::tolower(c);
-        });
-        if (iast_enabled != "true" && iast_enabled != "1") {
-            py::module::import("logging").attr("warning")("IAST not enabled but native module is being loaded");
-        }
+PYBIND11_MODULE(_native, m) {
+  const char *env_iast_enabled = std::getenv("DD_IAST_ENABLED");
+  if (env_iast_enabled == nullptr) {
+    py::module::import("logging").attr("warning")(
+        "IAST not enabled but native module is being loaded");
+  } else {
+    std::string iast_enabled = std::string(env_iast_enabled);
+    std::transform(iast_enabled.begin(), iast_enabled.end(),
+                   iast_enabled.begin(),
+                   [](unsigned char c) { return std::tolower(c); });
+    if (iast_enabled != "true" && iast_enabled != "1") {
+      py::module::import("logging").attr("warning")(
+          "IAST not enabled but native module is being loaded");
     }
+  }
 
-    initializer = make_unique<Initializer>();
-    initializer->create_context();
+  initializer = make_unique<Initializer>();
+  // Create an atexit callback to clean up the Initializer before the
+  // interpreter finishes
+  auto atexit_register = py::module_::import("atexit").attr("register");
+  atexit_register(py::cpp_function([]() {
+    initializer->reset_context();
+    initializer.reset();
+  }));
 
-    m.doc() = "Native Python module";
+  initializer->create_context();
 
-    py::module m_initializer = pyexport_m_initializer(m);
-    pyexport_m_taint_tracking(m);
+  m.doc() = "Native Python module";
 
-    pyexport_m_aspect_helpers(m);
+  py::module m_initializer = pyexport_m_initializer(m);
+  pyexport_m_taint_tracking(m);
 
-    // Note: the order of these definitions matter. For example,
-    // stacktrace_element definitions must be before the ones of the
-    // classes inheriting from it.
-    PyObject* hm_aspects = PyModule_Create(&aspects);
-    m.add_object("aspects", hm_aspects);
+  pyexport_m_aspect_helpers(m);
 
-    PyObject* hm_ops = PyModule_Create(&ops);
-    m.add_object("ops", hm_ops);
+  // Note: the order of these definitions matter. For example,
+  // stacktrace_element definitions must be before the ones of the
+  // classes inheriting from it.
+  PyObject *hm_aspects = PyModule_Create(&aspects);
+  m.add_object("aspects", hm_aspects);
+
+  PyObject *hm_ops = PyModule_Create(&ops);
+  m.add_object("ops", hm_ops);
 }
diff --git a/releasenotes/notes/initializer-atexit-b353744cf1af7917.yaml b/releasenotes/notes/initializer-atexit-b353744cf1af7917.yaml
@@ -0,0 +1,11 @@
+---
+fixes:
+  - |
+    Code Security: ensure the ``Initializer`` object is always reset and freed before the Python runtime.
+other:
+  - |
+    For any change which does not fall into any of the above categories. Since changes falling into this category are 
+    likely rare and not very similar to each other, no specific format other than a required scope is provided. 
+    The author is requested to use their best judgment to ensure a quality release note.
+    Format: <scope>: <add_release_note_here>.
+