From 94a070172254cf99bb8b1b8c93af138b2ce69484 Mon Sep 17 00:00:00 2001
From: Alberto Vara <alberto.vara@datadoghq.com>
Date: Mon, 23 Dec 2024 08:19:39 +0100
Subject: [PATCH] chore(iast): memory leak in aspect modulo (#11787)

a memory leak was introduced in #11601
This PR reverts it.

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- 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)

(cherry picked from commit 58ef74f405f3e9e4c24f31cccd6d6330e14d3c97)
---
 .../_taint_tracking/Aspects/AspectModulo.cpp  | 43 +++++++++++--------
 1 file changed, 25 insertions(+), 18 deletions(-)

diff --git a/ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectModulo.cpp b/ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectModulo.cpp
index b7454de26f8..a08f76d9f3d 100644
--- a/ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectModulo.cpp
+++ b/ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectModulo.cpp
@@ -2,7 +2,7 @@
 #include "Helpers.h"
 
 static PyObject*
-do_modulo(PyObject* text, PyObject* insert_tuple_or_obj, py::object py_candidate_text, py::object py_candidate_tuple)
+do_modulo(PyObject* text, PyObject* insert_tuple_or_obj)
 {
     PyObject* result = nullptr;
 
@@ -13,22 +13,18 @@ do_modulo(PyObject* text, PyObject* insert_tuple_or_obj, py::object py_candidate
         Py_INCREF(insert_tuple);
     } else {
         insert_tuple = PyTuple_Pack(1, insert_tuple_or_obj);
+        if (insert_tuple == nullptr) {
+            return nullptr;
+        }
     }
 
-    if (PyUnicode_Check(text) && insert_tuple != nullptr) {
+    if (PyUnicode_Check(text)) {
         result = PyUnicode_Format(text, insert_tuple);
+    } else if (PyBytes_Check(text) or PyByteArray_Check(text)) {
+        auto method_name = PyUnicode_FromString("__mod__");
+        result = PyObject_CallMethodObjArgs(text, method_name, insert_tuple, nullptr);
+        Py_DECREF(method_name);
     } else {
-        try {
-            py::object res_py = py_candidate_text.attr("__mod__")(py_candidate_tuple);
-            PyObject* res_pyo = res_py.ptr();
-            if (res_pyo != nullptr) {
-                Py_INCREF(res_pyo);
-            }
-            return res_pyo;
-        } catch (py::error_already_set& e) {
-            e.restore();
-            return nullptr;
-        }
     }
     Py_DECREF(insert_tuple);
     if (has_pyerr()) {
@@ -53,7 +49,21 @@ api_modulo_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
 
     // Lambda to get the result of the modulo operation
     auto get_result = [&]() -> PyObject* {
-        return do_modulo(candidate_text, candidate_tuple, py_candidate_text, py_candidate_tuple);
+        PyObject* res = do_modulo(candidate_text, candidate_tuple);
+        if (res == nullptr) {
+            try {
+                py::object res_py = py_candidate_text.attr("__mod__")(py_candidate_tuple);
+                PyObject* res_pyo = res_py.ptr();
+                if (res_pyo != nullptr) {
+                    Py_INCREF(res_pyo);
+                }
+                return res_pyo;
+            } catch (py::error_already_set& e) {
+                e.restore();
+                return nullptr;
+            }
+        }
+        return res;
     };
 
     TRY_CATCH_ASPECT("modulo_aspect", return get_result(), , {
@@ -97,10 +107,7 @@ api_modulo_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
         }
         py::tuple formatted_parameters(list_formatted_parameters);
 
-        PyObject* applied_params = do_modulo(StringToPyObject(fmttext, py_str_type).ptr(),
-                                             formatted_parameters.ptr(),
-                                             StringToPyObject(fmttext, py_str_type),
-                                             formatted_parameters);
+        PyObject* applied_params = do_modulo(StringToPyObject(fmttext, py_str_type).ptr(), formatted_parameters.ptr());
         if (applied_params == nullptr) {
             return get_result();
         }