fix(iast): potential string id collisions [backport 2.1] (#7324)

Backport 77540c1 from #7318 to 2.1.

Fix potential string id collisions that could cause false positives with
non tainted objects being marked as tainted.

## 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.

Co-authored-by: Alberto Vara <[email protected]>

Author: github-actions[bot]

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

@@ -19,17 +19,21 @@ api_extend_aspect(PyObject* self, PyObject* const* args, Py_ssize_t nargs)
     if (!PyByteArray_Check(to_add) and !PyBytes_Check(to_add)) {
         return nullptr;
     }
+
+    auto ctx_map = initializer->get_tainting_map();
+    const auto& to_candidate = get_tainted_object(candidate_text, ctx_map);
+    auto to_result = initializer->allocate_tainted_object_copy(to_candidate);
+    const auto& to_toadd = get_tainted_object(to_add, ctx_map);
+
+    // 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);
-    auto ctx_map = initializer->get_tainting_map();
-    if (not ctx_map or ctx_map->empty()) {
+
+    if (not ctx_map or ctx_map->empty() or to_result == nullptr) {
         Py_RETURN_NONE;
     }
 
-    const auto& to_candidate = get_tainted_object(candidate_text, ctx_map);
-    auto to_result = initializer->allocate_tainted_object(to_candidate);
-    const auto& to_toadd = get_tainted_object(to_add, ctx_map);
     if (to_toadd) {
         to_result->add_ranges_shifted(to_toadd, (long)len_candidate_text);
     }

ddtrace/appsec/_iast/_taint_tracking/TaintTracking/TaintRange.h

@@ -33,9 +33,13 @@ using TaintRangeMapType = std::map<uintptr_t, TaintedObjectPtr>;
 #endif // NDEBUG
 
 inline static uintptr_t
-get_unique_id(const PyObject* str)
+get_unique_id(PyObject* str)
 {
-    return uintptr_t(str);
+    if ((((PyASCIIObject*)str)->hash) == -1) {
+        Py_hash_t result = PyObject_Hash(str);
+    }
+    auto res = uintptr_t(str) + ((PyASCIIObject*)str)->hash;
+    return hash<uintptr_t>{}(res);
 }
 
 struct TaintRange

releasenotes/notes/iast-fix-id-collisions-87da1c4229876398.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    Vulnerability Management for Code-level (IAST): Fix potential string id collisions that could cause false 
+    positives with non tainted objects being marked as tainted.

tests/appsec/iast/aspects/test_bytearray_extend_aspect.py

@@ -86,7 +86,7 @@ def test_first_and_second_tainted(self):
             pyobject=bytearray(b"123"), source_name="test1", source_value="foo", source_origin=OriginType.PARAMETER
         )
         ba2 = taint_pyobject(
-            pyobject=bytearray(b"456"), source_name="test2", source_value="foo", source_origin=OriginType.PARAMETER
+            pyobject=bytearray(b"456"), source_name="test2", source_value="bar", source_origin=OriginType.BODY
         )
         result = mod.do_bytearray_extend(ba1, ba2)
         assert result == bytearray(b"123456")

tests/appsec/iast/taint_tracking/test_native_taint_range.py

@@ -23,6 +23,9 @@
     from ddtrace.appsec._iast._taint_tracking import shift_taint_range
     from ddtrace.appsec._iast._taint_tracking import shift_taint_ranges
     from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+    from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
+    from ddtrace.appsec._iast._taint_tracking.aspects import format_aspect
+    from ddtrace.appsec._iast._taint_tracking.aspects import join_aspect
 except (ImportError, AttributeError):
     pytest.skip("IAST not supported for this Python version", allow_module_level=True)
 
@@ -86,6 +89,70 @@ def test_unicode_fast_tainting():
         assert not is_notinterned_notfasttainted_unicode(c)
 
 
+def test_collisions():
+    not_tainted = []
+    tainted = []
+    mixed_tainted_ids = []
+    mixed_nottainted = []
+    mixed_tainted_and_nottainted = []
+
+    # Generate untainted strings
+    for i in range(10):
+        not_tainted.append("N%04d" % i)
+
+    # Generate tainted strings
+    for i in range(10000):
+        t = taint_pyobject(
+            "T%04d" % i, source_name="request_body", source_value="hello", source_origin=OriginType.PARAMETER
+        )
+        tainted.append(t)
+
+    for t in tainted:
+        assert len(get_ranges(t)) > 0
+
+    for n in not_tainted:
+        assert get_ranges(n) == []
+
+    # Do join and format operations mixing tainted and untainted strings, store in mixed_tainted_and_nottainted
+    for _ in range(10000):
+        n1 = add_aspect(add_aspect("_", random.choice(not_tainted)), "{}_")
+
+        t2 = random.choice(tainted)
+
+        n3 = random.choice(not_tainted)
+
+        t4 = random.choice(tainted)
+        mixed_tainted_ids.append(id(t4))
+
+        t5 = format_aspect(n1.format, n1, t4)
+        mixed_tainted_ids.append(id(t5))
+
+        t6 = join_aspect(t5.join, t5, [t2, n3])
+        mixed_tainted_ids.append(id(t6))
+
+    for t in mixed_tainted_and_nottainted:
+        assert len(get_ranges(t)) == 2
+
+    # Do join and format operations with only untainted strings, store in mixed_nottainted
+    for i in range(10):
+        n1 = add_aspect("===", not_tainted[i])
+
+        n2 = random.choice(not_tainted)
+
+        n3 = random.choice(not_tainted)
+
+        n4 = random.choice(not_tainted)
+
+        n5 = format_aspect(n1.format, n1, [n4])
+
+        n6 = join_aspect(n5.join, n5, [n2, n3])
+
+        mixed_nottainted.append(n6)
+
+    for n in mixed_nottainted:
+        assert len(get_ranges(n)) == 0, f"id {id(n)} in {(id(n) in mixed_tainted_ids)}"
+
+
 def test_set_get_ranges_str():
     s1 = "abcde😁"
     s2 = "defg"