chore(iast): inplace add aspect (#10311)

IAST: Aspect for inplace add (`+=` operator) function.
## 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)

benchmarks/bm/iast_fixtures/str_methods.py

@@ -53,6 +53,24 @@ def do_operator_add_params(a, b):
     return a + b
 
 
+def do_operator_add_inplace_params(a, b):
+    a += b
+    return a
+
+
+def do_operator_add_inplace_3_params(a, b, c):
+    a += b
+    a += c
+    return a
+
+
+def do_operator_add_inplace_3_times(a, b):
+    a += b
+    a += b
+    a += b
+    return a
+
+
 def do_string_assignment(a):
     b = a
     return b
@@ -971,15 +989,6 @@ def do_slice_negative(
     return s[-16:]
 
 
-def mult_two(a, b):  # type: (Any, Any) -> Any
-    return a * b
-
-
-def inplace_mult(a, b):
-    a *= b
-    return a
-
-
 class MyObject(object):
     def __init__(self, str_param):  # type: (str) -> None
         self.str_param = str_param
@@ -1219,3 +1228,17 @@ def do_customspec_formatspec():
 
 def do_fstring(a, b):
     return f"{a} + {b} = {a + b}"
+
+
+def _preprocess_lexer_input(text):
+    """Apply preprocessing such as decoding the input, removing BOM and normalizing newlines."""
+    # text now *is* a unicode string
+    text = text.replace("\r\n", "\n")
+    text = text.replace("\r", "\n")
+    text = text.strip()
+    text = text.strip("\n")
+
+    text = text.expandtabs(0)
+    text += "\n"
+
+    return text

ddtrace/appsec/_iast/_ast/visitor.py

@@ -87,6 +87,7 @@ def _mark_avoid_convert_recursively(node):
     },
     "operators": {
         ast.Add: "ddtrace_aspects.add_aspect",
+        "INPLACE_ADD": "ddtrace_aspects.add_inplace_aspect",
         "FORMAT_VALUE": "ddtrace_aspects.format_value_aspect",
         ast.Mod: "ddtrace_aspects.modulo_aspect",
         "BUILD_STRING": "ddtrace_aspects.build_string_aspect",
@@ -617,6 +618,51 @@ def visit_BinOp(self, call_node: ast.BinOp) -> Any:
 
         return call_node
 
+    def visit_AugAssign(self, augassign_node: ast.AugAssign) -> Any:
+        """
+        Replace an inplace add or multiply (+= / *=)
+        """
+        if isinstance(augassign_node.target, ast.Subscript):
+            # Can't augassign to function call, ignore this node
+            augassign_node.target.avoid_convert = True  # type: ignore[attr-defined]
+            self.generic_visit(augassign_node)
+            return augassign_node
+
+        self.generic_visit(augassign_node)
+        if augassign_node.op.__class__ == ast.Add:
+            # Optimization: ignore augassigns where the right side term is an integer since
+            # they can't apply to strings
+            if self._is_numeric_node(augassign_node.value):
+                return augassign_node
+
+            replacement_func = self._aspect_operators["INPLACE_ADD"]
+        else:
+            return augassign_node
+
+        # We must change the ctx of the target and value to Load() for using
+        # them as function arguments while keeping it as Store() for the
+        # Assign.targets, thus the manual copy
+
+        func_arg1 = copy.deepcopy(augassign_node.target)
+        func_arg1.ctx = ast.Load()  # type: ignore[attr-defined]
+        func_arg2 = copy.deepcopy(augassign_node.value)
+        func_arg2.ctx = ast.Load()  # type: ignore[attr-defined]
+
+        call_node = self._call_node(
+            augassign_node,
+            func=self._attr_node(augassign_node, replacement_func),
+            args=[func_arg1, func_arg2],
+        )
+
+        self.ast_modified = True
+        return self._node(
+            ast.Assign,
+            augassign_node,
+            targets=[augassign_node.target],
+            value=call_node,
+            type_comment=None,
+        )
+
     def visit_FormattedValue(self, fmt_value_node: ast.FormattedValue) -> Any:
         """
         Visit a FormattedValue node which are the constituent atoms for the
@@ -678,17 +724,6 @@ def visit_JoinedStr(self, joinedstr_node: ast.JoinedStr) -> Any:
         _set_metric_iast_instrumented_propagation()
         return call_node
 
-    def visit_AugAssign(self, augassign_node: ast.AugAssign) -> Any:
-        """Replace an inplace add or multiply."""
-        if isinstance(augassign_node.target, ast.Subscript):
-            # Can't augassign to function call, ignore this node
-            augassign_node.target.avoid_convert = True  # type: ignore[attr-defined]
-            self.generic_visit(augassign_node)
-            return augassign_node
-
-        # TODO: Replace an inplace add or multiply (+= / *=)
-        return augassign_node
-
     def visit_Assign(self, assign_node: ast.Assign) -> Any:
         """
         Add the ignore marks for left-side subscripts or list/tuples to avoid problems

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

@@ -49,11 +49,9 @@ add_aspect(PyObject* result_o,
 
     auto tainted = initializer->allocate_tainted_object_copy(to_candidate_text);
     tainted->add_ranges_shifted(to_text_to_add, static_cast<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);
+    set_tainted_object(result_o, tainted, tx_taint_map);
 
-    return res_new_id;
+    return result_o;
 }
 
 /**
@@ -115,4 +113,50 @@ api_add_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
         iast_taint_log_error(error_message);
         return result_o;
     }
+}
+
+PyObject*
+api_add_inplace_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
+{
+    PyObject* result_o = nullptr;
+
+    try {
+        if (nargs != 2) {
+            py::set_error(PyExc_ValueError, MSG_ERROR_N_PARAMS);
+            return nullptr;
+        }
+        PyObject* candidate_text = args[0];
+        PyObject* text_to_add = args[1];
+
+        result_o = PyNumber_InPlaceAdd(candidate_text, text_to_add);
+
+        const auto tx_map = initializer->get_tainting_map();
+        if (not tx_map or tx_map->empty()) {
+            return result_o;
+        }
+
+        if (not args_are_text_and_same_type(candidate_text, text_to_add)) {
+            return result_o;
+        }
+
+        // Quickly skip if both are noninterned-unicodes and not tainted
+        if (is_notinterned_notfasttainted_unicode(candidate_text) &&
+            is_notinterned_notfasttainted_unicode(text_to_add)) {
+            return result_o;
+        }
+        candidate_text = add_aspect(result_o, candidate_text, text_to_add, tx_map);
+        return candidate_text;
+    } catch (const py::error_already_set& e) {
+        const std::string error_message = "IAST propagation error in add_aspect. " + std::string(e.what());
+        iast_taint_log_error(error_message);
+        return result_o;
+    } catch (const std::exception& e) {
+        const std::string error_message = "IAST propagation error in add_aspect. " + std::string(e.what());
+        iast_taint_log_error(error_message);
+        return result_o;
+    } catch (...) {
+        const std::string error_message = "Unkown IAST propagation error in add_aspect. ";
+        iast_taint_log_error(error_message);
+        return result_o;
+    }
 }

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

@@ -3,3 +3,6 @@
 
 PyObject*
 api_add_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs);
+
+PyObject*
+api_add_inplace_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs);

ddtrace/appsec/_iast/_taint_tracking/_native.cpp

@@ -28,6 +28,7 @@ namespace py = pybind11;
 
 static PyMethodDef AspectsMethods[] = {
     { "add_aspect", ((PyCFunction)api_add_aspect), METH_FASTCALL, "aspect add" },
+    { "add_inplace_aspect", ((PyCFunction)api_add_inplace_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" },

ddtrace/appsec/_iast/_taint_tracking/aspects.py

@@ -52,13 +52,15 @@
 
 
 add_aspect = aspects.add_aspect
+add_inplace_aspect = aspects.add_inplace_aspect
 _extend_aspect = aspects.extend_aspect
 index_aspect = aspects.index_aspect
 _join_aspect = aspects.join_aspect
 _slice_aspect = aspects.slice_aspect
 
 __all__ = [
     "add_aspect",
+    "add_inplace_aspect",
     "str_aspect",
     "bytearray_extend_aspect",
     "decode_aspect",

tests/appsec/iast/_ast/test_ast_patching.py

@@ -82,6 +82,23 @@ def test_astpatch_module_changed_add_operator(module_name):
     assert "ddtrace_aspects.add_aspect(" in new_code
 
 
+@pytest.mark.parametrize(
+    "module_name",
+    [
+        ("tests.appsec.iast.fixtures.ast.add_operator.inplace"),
+    ],
+)
+def test_astpatch_module_changed_add_inplace_operator(module_name):
+    module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None]))
+    assert ("", "") != (module_path, new_source)
+    new_code = astunparse.unparse(new_source)
+    assert new_code.startswith(
+        "\nimport ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks"
+        "\nimport ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects"
+    )
+    assert "ddtrace_aspects.add_inplace_aspect(" in new_code
+
+
 @pytest.mark.parametrize(
     "module_name",
     [

tests/appsec/iast/aspects/test_add_aspect_fixtures.py

@@ -99,3 +99,15 @@ def test_string_operator_add_two_mixed_bytearray_bytes(self):  # type: () -> Non
         result = mod.do_operator_add_params(string_input, bar)
         assert result == bytearray(b"foobar")
         assert len(get_tainted_ranges(result)) == 0
+
+    def test_string_operator_add_bytearrays(self):  # type: () -> None
+        string_input = taint_pyobject(
+            pyobject=bytearray(b"foo"), source_name="foo", source_value="foo", source_origin=OriginType.PARAMETER
+        )
+        bar = taint_pyobject(
+            pyobject=bytearray(b"bar"), source_name="bar", source_value="bar", source_origin=OriginType.PARAMETER
+        )
+
+        result = mod.do_operator_add_params(string_input, bar)
+        assert result == bytearray(b"foobar")
+        assert get_tainted_ranges(result)