chore(asm): avoid double calls on some aspect front-ends (#9344)

## Description

Some aspects (index, slice, ljust, format, encode, decode and replace)
were calling the original function and then potentially the native
aspect, which would call again the original function. This reworks these
aspect front-ends so only one call is ever done, which should improve
performance a little.

This refactor also arised and error where we were not checking the
CPython error with the new `has_pyerror()` helper in the native
`api_aspect_index` so it fixed that too.

## 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] Risks are described (performance impact, potential for breakage,
maintainability)
- [X] Change is maintainable (easy to change, telemetry, documentation)
- [X] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed or label `changelog/no-changelog` is set
- [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))
- [X] If this PR changes the public interface, I've notified
`@DataDog/apm-tees`.

## Reviewer Checklist

- [x] Title is accurate
- [x] All changes are related to the pull request's stated goal
- [x] Description motivates each change
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- [x] Testing strategy adequately addresses listed risks
- [x] Change is maintainable (easy to change, telemetry, documentation)
- [x] Release note makes sense to a user of the library
- [x] Author has 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)

---------

Signed-off-by: Juanjo Alvarez <[email protected]>
Co-authored-by: Alberto Vara <[email protected]>

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

@@ -1,4 +1,5 @@
 #include "AspectIndex.h"
+#include "Helpers.h"
 
 /**
  * @brief Index aspect
@@ -49,6 +50,9 @@ api_index_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
     PyObject* idx = args[1];
 
     PyObject* result_o = PyObject_GetItem(candidate_text, idx);
+    if (has_pyerr()) {
+        return nullptr;
+    }
 
     if (not ctx_map or ctx_map->empty()) {
         return result_o;

ddtrace/appsec/_iast/_taint_tracking/aspects.py

@@ -210,16 +210,13 @@ def join_aspect(orig_function: Optional[Callable], flag_added_args: int, *args:
 
 
 def index_aspect(candidate_text: Text, index: int) -> Text:
-    result = candidate_text[index]
+    if isinstance(candidate_text, IAST.TEXT_TYPES) and isinstance(index, int):
+        try:
+            return _index_aspect(candidate_text, index)
+        except Exception as e:
+            iast_taint_log_error("IAST propagation error. index_aspect. {}".format(e))
 
-    if not isinstance(candidate_text, IAST.TEXT_TYPES) or not isinstance(index, int):
-        return result
-
-    try:
-        return _index_aspect(candidate_text, index)
-    except Exception as e:
-        iast_taint_log_error("IAST propagation error. index_aspect. {}".format(e))
-    return result
+    return candidate_text[index]
 
 
 def slice_aspect(candidate_text: Text, start: int, stop: int, step: int) -> Text:
@@ -230,15 +227,12 @@ def slice_aspect(candidate_text: Text, start: int, stop: int, step: int) -> Text
         or (step is not None and not isinstance(step, int))
     ):
         return candidate_text[start:stop:step]
-    result = candidate_text[start:stop:step]
+
     try:
-        new_result = _slice_aspect(candidate_text, start, stop, step)
-        if new_result != result:
-            raise Exception("Propagation result %r is different to candidate_text[slice] %r" % (new_result, result))
-        return new_result
+        return _slice_aspect(candidate_text, start, stop, step)
     except Exception as e:
         iast_taint_log_error("IAST propagation error. slice_aspect. {}".format(e))
-    return result
+    return candidate_text[start:stop:step]
 
 
 def bytearray_extend_aspect(orig_function: Optional[Callable], flag_added_args: int, *args: Any, **kwargs: Any) -> Any:
@@ -319,29 +313,25 @@ def ljust_aspect(
     candidate_text = args[0]
     args = args[flag_added_args:]
 
-    result = candidate_text.ljust(*args, **kwargs)
-
-    if not isinstance(candidate_text, IAST.TEXT_TYPES):
-        return result
-
-    try:
-        ranges_new = get_ranges(candidate_text)
-        fillchar = parse_params(1, "fillchar", " ", *args, **kwargs)
-        fillchar_ranges = get_ranges(fillchar)
-        if ranges_new is None or (not ranges_new and not fillchar_ranges):
+    if isinstance(candidate_text, IAST.TEXT_TYPES):
+        try:
+            ranges_new = get_ranges(candidate_text)
+            fillchar = parse_params(1, "fillchar", " ", *args, **kwargs)
+            fillchar_ranges = get_ranges(fillchar)
+            if ranges_new is None or (not ranges_new and not fillchar_ranges):
+                return candidate_text.ljust(*args, **kwargs)
+
+            if fillchar_ranges:
+                # Can only be one char, so we create one range to cover from the start to the end
+                ranges_new = ranges_new + [shift_taint_range(fillchar_ranges[0], len(candidate_text))]
+
+            result = candidate_text.ljust(parse_params(0, "width", None, *args, **kwargs), fillchar)
+            taint_pyobject_with_ranges(result, ranges_new)
             return result
+        except Exception as e:
+            iast_taint_log_error("IAST propagation error. ljust_aspect. {}".format(e))
 
-        if fillchar_ranges:
-            # Can only be one char, so we create one range to cover from the start to the end
-            ranges_new = ranges_new + [shift_taint_range(fillchar_ranges[0], len(candidate_text))]
-
-        new_result = candidate_text.ljust(parse_params(0, "width", None, *args, **kwargs), fillchar)
-        taint_pyobject_with_ranges(new_result, ranges_new)
-        return new_result
-    except Exception as e:
-        iast_taint_log_error("IAST propagation error. ljust_aspect. {}".format(e))
-
-    return result
+    return candidate_text.ljust(*args, **kwargs)
 
 
 def zfill_aspect(
@@ -410,16 +400,14 @@ def format_aspect(
     if not isinstance(candidate_text, IAST.TEXT_TYPES):
         return result
 
-    try:
-        params = tuple(args) + tuple(kwargs.values())
-        new_result = _format_aspect(candidate_text, params, *args, **kwargs)
-        if new_result != result:
-            raise Exception("Propagation result %r is different to candidate_text.format %r" % (new_result, result))
-        return new_result
-    except Exception as e:
-        iast_taint_log_error("IAST propagation error. format_aspect. {}".format(e))
+    if isinstance(candidate_text, IAST.TEXT_TYPES):
+        try:
+            params = tuple(args) + tuple(kwargs.values())
+            return _format_aspect(candidate_text, params, *args, **kwargs)
+        except Exception as e:
+            iast_taint_log_error("IAST propagation error. format_aspect. {}".format(e))
 
-    return result
+    return candidate_text.format(*args, **kwargs)
 
 
 def format_map_aspect(
@@ -684,18 +672,15 @@ def decode_aspect(
     self = args[0]
     args = args[(flag_added_args or 1) :]
     # Assume we call decode method of the first argument
-    result = self.decode(*args, **kwargs)
 
-    if not is_pyobject_tainted(self) or not isinstance(self, bytes):
-        return result
-
-    try:
-        codec = args[0] if args else "utf-8"
-        inc_dec = codecs.getincrementaldecoder(codec)(**kwargs)
-        return incremental_translation(self, inc_dec, inc_dec.decode, "")
-    except Exception as e:
-        iast_taint_log_error("IAST propagation error. decode_aspect. {}".format(e))
-    return result
+    if is_pyobject_tainted(self) and isinstance(self, bytes):
+        try:
+            codec = args[0] if args else "utf-8"
+            inc_dec = codecs.getincrementaldecoder(codec)(**kwargs)
+            return incremental_translation(self, inc_dec, inc_dec.decode, "")
+        except Exception as e:
+            iast_taint_log_error("IAST propagation error. decode_aspect. {}".format(e))
+    return self.decode(*args, **kwargs)
 
 
 def encode_aspect(
@@ -708,18 +693,15 @@ def encode_aspect(
 
     self = args[0]
     args = args[(flag_added_args or 1) :]
-    # Assume we call encode method of the first argument
-    result = self.encode(*args, **kwargs)
-
-    if not is_pyobject_tainted(self) or not isinstance(self, str):
-        return result
 
-    try:
-        codec = args[0] if args else "utf-8"
-        inc_enc = codecs.getincrementalencoder(codec)(**kwargs)
-        return incremental_translation(self, inc_enc, inc_enc.encode, b"")
-    except Exception as e:
-        iast_taint_log_error("IAST propagation error. encode_aspect. {}".format(e))
+    if is_pyobject_tainted(self) and isinstance(self, str):
+        try:
+            codec = args[0] if args else "utf-8"
+            inc_enc = codecs.getincrementalencoder(codec)(**kwargs)
+            return incremental_translation(self, inc_enc, inc_enc.encode, b"")
+        except Exception as e:
+            iast_taint_log_error("IAST propagation error. encode_aspect. {}".format(e))
+    result = self.encode(*args, **kwargs)
     return result
 
 

tests/appsec/iast/aspects/test_encode_decode_aspect.py

@@ -131,7 +131,7 @@ def test_encode_and_add_aspect(infix, args, kwargs, should_be_tainted, prefix, s
         assert list_ranges[0].length == len_infix
 
 
-def test_encode_error_and_no_log_metric(telemetry_writer):
+def test_encode_error_with_tainted_gives_one_log_metric(telemetry_writer):
     string_input = taint_pyobject(
         pyobject="abcde",
         source_name="test_add_aspect_tainting_left_hand",
@@ -141,11 +141,22 @@ def test_encode_error_and_no_log_metric(telemetry_writer):
     with pytest.raises(LookupError):
         mod.do_encode(string_input, "encoding-not-exists")
 
+    list_metrics_logs = list(telemetry_writer._logs)
+    assert len(list_metrics_logs) == 1
+    assert "message" in list_metrics_logs[0]
+    assert "IAST propagation error. encode_aspect. " in list_metrics_logs[0]["message"]
+
+
+def test_encode_error_with_not_tainted_gives_no_log_metric(telemetry_writer):
+    string_input = "abcde"
+    with pytest.raises(LookupError):
+        mod.do_encode(string_input, "encoding-not-exists")
+
     list_metrics_logs = list(telemetry_writer._logs)
     assert len(list_metrics_logs) == 0
 
 
-def test_dencode_error_and_no_log_metric(telemetry_writer):
+def test_dencode_error_with_tainted_gives_one_log_metric(telemetry_writer):
     string_input = taint_pyobject(
         pyobject=b"abcde",
         source_name="test_add_aspect_tainting_left_hand",
@@ -155,5 +166,16 @@ def test_dencode_error_and_no_log_metric(telemetry_writer):
     with pytest.raises(LookupError):
         mod.do_decode(string_input, "decoding-not-exists")
 
+    list_metrics_logs = list(telemetry_writer._logs)
+    assert len(list_metrics_logs) == 1
+    assert "message" in list_metrics_logs[0]
+    assert "IAST propagation error. decode_aspect. " in list_metrics_logs[0]["message"]
+
+
+def test_dencode_error_with_not_tainted_gives_no_log_metric(telemetry_writer):
+    string_input = b"abcde"
+    with pytest.raises(LookupError):
+        mod.do_decode(string_input, "decoding-not-exists")
+
     list_metrics_logs = list(telemetry_writer._logs)
     assert len(list_metrics_logs) == 0

tests/appsec/iast/aspects/test_index_aspect_fixtures.py

@@ -64,7 +64,7 @@ def test_string_index(input_str, index_pos, expected_result, tainted):
 
 
 @pytest.mark.skipif(sys.version_info < (3, 9, 0), reason="Python version not supported by IAST")
-def test_index_error_and_no_log_metric(telemetry_writer):
+def test_index_error_with_tainted_gives_one_log_metric(telemetry_writer):
     string_input = taint_pyobject(
         pyobject="abcde",
         source_name="test_add_aspect_tainting_left_hand",
@@ -75,7 +75,8 @@ def test_index_error_and_no_log_metric(telemetry_writer):
         mod.do_index(string_input, 100)
 
     list_metrics_logs = list(telemetry_writer._logs)
-    assert len(list_metrics_logs) == 0
+    assert len(list_metrics_logs) == 1
+    assert "IAST propagation error. index_aspect" in list_metrics_logs[0]["message"]
 
 
 @pytest.mark.skip_iast_check_logs

tests/appsec/iast/aspects/test_str_aspect.py

@@ -527,13 +527,13 @@ def test_aspect_ljust_str_tainted(self):
         ljusted = mod.do_ljust(string_input, 4)  # pylint: disable=no-member
         assert as_formatted_evidence(ljusted) == ":+-foo-+: "
 
-    def test_aspect_ljust_error_and_no_log_metric(self, telemetry_writer):
+    def test_aspect_ljust_error_with_tainted_gives_one_log_metric(self, telemetry_writer):
         string_input = create_taint_range_with_format(":+-foo-+:")
         with pytest.raises(TypeError):
             mod.do_ljust(string_input, "aaaaa")
 
         list_metrics_logs = list(telemetry_writer._logs)
-        assert len(list_metrics_logs) == 0
+        assert len(list_metrics_logs) == 1
 
     def test_zfill(self):
         # Not tainted
@@ -572,4 +572,4 @@ def test_format(self):
         result = mod.do_format_fill(string_input)  # pylint: disable=no-member
         # TODO format with params doesn't work correctly the assert should be
         #  assert as_formatted_evidence(result) == ":+-foo       -+:"
-        assert as_formatted_evidence(result) == "foo       "
+        assert as_formatted_evidence(result) == ":+-foo-+:"