chore(asm): move iast taint tracking dict to asm context (#5397)

IAST: Move the taint tracking information to the asm context so its
scope is the current request.

This way we resolve the problem of the memory leaked from tainted
objects.
For details, compare appsec-iast-latest vs appsec-iast-staging in
reliability environment (staging being this branch)

## 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/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Author is aware of the performance implications of this PR as
reported in the benchmarks PR comment.

## 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 is aware of, and discussed the performance implications
of this PR as reported in the benchmarks PR comment.

Author: gnufede

ddtrace/appsec/_asm_request_context.py

@@ -24,7 +24,6 @@
 thread will have a different context.
 """
 
-
 # FIXME: remove these and use the new context API once implemented and allowing
 # contexts without spans
 
@@ -37,6 +36,7 @@
 _DD_WAF_CALLBACK = contextvars.ContextVar("datadog_early_waf_callback", default=None)
 _DD_WAF_RESULTS = contextvars.ContextVar("datadog_early_waf_results", default=([[], [], []]))
 _DD_WAF_SENT = contextvars.ContextVar("datadog_waf_adress_sent", default=None)
+_DD_IAST_TAINT_DICT = contextvars.ContextVar("datadog_iast_taint_dict", default={})
 
 
 def reset():  # type: () -> None
@@ -45,6 +45,7 @@ def reset():  # type: () -> None
     _DD_EARLY_HEADERS_CASE_SENSITIVE_CONTEXTVAR.set(False)
     _DD_BLOCK_REQUEST_CALLABLE.set(None)
     _DD_WAF_SENT.set(set())
+    _DD_IAST_TAINT_DICT.set({})
 
 
 def set_ip(ip):  # type: (Optional[str]) -> None
@@ -55,6 +56,14 @@ def get_ip():  # type: () -> Optional[str]
     return _DD_EARLY_IP_CONTEXTVAR.get()
 
 
+def set_taint_dict(taint_dict):  # type: (dict) -> None
+    _DD_IAST_TAINT_DICT.set(taint_dict)
+
+
+def get_taint_dict():  # type: () -> dict
+    return _DD_IAST_TAINT_DICT.get()
+
+
 # Note: get/set headers use Any since we just carry the headers here without changing or using them
 # and different frameworks use different types that we don't want to force it into a Mapping at the
 # early point set_headers is usually called

ddtrace/appsec/iast/_ast/aspects/__init__.py

@@ -4,11 +4,11 @@
 import codecs
 
 from ddtrace.appsec.iast._input_info import Input_info
-from ddtrace.appsec.iast._taint_tracking import add_taint_pyobject  # type: ignore[attr-defined]
-from ddtrace.appsec.iast._taint_tracking import get_tainted_ranges  # type: ignore[attr-defined]
-from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted  # type: ignore[attr-defined]
-from ddtrace.appsec.iast._taint_tracking import set_tainted_ranges  # type: ignore[attr-defined]
-from ddtrace.appsec.iast._taint_tracking import taint_pyobject  # type: ignore[attr-defined]
+from ddtrace.appsec.iast._taint_tracking import add_taint_pyobject
+from ddtrace.appsec.iast._taint_tracking import get_tainted_ranges
+from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
+from ddtrace.appsec.iast._taint_tracking import set_tainted_ranges
+from ddtrace.appsec.iast._taint_tracking import taint_pyobject
 
 
 def str_aspect(*args, **kwargs):

ddtrace/appsec/iast/_taint_tracking/__init__.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+from typing import TYPE_CHECKING
+
+from ddtrace.appsec._asm_request_context import get_taint_dict
+from ddtrace.appsec._asm_request_context import set_taint_dict
+from ddtrace.appsec.iast._taint_tracking._native import new_pyobject_id
+from ddtrace.appsec.iast._taint_tracking._native import setup  # noqa: F401
+
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from ddtrace.appsec.iast._input_info import Input_info
+
+
+def add_taint_pyobject(pyobject, op1, op2):  # type: (Any, Any, Any) -> Any
+    if not (is_pyobject_tainted(op1) or is_pyobject_tainted(op2)):
+        return pyobject
+
+    pyobject = new_pyobject_id(pyobject, len(pyobject))
+    taint_dict = get_taint_dict()
+    new_ranges = []
+    if is_pyobject_tainted(op1):
+        new_ranges = list(taint_dict[id(op1)])
+    if is_pyobject_tainted(op2):
+        offset = len(op1)
+        for input_info, start, size in taint_dict[id(op2)]:
+            new_ranges.append((input_info, start + offset, size))
+
+    taint_dict[id(pyobject)] = tuple(new_ranges)
+    set_taint_dict(taint_dict)
+    return pyobject
+
+
+def taint_pyobject(pyobject, input_info):  # type: (Any, Input_info) -> Any
+    if not pyobject:  # len(pyobject) < 1
+        return pyobject
+    assert input_info is not None
+    len_pyobject = len(pyobject)
+    pyobject = new_pyobject_id(pyobject, len_pyobject)
+    taint_dict = get_taint_dict()
+    taint_dict[id(pyobject)] = ((input_info, 0, len_pyobject),)
+    set_taint_dict(taint_dict)
+    return pyobject
+
+
+def is_pyobject_tainted(pyobject):  # type: (Any) -> bool
+    return id(pyobject) in get_taint_dict()
+
+
+def set_tainted_ranges(pyobject, ranges):  # type: (Any, tuple) -> None
+    taint_dict = get_taint_dict()
+    assert pyobject not in taint_dict
+    taint_dict[id(pyobject)] = ranges
+    set_taint_dict(taint_dict)
+
+
+def get_tainted_ranges(pyobject):  # type: (Any) -> tuple
+    return get_taint_dict().get(id(pyobject), tuple())
+
+
+def clear_taint_mapping():  # type: () -> None
+    set_taint_dict({})

ddtrace/appsec/iast/_taint_tracking/_native.cpp

@@ -0,0 +1,64 @@
+#include <Python.h>
+#include <iostream>
+#include <tuple>
+
+PyObject *bytes_join = NULL;
+PyObject *bytearray_join = NULL;
+PyObject *empty_bytes = NULL;
+PyObject *empty_bytearray = NULL;
+PyObject *empty_unicode = NULL;
+
+static PyObject *setup(PyObject *Py_UNUSED(module), PyObject *args) {
+  PyArg_ParseTuple(args, "OO", &bytes_join, &bytearray_join);
+  empty_bytes = PyBytes_FromString("");
+  empty_bytearray = PyByteArray_FromObject(empty_bytes);
+  empty_unicode = PyUnicode_New(0, 127);
+  Py_RETURN_NONE;
+}
+
+static PyObject *new_pyobject_id(PyObject *Py_UNUSED(module), PyObject *args) {
+  PyObject *tainted_object;
+  Py_ssize_t object_length;
+  PyArg_ParseTuple(args, "On", &tainted_object, &object_length);
+  if (PyUnicode_Check(tainted_object)) {
+    if (PyUnicode_CHECK_INTERNED(tainted_object) == 0) { // SSTATE_NOT_INTERNED
+      Py_INCREF(tainted_object);
+      return tainted_object;
+    }
+    return PyUnicode_Join(empty_unicode,
+                          Py_BuildValue("(OO)", tainted_object, empty_unicode));
+  } else if (object_length > 1) {
+    // Bytes and bytearrays with length > 1 are not interned
+    Py_INCREF(tainted_object);
+    return tainted_object;
+  } else if (PyBytes_Check(tainted_object)) {
+    return PyObject_CallFunctionObjArgs(
+        bytes_join, empty_bytes,
+        Py_BuildValue("(OO)", tainted_object, empty_bytes), NULL);
+  } else {
+    return PyObject_CallFunctionObjArgs(
+        bytearray_join, empty_bytearray,
+        Py_BuildValue("(OO)", tainted_object, empty_bytearray), NULL);
+  }
+}
+
+static PyMethodDef TaintTrackingMethods[] = {
+    // We are using  METH_VARARGS because we need compatibility with
+    // python 3.5, 3.6. but METH_FASTCALL could be used instead for python
+    // >= 3.7
+    {"setup", (PyCFunction)setup, METH_VARARGS, "setup tainting module"},
+    {"new_pyobject_id", (PyCFunction)new_pyobject_id, METH_VARARGS,
+     "new_pyobject_id"},
+    {NULL, NULL, 0, NULL}};
+
+static struct PyModuleDef taint_tracking = {
+    PyModuleDef_HEAD_INIT, "ddtrace.appsec.iast._taint_tracking._native",
+    "taint tracking module", -1, TaintTrackingMethods};
+
+PyMODINIT_FUNC PyInit__native(void) {
+  PyObject *m;
+  m = PyModule_Create(&taint_tracking);
+  if (m == NULL)
+    return NULL;
+  return m;
+}

ddtrace/appsec/iast/_taint_tracking/_taint_tracking.cpp

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 from ddtrace.appsec.iast._input_info import Input_info
-from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted  # type: ignore[attr-defined]
-from ddtrace.appsec.iast._taint_tracking import taint_pyobject  # type: ignore[attr-defined]
+from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
+from ddtrace.appsec.iast._taint_tracking import taint_pyobject
 
 
 DBAPI_INTEGRATIONS = ("sqlite", "psycopg", "mysql", "mariadb")

ddtrace/appsec/iast/_taint_utils.py

@@ -178,7 +178,7 @@ def drop(module_name):
 
             from ddtrace.appsec.iast._ast.ast_patching import _should_iast_patch
             from ddtrace.appsec.iast._loader import _exec_iast_patched_module
-            from ddtrace.appsec.iast._taint_tracking import setup  # type:ignore[attr-defined]
+            from ddtrace.appsec.iast._taint_tracking import setup
             from ddtrace.internal.module import ModuleWatchdog
 
             setup(bytes.join, bytearray.join)

ddtrace/bootstrap/sitecustomize.py

@@ -109,7 +109,7 @@ def taint_request_init(wrapped, instance, args, kwargs):
     wrapped(*args, **kwargs)
     if _is_iast_enabled():
         from ddtrace.appsec.iast._input_info import Input_info
-        from ddtrace.appsec.iast._taint_tracking import taint_pyobject  # type: ignore[attr-defined]
+        from ddtrace.appsec.iast._taint_tracking import taint_pyobject
 
         taint_pyobject(
             instance.query_string,

ddtrace/contrib/flask/patch.py

@@ -308,7 +308,7 @@ def get_exts_for(name):
         if sys.version_info >= (3, 6, 0):
             ext_modules.append(
                 Extension(
-                    "ddtrace.appsec.iast._taint_tracking",
+                    "ddtrace.appsec.iast._taint_tracking._native",
                     # Sort source files for reproducibility
                     sources=sorted(
                         glob.glob(