chore(iast): optimize modulo aspect more (3/3)

ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectIndex.cpp

@@ -52,7 +52,7 @@ api_index_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
     if (!is_text(candidate_text) or !is_some_number(idx)) {
         return result_o;
     }
-    TRY_CATCH_ASPECT("index_aspect", , {
+    TRY_CATCH_ASPECT("index_aspect", return result_o, , {
         const auto ctx_map = Initializer::get_tainting_map();
         if (not ctx_map or ctx_map->empty()) {
             return result_o;

ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectModulo.cpp

@@ -1,41 +1,27 @@
 #include "AspectModulo.h"
 #include "Helpers.h"
 
-py::object
-api_modulo_aspect_pyobject(py::object candidate_text, py::object candidate_tuple)
-{
-    return candidate_text.attr("__mod__")(candidate_tuple);
-}
-
 static PyObject*
 do_modulo(PyObject* text, PyObject* insert_tuple_or_obj)
 {
     PyObject* result = nullptr;
 
     // If the second argument is not a tuple and not a mapping, wrap it in a tuple
     PyObject* insert_tuple = nullptr;
-    if (PyMapping_Check(insert_tuple_or_obj)) {
+    if (PyMapping_Check(insert_tuple_or_obj) or PyTuple_Check(insert_tuple_or_obj)) {
         insert_tuple = insert_tuple_or_obj;
-        Py_INCREF(insert_tuple); // Increment the reference count since we'll be using this directly
-    } else if (PyTuple_Check(insert_tuple_or_obj)) {
-        insert_tuple = insert_tuple_or_obj;
-        Py_INCREF(insert_tuple); // Increment the reference count since we'll be using this directly
+        Py_INCREF(insert_tuple);
     } else {
         insert_tuple = PyTuple_Pack(1, insert_tuple_or_obj);
         if (insert_tuple == nullptr) {
-            return nullptr; // Return if PyTuple_Pack fails
+            return nullptr;
         }
     }
 
-    // Check if text is a Unicode object
     if (PyUnicode_Check(text)) {
         result = PyUnicode_Format(text, insert_tuple);
-    }
-    // Check if text is a bytes object
-    else if (PyBytes_Check(text)) {
-        // Convert bytes to str, format, and convert back to bytes
-        PyObject* text_unicode = PyUnicode_FromEncodedObject(text, "utf-8", "strict");
-        if (text_unicode != nullptr) {
+    } else if (PyBytes_Check(text)) {
+        if (PyObject* text_unicode = PyUnicode_FromEncodedObject(text, "utf-8", "strict"); text_unicode != nullptr) {
             result = PyUnicode_Format(text_unicode, insert_tuple);
             Py_DECREF(text_unicode);
 
@@ -45,11 +31,9 @@ do_modulo(PyObject* text, PyObject* insert_tuple_or_obj)
                 result = encoded_result;
             }
         }
-    }
-    // Check if text is a bytearray object
-    else if (PyByteArray_Check(text)) {
-        PyObject* text_bytes = PyBytes_FromStringAndSize(PyByteArray_AsString(text), PyByteArray_Size(text));
-        if (text_bytes != nullptr) {
+    } else if (PyByteArray_Check(text)) {
+        if (PyObject* text_bytes = PyBytes_FromStringAndSize(PyByteArray_AsString(text), PyByteArray_Size(text));
+            text_bytes != nullptr) {
             PyObject* text_unicode = PyUnicode_FromEncodedObject(text_bytes, "utf-8", "strict");
             Py_DECREF(text_bytes);
             if (text_unicode != nullptr) {
@@ -69,9 +53,7 @@ do_modulo(PyObject* text, PyObject* insert_tuple_or_obj)
             Py_DECREF(result);
             result = result_bytearray;
         }
-    }
-    // If text is not one of the expected types, raise an error
-    else {
+    } else {
         Py_DECREF(insert_tuple);
         return nullptr;
     }
@@ -80,84 +62,89 @@ do_modulo(PyObject* text, PyObject* insert_tuple_or_obj)
     return result;
 }
 
-template<class StrType>
-StrType
-api_modulo_aspect(StrType candidate_text, py::object candidate_tuple)
+PyObject*
+api_modulo_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
 {
-    if (not is_text(candidate_text.ptr())) {
-        return candidate_text.attr("__mod__")(candidate_tuple);
-    }
-
-    PyObject* pyo_result_o = do_modulo(candidate_text.ptr(), candidate_tuple.ptr());
-    if (pyo_result_o == nullptr) {
-        return candidate_text.attr("__mod__")(candidate_tuple);
+    if (nargs != 2) {
+        py::set_error(PyExc_ValueError, MSG_ERROR_N_PARAMS);
+        return nullptr;
     }
+    PyObject* candidate_text = args[0];
+    PyObject* candidate_tuple = args[1];
+
+    const auto py_candidate_text = py::reinterpret_borrow<py::object>(candidate_text);
+    auto py_candidate_tuple = py::reinterpret_borrow<py::object>(candidate_tuple);
+
+    // Lambda to get the result of the modulo operation
+    auto get_result = [&]() -> PyObject* {
+        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(), , {
+        const auto py_str_type = get_pytext_type(args[0]);
+        if (py_str_type == PyTextType::OTHER) {
+            try {
+                return get_result();
+            } catch (py::error_already_set& e) {
+                e.restore();
+                return nullptr;
+            }
+        }
 
-    StrType result_o = py::reinterpret_steal<StrType>(pyo_result_o);
-    const py::tuple parameters =
-      py::isinstance<py::tuple>(candidate_tuple) ? candidate_tuple : py::make_tuple(candidate_tuple);
+        const py::tuple parameters =
+          py::isinstance<py::tuple>(py_candidate_tuple) ? py_candidate_tuple : py::make_tuple(py_candidate_tuple);
 
-    const auto tx_map = Initializer::get_tainting_map();
-    if (!tx_map || tx_map->empty()) {
-        return result_o;
-    }
+        const auto tx_map = Initializer::get_tainting_map();
+        if (!tx_map || tx_map->empty()) {
+            return get_result();
+        }
 
-    TRY_CATCH_ASPECT("modulo_aspect", , {
-        auto [ranges_orig, candidate_text_ranges] = are_all_text_all_ranges(candidate_text.ptr(), parameters);
+        auto [ranges_orig, candidate_text_ranges] = are_all_text_all_ranges(candidate_text, parameters);
 
         if (ranges_orig.empty()) {
-            return result_o;
+            return get_result();
         }
 
-        StrType fmttext = as_formatted_evidence(candidate_text, candidate_text_ranges, TagMappingMode::Mapper);
+        auto std_candidate_text = py_candidate_text.cast<string>();
+        auto fmttext = as_formatted_evidence(std_candidate_text, candidate_text_ranges, TagMappingMode::Mapper);
         py::list list_formatted_parameters;
 
         for (const py::handle& param_handle : parameters) {
-            auto param_strtype = py::reinterpret_borrow<py::object>(param_handle).cast<StrType>();
             if (is_text(param_handle.ptr())) {
                 auto [ranges, ranges_error] = get_ranges(param_handle.ptr(), tx_map);
-                StrType n_parameter = as_formatted_evidence(param_strtype, ranges, TagMappingMode::Mapper, nullopt);
-                list_formatted_parameters.append(n_parameter);
+                string n_parameter =
+                  as_formatted_evidence(AnyTextObjectToString(param_handle), ranges, TagMappingMode::Mapper, nullopt);
+                list_formatted_parameters.append(StringToPyObject(n_parameter, py_str_type));
             } else {
                 list_formatted_parameters.append(param_handle);
             }
         }
         py::tuple formatted_parameters(list_formatted_parameters);
 
-        PyObject* pyo_applied_params = do_modulo(fmttext.ptr(), formatted_parameters.ptr());
-
-        StrType applied_params;
-        if (pyo_applied_params == nullptr) {
-            return result_o;
-        } else {
-            applied_params = py::reinterpret_steal<StrType>(pyo_applied_params);
+        PyObject* applied_params = do_modulo(StringToPyObject(fmttext, py_str_type).ptr(), formatted_parameters.ptr());
+        if (applied_params == nullptr) {
+            return get_result();
         }
 
-        return api_convert_escaped_text_to_taint_text(applied_params, ranges_orig);
+        auto res_pyobject = api_convert_escaped_text_to_taint_text(applied_params, ranges_orig, py_str_type);
+        Py_DECREF(applied_params);
+        if (res_pyobject == nullptr) {
+            return get_result();
+        }
+        return res_pyobject;
     });
 }
-
-void
-pyexport_aspect_modulo(py::module& m)
-{
-    m.def("_aspect_modulo",
-          &api_modulo_aspect<py::str>,
-          "candidate_text"_a,
-          "candidate_tuple"_a,
-          py::return_value_policy::move);
-    m.def("_aspect_modulo",
-          &api_modulo_aspect<py::bytes>,
-          "candidate_text"_a,
-          "candidate_tuple"_a,
-          py::return_value_policy::move);
-    m.def("_aspect_modulo",
-          &api_modulo_aspect<py::bytearray>,
-          "candidate_text"_a,
-          "candidate_tuple"_a,
-          py::return_value_policy::move);
-    m.def("_aspect_modulo",
-          &api_modulo_aspect_pyobject,
-          "candidate_text"_a,
-          "candidate_tuple"_a,
-          py::return_value_policy::move);
-}

ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectModulo.h

@@ -3,9 +3,5 @@
 
 namespace py = pybind11;
 
-template<class StrType>
-StrType
-api_modulo_aspect(StrType candidate_text, py::object candidate_tuple);
-
-void
-pyexport_aspect_modulo(py::module& m);
+PyObject*
+api_modulo_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs);

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

@@ -85,7 +85,7 @@ api_add_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
     // PyNumber_Add actually works for any type!
     result_o = PyNumber_Add(candidate_text, text_to_add);
 
-    TRY_CATCH_ASPECT("add_aspect", , {
+    TRY_CATCH_ASPECT("add_aspect", return result_o, , {
         const auto tx_map = Initializer::get_tainting_map();
         if (not tx_map or tx_map->empty()) {
             return result_o;
@@ -119,7 +119,7 @@ api_add_inplace_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
 
     result_o = PyNumber_InPlaceAdd(candidate_text, text_to_add);
 
-    TRY_CATCH_ASPECT("add_inplace_aspect", , {
+    TRY_CATCH_ASPECT("add_inplace_aspect", return result_o, , {
         const auto tx_map = Initializer::get_tainting_map();
         if (not tx_map or tx_map->empty()) {
             return result_o;

ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectSlice.cpp

@@ -145,7 +145,7 @@ api_slice_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
 
     PyObject* result_o = PyObject_GetItem(candidate_text, slice);
 
-    TRY_CATCH_ASPECT("slice_aspect", Py_XDECREF(slice), {
+    TRY_CATCH_ASPECT("slice_aspect", return result_o, Py_XDECREF(slice), {
         // If no result or the params are not None|Number or the result is the same as the candidate text, nothing
         // to taint
         if (result_o == nullptr or (!is_text(candidate_text)) or (start != Py_None and !PyLong_Check(start)) or

ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectSplit.cpp

@@ -62,7 +62,7 @@ split_text_common(const py::object& orig_function,
         return result_o;
     }
 
-    TRY_CATCH_ASPECT("split_aspect", , {
+    TRY_CATCH_ASPECT("split_aspect", return result_o, , {
         if (split_func == "split") {
             if (auto re_split_result = handle_potential_re_split(args_tuple, sliced_args, kwargs, tx_map);
                 re_split_result.has_value()) {
@@ -103,7 +103,7 @@ api_splitlines_text(const py::object& orig_function,
         return result_o;
     }
 
-    TRY_CATCH_ASPECT("split_aspect", , {
+    TRY_CATCH_ASPECT("split_aspect", return result_o, , {
         auto [ranges, ranges_error] = get_ranges(text.ptr(), tx_map);
         if (ranges_error || ranges.empty()) {
             return result_o;

ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectsOsPath.cpp

@@ -23,7 +23,7 @@ api_ospathjoin_aspect(StrType& first_part, const py::args& args)
         return result_o;
     }
 
-    TRY_CATCH_ASPECT("ospathjoin_aspect", , {
+    TRY_CATCH_ASPECT("ospathjoin_aspect", return result_o, , {
         const auto separator = ospath.attr("sep").cast<std::string>();
         const auto sepsize = separator.size();
 
@@ -105,7 +105,7 @@ api_ospathbasename_aspect(const StrType& path)
     auto basename = ospath.attr("basename");
     auto result_o = basename(path);
 
-    TRY_CATCH_ASPECT("ospathbasename_aspect", , {
+    TRY_CATCH_ASPECT("ospathbasename_aspect", return result_o, , {
         const auto tx_map = Initializer::get_tainting_map();
         if (not tx_map or tx_map->empty() or py::len(result_o) == 0) {
             return result_o;
@@ -138,7 +138,7 @@ api_ospathdirname_aspect(const StrType& path)
     auto dirname = ospath.attr("dirname");
     auto result_o = dirname(path);
 
-    TRY_CATCH_ASPECT("ospathdirname_aspect", , {
+    TRY_CATCH_ASPECT("ospathdirname_aspect", return result_o, , {
         const auto tx_map = Initializer::get_tainting_map();
         if (not tx_map or tx_map->empty() or py::len(result_o) == 0) {
             return result_o;
@@ -171,7 +171,7 @@ forward_to_set_ranges_on_splitted(const char* function_name, const StrType& path
     auto function = ospath.attr(function_name);
     auto result_o = function(path);
 
-    TRY_CATCH_ASPECT("forward_to_set_ranges_on_splitted", , {
+    TRY_CATCH_ASPECT("forward_to_set_ranges_on_splitted", return result_o, , {
         const auto tx_map = Initializer::get_tainting_map();
         if (not tx_map or tx_map->empty() or py::len(result_o) == 0) {
             return result_o;
@@ -223,7 +223,7 @@ api_ospathnormcase_aspect(const StrType& path)
     auto normcase = ospath.attr("normcase");
     auto result_o = normcase(path);
 
-    TRY_CATCH_ASPECT("ospathnormcase_aspect", , {
+    TRY_CATCH_ASPECT("ospathnormcase_aspect", return result_o, , {
         const auto tx_map = Initializer::get_tainting_map();
         if (not tx_map or tx_map->empty()) {
             return result_o;