fix(iast): fix memory leak error [backport #7788 to 1.20] (#7820)

This backports below backports to 1.20.
- fix(iast): iast leak for python 3.11:
#7788
- chore(iast): add propagation for str aspect join:
#6940
- chore(iast): add collision tests for bytes and bytearray:
#7367
- chore(iast): store hash and id separately in tx_taint_map:
#7340

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

Author: avara1986

ddtrace/appsec/iast/_stacktrace.c

@@ -17,18 +17,22 @@
 #define GET_LINENO(frame) PyFrame_GetLineNumber((PyFrameObject*)frame)
 #define GET_FRAME(tstate) PyThreadState_GetFrame(tstate)
 #define GET_PREVIOUS(frame) PyFrame_GetBack(frame)
-#define FRAME_DECREF(frame) Py_DECREF(frame)
+#define FRAME_DECREF(frame) Py_DecRef(frame)
 #define FRAME_XDECREF(frame) Py_XDECREF(frame)
-#define FILENAME_DECREF(filename) Py_DECREF(filename)
-#define FILENAME_XDECREF(filename) Py_XDECREF(filename)
+#define FILENAME_DECREF(filename) Py_DecRef(filename)
+#define FILENAME_XDECREF(filename)                                                                                     \
+    if (filename)                                                                                                      \
+    Py_DecRef(filename)
 static inline PyObject*
 GET_FILENAME(PyFrameObject* frame)
 {
     PyCodeObject* code = PyFrame_GetCode(frame);
     if (!code) {
         return NULL;
     }
-    return PyObject_GetAttrString((PyObject*)code, "co_filename");
+    PyObject* filename = PyObject_GetAttrString((PyObject*)code, "co_filename");
+    Py_DecRef(code);
+    return filename;
 }
 #else
 #define GET_FRAME(tstate) tstate->frame
@@ -59,7 +63,7 @@ get_file_and_line(PyObject* Py_UNUSED(module), PyObject* cwd_obj)
 {
     PyThreadState* tstate = PyThreadState_Get();
     if (!tstate) {
-        return NULL;
+        goto exit_0;
     }
 
     int line;
@@ -69,18 +73,16 @@ get_file_and_line(PyObject* Py_UNUSED(module), PyObject* cwd_obj)
     char* cwd = NULL;
 
     if (!PyUnicode_FSConverter(cwd_obj, &cwd_bytes)) {
-        return NULL;
+        goto exit_0;
     }
     cwd = PyBytes_AsString(cwd_bytes);
     if (!cwd) {
-        Py_DECREF(cwd_bytes);
-        return NULL;
+        goto exit_0;
     }
 
     PyFrameObject* frame = GET_FRAME(tstate);
     if (!frame) {
-        Py_DECREF(cwd_bytes);
-        return NULL;
+        goto exit_0;
     }
 
     while (NULL != frame) {
@@ -109,10 +111,25 @@ get_file_and_line(PyObject* Py_UNUSED(module), PyObject* cwd_obj)
         result = PyTuple_Pack(2, filename_o, line_obj);
         break;
     }
+    if (result == NULL) {
+        goto exit_0;
+    }
+
 exit:
-    Py_DECREF(cwd_bytes);
+    Py_DecRef(cwd_bytes);
+    FRAME_XDECREF(frame);
+    FILENAME_XDECREF(filename_o);
+    return result;
+
+exit_0:; // fix: "a label can only be part of a statement and a declaration is not a statement" error
+    // Return "", 0
+    PyObject* line_obj = Py_BuildValue("i", 0);
+    filename_o = PyUnicode_FromString("");
+    result = PyTuple_Pack(2, filename_o, line_obj);
+    Py_DecRef(cwd_bytes);
     FRAME_XDECREF(frame);
     FILENAME_XDECREF(filename_o);
+    Py_DecRef(line_obj);
     return result;
 }
 

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectExtend.cpp

@@ -1,5 +1,4 @@
 #include "AspectExtend.h"
-#include "Initializer/Initializer.h"
 
 PyObject*
 api_extend_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
@@ -28,7 +27,7 @@ api_extend_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
     // Ensure no returns are done before this method call
     auto method_name = PyUnicode_FromString("extend");
     PyObject_CallMethodObjArgs(candidate_text, method_name, to_add, nullptr);
-    Py_DECREF(method_name);
+    Py_DecRef(method_name);
 
     if (not ctx_map or ctx_map->empty() or to_result == nullptr) {
         Py_RETURN_NONE;

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectExtend.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "Initializer/Initializer.h"
 #include <Python.h>
 
 PyObject*

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectFormat.cpp

@@ -0,0 +1,46 @@
+#pragma once
+#include "Aspects/AspectFormat.h"
+
+template<class StrType>
+StrType
+api_format_aspect(StrType& candidate_text,
+                  const py::tuple& parameter_list,
+                  const py::args& args,
+                  const py::kwargs& kwargs)
+{
+    auto [ranges_orig, candidate_text_ranges] = are_all_text_all_ranges(candidate_text.ptr(), parameter_list);
+
+    if (!ranges_orig.empty() or !candidate_text_ranges.empty()) {
+        auto new_template =
+          _int_as_formatted_evidence<StrType>(candidate_text, candidate_text_ranges, TagMappingMode::Mapper);
+
+        py::list new_args;
+        py::dict new_kwargs;
+        for (const auto arg : args) {
+            auto str_arg = py::cast<py::str>(arg);
+            auto n_arg = _all_as_formatted_evidence<py::str>(str_arg, TagMappingMode::Mapper);
+            new_args.append(n_arg);
+        }
+        for (auto [key, value] : new_kwargs) {
+            auto str_value = py::cast<py::str>(value);
+            auto n_value = _all_as_formatted_evidence<py::str>(str_value, TagMappingMode::Mapper);
+            new_kwargs[key] = n_value;
+        }
+
+        StrType new_template_format =
+          py::getattr(new_template, "format")(*(py::cast<py::tuple>(new_args)), **new_kwargs);
+        std::tuple result = _convert_escaped_text_to_taint_text<StrType>(new_template_format, ranges_orig);
+        StrType result_text = get<0>(result);
+        TaintRangeRefs result_ranges = get<1>(result);
+        PyObject* new_result = new_pyobject_id(result_text.ptr());
+        set_ranges(new_result, result_ranges);
+        return py::reinterpret_steal<StrType>(new_result);
+    }
+    return py::getattr(candidate_text, "format")(*args, **kwargs);
+}
+
+void
+pyexport_format_aspect(py::module& m)
+{
+    m.def("_format_aspect", &api_format_aspect<py::str>, "candidate_text"_a, "parameter_list"_a);
+}

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectFormat.h

@@ -0,0 +1,16 @@
+#pragma once
+#include "Aspects/Helpers.h"
+#include "Initializer/Initializer.h"
+#include "TaintTracking/Source.h"
+#include "TaintTracking/TaintRange.h"
+#include "TaintTracking/TaintedObject.h"
+
+template<class StrType>
+StrType
+api_format_aspect(StrType& candidate_text,
+                  const py::tuple& parameter_list,
+                  const py::args& args,
+                  const py::kwargs& kwargs);
+
+void
+pyexport_format_aspect(py::module& m);

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectJoin.cpp

@@ -1,5 +1,62 @@
 #include "AspectJoin.h"
 
+PyObject*
+aspect_join_str(PyObject* sep,
+                PyObject* result,
+                PyObject* iterable_str,
+                size_t len_iterable,
+                TaintRangeMapType* tx_taint_map)
+{
+    // This is the special case for unicode str and unicode iterable_str.
+    // The iterable elements string will be split into 1 char-length strings.
+    const auto& to_iterable_str = get_tainted_object(iterable_str, tx_taint_map);
+    const auto& to_joiner = get_tainted_object(sep, tx_taint_map);
+
+    if (to_joiner == nullptr and to_iterable_str == nullptr) {
+        // No taints at all: return original result
+        return result;
+    }
+
+    const size_t& len_sep = PyUnicode_GET_LENGTH(sep);
+    const size_t& element_len = 1;
+    unsigned long current_pos{ 0L };
+    TaintedObjectPtr result_to;
+
+    if (len_sep == 0 and to_iterable_str) {
+        // Empty separator: the result is identical to iterable_str
+        result_to = initializer->allocate_tainted_object_copy(to_iterable_str);
+    } else if (len_sep > 0 and to_joiner and to_iterable_str == nullptr) {
+        // Iterable str is not tainted, only add n-times the joiner taints
+        result_to = initializer->allocate_tainted_object();
+        for (size_t i = 0; i < len_iterable - 1; i++) {
+            current_pos += element_len;
+            result_to->add_ranges_shifted(to_joiner, current_pos);
+            current_pos += len_sep;
+        }
+    } else {
+        // General case, iterable and joiner may be tainted
+        result_to = initializer->allocate_tainted_object();
+        for (size_t i = 0; i < len_iterable; i++) {
+            if (to_iterable_str) {
+                result_to->add_ranges_shifted(to_iterable_str, current_pos, element_len, i);
+            }
+
+            current_pos += element_len;
+            if (i < len_iterable - 1) {
+                if (to_joiner) {
+                    result_to->add_ranges_shifted(to_joiner, current_pos);
+                }
+                current_pos += len_sep;
+            }
+        }
+    }
+
+    PyObject* new_result{ new_pyobject_id(result) };
+    set_tainted_object(new_result, result_to, tx_taint_map);
+    Py_DecRef(result);
+    return new_result;
+}
+
 PyObject*
 aspect_join(PyObject* sep, PyObject* result, PyObject* iterable_elements, TaintRangeMapType* tx_taint_map)
 {
@@ -13,6 +70,12 @@ aspect_join(PyObject* sep, PyObject* result, PyObject* iterable_elements, TaintR
     } else if (PyTuple_Check(iterable_elements)) {
         len_iterable = PyTuple_Size(iterable_elements);
         GetElement = PyTuple_GetItem;
+    } else if (PyUnicode_Check(sep) and PyUnicode_Check(iterable_elements)) {
+        len_iterable = PyUnicode_GET_LENGTH(iterable_elements);
+        if (len_iterable)
+            return aspect_join_str(sep, result, iterable_elements, len_iterable, tx_taint_map);
+        else
+            return result; // Empty string is returned if empty iterable, so no tainted
     }
 
     unsigned long current_pos{ 0L };
@@ -69,9 +132,9 @@ aspect_join(PyObject* sep, PyObject* result, PyObject* iterable_elements, TaintR
         }
     }
 
-    PyObject* new_result{ new_pyobject_id(result, get_pyobject_size(result)) };
+    PyObject* new_result{ new_pyobject_id(result) };
     set_tainted_object(new_result, result_to, tx_taint_map);
-    Py_DECREF(result);
+    Py_DecRef(result);
     return new_result;
 }
 
@@ -97,7 +160,7 @@ api_join_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
                 PyList_Append(list_aux, item);
             }
             arg0 = list_aux;
-            Py_DECREF(iterator);
+            Py_DecRef(iterator);
             decref_arg0 = true;
         }
     }
@@ -118,7 +181,7 @@ api_join_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
     if (get_pyobject_size(result) == 0) {
         // Empty result cannot have taint ranges
         if (decref_arg0) {
-            Py_DECREF(arg0);
+            Py_DecRef(arg0);
         }
         return result;
     }
@@ -129,7 +192,7 @@ api_join_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
     }
     auto res = aspect_join(sep, result, arg0, ctx_map);
     if (decref_arg0) {
-        Py_DECREF(arg0);
+        Py_DecRef(arg0);
     }
     return res;
 }

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectJoin.h

@@ -1,11 +1,7 @@
 #pragma once
-#include "Aspects/Helpers.h"
 #include "Initializer/Initializer.h"
-#include "TaintTracking//TaintedObject.h"
 #include "TaintTracking/TaintRange.h"
-#include "TaintedOps/TaintedOps.h"
-#include <Python.h>
-#include <pybind11/pybind11.h>
+#include "TaintTracking/TaintedObject.h"
 
 namespace py = pybind11;
 

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectOperatorAdd.cpp

@@ -1,11 +1,20 @@
 #include "AspectOperatorAdd.h"
 
+/**
+ * This function updates result_o object with taint information of candidate_text and/or text_to_add
+ *
+ * @param result_o The result object to which the aspect will be added.
+ * @param candidate_text The candidate text object to which the aspect will be added.
+ * @param text_to_add The text aspect to be added.
+ * @param tx_taint_map The taint range map that stores taint information.
+ *
+ * @return A new result object with the taint information.
+ */
 PyObject*
 add_aspect(PyObject* result_o, PyObject* candidate_text, PyObject* text_to_add, TaintRangeMapType* tx_taint_map)
 {
     size_t len_candidate_text{ get_pyobject_size(candidate_text) };
     size_t len_text_to_add{ get_pyobject_size(text_to_add) };
-    size_t len_result_o{ get_pyobject_size(result_o) };
 
     if (len_text_to_add == 0 and len_candidate_text > 0) {
         return candidate_text;
@@ -16,8 +25,8 @@ add_aspect(PyObject* result_o, PyObject* candidate_text, PyObject* text_to_add,
 
     const auto& to_candidate_text = get_tainted_object(candidate_text, tx_taint_map);
     if (to_candidate_text and to_candidate_text->get_ranges().size() >= TaintedObject::TAINT_RANGE_LIMIT) {
-        const auto& res_new_id = new_pyobject_id(result_o, len_result_o);
-        Py_DECREF(result_o);
+        const auto& res_new_id = new_pyobject_id(result_o);
+        Py_DecRef(result_o);
         // If left side is already at the maximum taint ranges, we just reuse its
         // ranges, we don't need to look at left side.
         set_tainted_object(res_new_id, to_candidate_text, tx_taint_map);
@@ -29,21 +38,35 @@ add_aspect(PyObject* result_o, PyObject* candidate_text, PyObject* text_to_add,
         return result_o;
     }
     if (!to_text_to_add) {
-        const auto& res_new_id = new_pyobject_id(result_o, len_result_o);
-        Py_DECREF(result_o);
+        const auto& res_new_id = new_pyobject_id(result_o);
+        Py_DecRef(result_o);
         set_tainted_object(res_new_id, to_candidate_text, tx_taint_map);
         return res_new_id;
     }
 
-    auto tainted = initializer->allocate_tainted_object(to_candidate_text);
-    tainted->add_ranges_shifted(to_text_to_add, (long)len_candidate_text);
-    const auto& res_new_id = new_pyobject_id(result_o, len_result_o);
-    Py_DECREF(result_o);
+    auto tainted = initializer->allocate_tainted_object_copy(to_candidate_text);
+    tainted->add_ranges_shifted(to_text_to_add, (RANGE_START)len_candidate_text);
+    const auto res_new_id = new_pyobject_id(result_o);
+    Py_DecRef(result_o);
     set_tainted_object(res_new_id, tainted, tx_taint_map);
 
     return res_new_id;
 }
 
+/**
+ * Adds aspect, override all python Add operations.
+ *
+ * The AST Visitor (ddtrace/appsec/_iast/_ast/visitor.py) replaces all "+" operations in Python code with this function.
+ * This function takes 2 arguments. If the operation is 'a = b + c', this function should be 'a = api_add_aspect(b, c)'.
+ * This function connects Python with the C++ function 'add_aspect'.
+ *
+ * @param self The Python extension module.
+ * @param args An array of Python objects containing the candidate text and text aspect.
+ * @param nargs The number of arguments in the 'args' array.
+ *
+ * @return A new Python object representing the result of adding the aspect to the candidate text, considering taint
+ * information.
+ */
 PyObject*
 api_add_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
 {

ddtrace/appsec/iast/_taint_tracking/Aspects/AspectOperatorAdd.h

@@ -1,13 +1,7 @@
-#ifndef _NATIVE_ASPECTOPERATORADD_H
-#define _NATIVE_ASPECTOPERATORADD_H
-#include "Aspects/Helpers.h"
+#pragma once
 #include "Initializer/Initializer.h"
 #include "TaintTracking/TaintRange.h"
 #include "TaintTracking/TaintedObject.h"
-#include "TaintedOps/TaintedOps.h"
-#include <pybind11/pybind11.h>
 
 PyObject*
 api_add_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs);
-
-#endif //_NATIVE_ASPECTOPERATORADD_H