From f2c630ed267b58fe405d33dfbadc792fd9fcf953 Mon Sep 17 00:00:00 2001
From: Santiago Mola <santiago.mola@datadoghq.com>
Date: Fri, 25 Apr 2025 08:27:32 +0200
Subject: [PATCH] fix(iast): avoid excessive filtering of stacktrace locations

---
 ddtrace/appsec/_iast/_stacktrace.c            | 247 ++++++++++++------
 ddtrace/appsec/_iast/_stacktrace.pyi          |   2 +-
 ddtrace/appsec/_iast/taint_sinks/_base.py     |  23 +-
 ...st-stacktrace-filter-6afd6a9568f36d99.yaml |   4 +
 .../appsec/iast/taint_sinks/test_weak_hash.py |   4 +-
 tests/appsec/iast_memcheck/_stacktrace_py.py  |  42 ---
 .../iast_memcheck/fixtures/stacktrace.py      |  15 +-
 .../iast_memcheck/test_iast_mem_check.py      |  74 +-----
 8 files changed, 201 insertions(+), 210 deletions(-)
 create mode 100644 releasenotes/notes/fix-iast-stacktrace-filter-6afd6a9568f36d99.yaml
 delete mode 100644 tests/appsec/iast_memcheck/_stacktrace_py.py

diff --git a/ddtrace/appsec/_iast/_stacktrace.c b/ddtrace/appsec/_iast/_stacktrace.c
index 893d2c5a076..057e9d5a803 100644
--- a/ddtrace/appsec/_iast/_stacktrace.c
+++ b/ddtrace/appsec/_iast/_stacktrace.c
@@ -9,6 +9,7 @@ static __thread int in_stacktrace = 0;
 #include <Python.h>
 #include <frameobject.h>
 #include <patchlevel.h>
+#include <stdbool.h>
 
 #ifdef _WIN32
 #define DD_TRACE_INSTALLED_PREFIX "\\ddtrace\\"
@@ -25,6 +26,7 @@ static __thread int in_stacktrace = 0;
 #define GET_LINENO(frame) PyFrame_GetLineNumber((PyFrameObject*)frame)
 #define GET_FRAME(tstate) PyThreadState_GetFrame(tstate)
 #define GET_PREVIOUS(frame) PyFrame_GetBack(frame)
+#define FRAME_INCREF(frame) Py_INCREF((PyObject*)frame)
 #define FRAME_DECREF(frame) Py_DecRef((PyObject*)frame)
 #define FRAME_XDECREF(frame) Py_XDECREF((PyObject*)frame)
 #define FILENAME_DECREF(filename) Py_DecRef(filename)
@@ -68,6 +70,7 @@ GET_FUNCTION(PyFrameObject* frame)
 #define GET_FRAME(tstate) tstate->frame
 #define GET_PREVIOUS(frame) frame->f_back
 #define GET_FILENAME(frame) ((PyObject*)(frame->f_code->co_filename))
+#define FRAME_INCREF(frame)
 #define FRAME_DECREF(frame)
 #define FRAME_XDECREF(frame)
 #define FILENAME_DECREF(filename)
@@ -88,6 +91,14 @@ GET_FUNCTION(PyFrameObject* frame)
 #endif
 #endif
 
+// Python standard library path
+static char* STDLIB_PATH = NULL;
+static ssize_t STDLIB_PATH_LEN = 0;
+
+// Python site-packages path
+static char* PURELIB_PATH = NULL;
+static ssize_t PURELIB_PATH_LEN = 0;
+
 static inline PyObject*
 SAFE_GET_LOCALS(PyFrameObject* frame)
 {
@@ -121,115 +132,191 @@ GET_CLASS(PyFrameObject* frame)
 }
 
 /**
- * get_file_and_line
- *
- * Get the filename, line number, function name and class name of the original wrapped
- * function to report it.
- *
- * Returns a tuple:
- *     (filename, line_number, function name, class name)
- **/
-static PyObject*
-get_file_and_line(PyObject* Py_UNUSED(module), PyObject* cwd_obj)
+ * Checks if the filename is special.
+ * For example, a frozen module (`<frozen 'os'>`), a template (`<template>`), etc.
+ */
+static inline bool
+_is_special_frame(const char* filename)
 {
-    // Mark that we are now capturing a stack trace to avoid reentrant calls on GET_LOCALS
-    in_stacktrace = 1;
+    return filename && strncmp(filename, "<", strlen("<")) == 0;
+}
 
-    PyThreadState* tstate = PyThreadState_Get();
-    if (!tstate) {
-        goto exit_0;
-    }
+static inline bool
+_is_ddtrace_filename(const char* filename)
+{
+    return filename && strstr(filename, DD_TRACE_INSTALLED_PREFIX) != NULL && strstr(filename, TESTS_PREFIX) == NULL;
+}
 
-    int line;
-    PyObject* filename_o = NULL;
-    PyObject* result = NULL;
-    PyObject* cwd_bytes = NULL;
-    char* cwd = NULL;
+static inline bool
+_is_site_packages_filename(const char* filename)
+{
+    const bool res = filename && PURELIB_PATH && strncmp(filename, PURELIB_PATH, PURELIB_PATH_LEN) == 0;
+    return res;
+}
+
+static inline bool
+_is_stdlib_filename(const char* filename)
+{
+    // site-packages is often a subdirectory of stdlib directory, so stdlib
+    // path is defined as prefixed by stdlib and not prefixed by purelib.
+    // TODO: As of Python 3.10, we could use sys.stdlib_module_names.
+    const bool res = filename && STDLIB_PATH && !_is_site_packages_filename(filename) &&
+                     strncmp(filename, STDLIB_PATH, STDLIB_PATH_LEN) == 0;
+    return res;
+}
 
-    if (!PyUnicode_FSConverter(cwd_obj, &cwd_bytes)) {
-        goto exit_0;
+static char*
+get_sysconfig_path(const char* name)
+{
+    PyObject* sysconfig_mod = PyImport_ImportModule("sysconfig");
+    if (!sysconfig_mod) {
+        return NULL;
     }
-    cwd = PyBytes_AsString(cwd_bytes);
-    if (!cwd) {
-        goto exit_0;
+
+    PyObject* path = PyObject_CallMethod(sysconfig_mod, "get_path", "s", name);
+    if (!path) {
+        Py_DECREF(sysconfig_mod);
+        return NULL;
     }
 
-    PyFrameObject* frame = GET_FRAME(tstate);
-    if (!frame) {
-        goto exit_0;
+    const char* path_str = PyUnicode_AsUTF8(path);
+    char* res = NULL;
+    if (path_str) {
+        res = strdup(path_str);
     }
+    Py_DECREF(path);
+    Py_DECREF(sysconfig_mod);
+    return res;
+}
 
+/**
+ * Gets a reference to a PyFrameObject and walks up the stack until a relevant frame is found.
+ *
+ * Returns a new reference to the PyFrameObject.
+ *
+ * The caller is not responsible for DECREF'ing the given PyFrameObject, but it is responsible for
+ * DECREF'ing the returned PyFrameObject.
+ */
+static PyFrameObject*
+_find_relevant_frame(PyFrameObject* frame, bool allow_site_packages)
+{
     while (NULL != frame) {
-        filename_o = GET_FILENAME(frame);
+        PyObject* filename_o = GET_FILENAME(frame);
         if (!filename_o) {
-            goto exit;
+            FRAME_DECREF(frame);
+            return NULL;
         }
         const char* filename = PyUnicode_AsUTF8(filename_o);
-        if (((strstr(filename, DD_TRACE_INSTALLED_PREFIX) != NULL && strstr(filename, TESTS_PREFIX) == NULL)) ||
-            (strstr(filename, SITE_PACKAGES_PREFIX) != NULL || strstr(filename, cwd) == NULL)) {
+        if (_is_special_frame(filename) || _is_ddtrace_filename(filename) || _is_stdlib_filename(filename) ||
+            (!allow_site_packages && _is_site_packages_filename(filename))) {
             PyFrameObject* prev_frame = GET_PREVIOUS(frame);
             FRAME_DECREF(frame);
             FILENAME_DECREF(filename_o);
             frame = prev_frame;
             continue;
         }
-        /*
-         frame->f_lineno will not always return the correct line number
-         you need to call PyCode_Addr2Line().
-        */
-        line = GET_LINENO(frame);
-        PyObject* line_obj = Py_BuildValue("i", line);
-        if (!line_obj) {
-            goto exit;
-        }
-        PyObject* func_name = GET_FUNCTION(frame);
-        if (!func_name) {
-            Py_DecRef(line_obj);
-            goto exit;
-        }
-        PyObject* class_name = GET_CLASS(frame);
-        if (!class_name) {
-            Py_DecRef(line_obj);
-            Py_DecRef(func_name);
-            goto exit;
-        }
-        result = PyTuple_Pack(4, filename_o, line_obj, func_name, class_name);
-        Py_DecRef(func_name);
-        Py_DecRef(class_name);
-        Py_DecRef(line_obj);
+        FILENAME_DECREF(filename_o);
         break;
     }
-    if (result == NULL) {
-        goto exit_0;
+    return frame;
+}
+
+static PyObject*
+_get_result_tuple(PyFrameObject* frame)
+{
+    PyObject* result = NULL;
+    PyObject* filename_o = NULL;
+    PyObject* line_o = NULL;
+    PyObject* funcname_o = NULL;
+    PyObject* classname_o = NULL;
+
+    filename_o = GET_FILENAME(frame);
+    if (!filename_o) {
+        goto error;
     }
 
-exit:
-    Py_DecRef(cwd_bytes);
-    FRAME_XDECREF(frame);
+    // frame->f_lineno will not always return the correct line number
+    // you need to call PyCode_Addr2Line().
+    int line = GET_LINENO(frame);
+    line_o = Py_BuildValue("i", line);
+    if (!line_o) {
+        goto error;
+    }
+    funcname_o = GET_FUNCTION(frame);
+    if (!funcname_o) {
+        goto error;
+    }
+    classname_o = GET_CLASS(frame);
+    if (!classname_o) {
+        goto error;
+    }
+    result = PyTuple_Pack(4, filename_o, line_o, funcname_o, classname_o);
+
+error:
     FILENAME_XDECREF(filename_o);
-    in_stacktrace = 0;
+    Py_XDECREF(line_o);
+    Py_XDECREF(funcname_o);
+    Py_XDECREF(classname_o);
     return result;
+}
+
+/**
+ * get_file_and_line
+ *
+ * Get the filename, line number, function name and class name of the original wrapped
+ * function to report it.
+ *
+ * Returns a tuple:
+ *     (filename, line_number, function name, class name)
+ **/
+static PyObject*
+get_file_and_line(PyObject* Py_UNUSED(module), PyObject* Py_UNUSED(args))
+{
+    // Mark that we are now capturing a stack trace to avoid reentrant calls on GET_LOCALS
+    in_stacktrace = 1;
+    PyFrameObject* frame = NULL;
+    PyFrameObject* backup_frame = NULL;
+    PyObject* result = NULL;
+    PyThreadState* tstate = PyThreadState_Get();
+    if (!tstate) {
+        goto exit;
+    }
 
-exit_0:; /* Label must be followed by a statement */
-    // Return "", -1, "", ""
-    PyObject* line_obj = Py_BuildValue("i", -1);
-    filename_o = PyUnicode_FromString("");
-    PyObject* func_name = PyUnicode_FromString("");
-    PyObject* class_name = PyUnicode_FromString("");
-    result = PyTuple_Pack(4, filename_o, line_obj, func_name, class_name);
-    Py_DecRef(cwd_bytes);
+    frame = GET_FRAME(tstate);
+    if (!frame) {
+        goto exit;
+    }
+
+    // Skip all frames until the first non-ddtrace and non-stdlib frame.
+    // Store that frame as backup (if any). If there is no better frame, fallback to this.
+    // This happens, for example, when the vulnerability is in a package installed in site-packages.
+    frame = _find_relevant_frame(frame, true);
+    if (NULL == frame) {
+        goto exit;
+    }
+    backup_frame = frame;
+    FRAME_INCREF(backup_frame);
+
+    // Continue skipping until we find a frame that is both non-ddtrace and non-site-packages.
+    frame = _find_relevant_frame(frame, false);
+    if (NULL == frame) {
+        frame = backup_frame;
+        backup_frame = NULL;
+    } else {
+        FRAME_DECREF(backup_frame);
+    }
+
+    result = _get_result_tuple(frame);
+
+exit:
     FRAME_XDECREF(frame);
-    FILENAME_XDECREF(filename_o);
-    Py_DecRef(line_obj);
-    Py_DecRef(func_name);
-    Py_DecRef(class_name);
     in_stacktrace = 0;
     return result;
 }
 
 static PyMethodDef StacktraceMethods[] = { { "get_info_frame",
                                              (PyCFunction)get_file_and_line,
-                                             METH_O,
+                                             METH_NOARGS,
                                              "Stacktrace function: returns (filename, line, method, class)" },
                                            { NULL, NULL, 0, NULL } };
 
@@ -245,5 +332,13 @@ PyInit__stacktrace(void)
     PyObject* m = PyModule_Create(&stacktrace);
     if (m == NULL)
         return NULL;
+    STDLIB_PATH = get_sysconfig_path("stdlib");
+    if (STDLIB_PATH) {
+        STDLIB_PATH_LEN = strlen(STDLIB_PATH);
+    }
+    PURELIB_PATH = get_sysconfig_path("purelib");
+    if (PURELIB_PATH) {
+        PURELIB_PATH_LEN = strlen(PURELIB_PATH);
+    }
     return m;
 }
diff --git a/ddtrace/appsec/_iast/_stacktrace.pyi b/ddtrace/appsec/_iast/_stacktrace.pyi
index e50336454a0..46c1e435ac0 100644
--- a/ddtrace/appsec/_iast/_stacktrace.pyi
+++ b/ddtrace/appsec/_iast/_stacktrace.pyi
@@ -1 +1 @@
-def get_info_frame(cwd_obj): ...
+def get_info_frame(): ...
diff --git a/ddtrace/appsec/_iast/taint_sinks/_base.py b/ddtrace/appsec/_iast/taint_sinks/_base.py
index 50dc8c3e910..a4051069945 100644
--- a/ddtrace/appsec/_iast/taint_sinks/_base.py
+++ b/ddtrace/appsec/_iast/taint_sinks/_base.py
@@ -1,4 +1,5 @@
 import os
+import sysconfig
 from typing import Any
 from typing import Callable
 from typing import Optional
@@ -30,6 +31,9 @@
 
 TEXT_TYPES = Union[str, bytes, bytearray]
 
+PURELIB_PATH = sysconfig.get_path("purelib")
+STDLIB_PATH = sysconfig.get_path("stdlib")
+
 
 class taint_sink_deduplication(deduplication):
     def _check_deduplication(self):
@@ -130,21 +134,32 @@ def _prepare_report(
     @classmethod
     def _compute_file_line(cls) -> Tuple[Optional[str], Optional[int], Optional[str], Optional[str]]:
         file_name = line_number = function_name = class_name = None
-
-        frame_info = get_info_frame(CWD)
+        frame_info = get_info_frame()
         if not frame_info or frame_info[0] in ("", -1):
             return file_name, line_number, function_name, class_name
 
         file_name, line_number, function_name, class_name = frame_info
 
-        if file_name.startswith(CWD):
-            file_name = os.path.relpath(file_name, start=CWD)
+        file_name = cls._rel_path(file_name)
+        if not file_name:
+            log.debug("Could not relativize vulnerability location path: %s", frame_info[0])
+            return None, None, None, None
 
         if not cls.is_not_reported(file_name, line_number):
             return None, None, None, None
 
         return file_name, line_number, function_name, class_name
 
+    @staticmethod
+    def _rel_path(file_name: str) -> str:
+        if file_name.startswith(PURELIB_PATH):
+            return os.path.relpath(file_name, start=PURELIB_PATH)
+        if file_name.startswith(STDLIB_PATH):
+            return os.path.relpath(file_name, start=STDLIB_PATH)
+        if file_name.startswith(CWD):
+            return os.path.relpath(file_name, start=CWD)
+        return ""
+
     @classmethod
     def _create_evidence_and_report(
         cls,
diff --git a/releasenotes/notes/fix-iast-stacktrace-filter-6afd6a9568f36d99.yaml b/releasenotes/notes/fix-iast-stacktrace-filter-6afd6a9568f36d99.yaml
new file mode 100644
index 00000000000..4258704651c
--- /dev/null
+++ b/releasenotes/notes/fix-iast-stacktrace-filter-6afd6a9568f36d99.yaml
@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Code Security: IAST: Avoid excessive filtering of stacktrace locations when finding vulnerabilities. After this change, vulnerabilities that were previously discarded will now be reported. In particular, if they were found within code in site-packages or outside of the working directory.
diff --git a/tests/appsec/iast/taint_sinks/test_weak_hash.py b/tests/appsec/iast/taint_sinks/test_weak_hash.py
index 36efd5c1e69..6ced416a427 100644
--- a/tests/appsec/iast/taint_sinks/test_weak_hash.py
+++ b/tests/appsec/iast/taint_sinks/test_weak_hash.py
@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+import os
 import sys
 from unittest import mock
 
@@ -95,9 +96,10 @@ def test_weak_hash_hashlib(iast_context_defaults, hash_func, method):
     ],
 )
 def test_ensure_line_reported_is_minus_one_for_edge_cases(iast_context_defaults, hash_func, method, fake_line):
+    absolute_path = os.path.abspath(WEAK_ALGOS_FIXTURES_PATH)
     with mock.patch(
         "ddtrace.appsec._iast.taint_sinks._base.get_info_frame",
-        return_value=(WEAK_ALGOS_FIXTURES_PATH, fake_line, "", ""),
+        return_value=(absolute_path, fake_line, "", ""),
     ):
         parametrized_weak_hash(hash_func, method)
 
diff --git a/tests/appsec/iast_memcheck/_stacktrace_py.py b/tests/appsec/iast_memcheck/_stacktrace_py.py
deleted file mode 100644
index 6135461e8b8..00000000000
--- a/tests/appsec/iast_memcheck/_stacktrace_py.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import inspect
-import os
-from typing import TYPE_CHECKING
-
-
-if TYPE_CHECKING:  # pragma: no cover
-    from typing import Optional  # noqa:F401
-    from typing import Text  # noqa:F401
-    from typing import Tuple  # noqa:F401
-
-
-FIRST_FRAME_NO_DDTRACE = 1
-
-DD_TRACE_INSTALLED_PREFIX = os.sep + "ddtrace" + os.sep
-SITE_PACKAGES_PREFIX = os.sep + "site-packages" + os.sep
-TESTS_PREFIX = os.sep + "tests" + os.sep
-
-
-def get_info_frame(cwd):
-    # type: (Text) -> Optional[Tuple[Text, int, Text, Text]]
-    """Get the filename (path + filename) and line number of the original wrapped function to report it.
-
-    CAVEAT: We should migrate this function to native code to improve the performance.
-    """
-    stack = inspect.stack()
-    for frame in stack[FIRST_FRAME_NO_DDTRACE:]:
-        filename = frame.filename
-        lineno = frame.lineno
-        function_name = frame.function
-        instance = frame.frame.f_locals.get("self")
-        class_name = instance.__class__.__name__ if instance else None
-
-        if (
-            (DD_TRACE_INSTALLED_PREFIX in filename and TESTS_PREFIX not in filename)
-            or (cwd not in filename)
-            or (SITE_PACKAGES_PREFIX in filename)
-        ):
-            continue
-
-        return filename, lineno, function_name, class_name
-
-    return None
diff --git a/tests/appsec/iast_memcheck/fixtures/stacktrace.py b/tests/appsec/iast_memcheck/fixtures/stacktrace.py
index fa08bcc334e..6ba133dcf27 100644
--- a/tests/appsec/iast_memcheck/fixtures/stacktrace.py
+++ b/tests/appsec/iast_memcheck/fixtures/stacktrace.py
@@ -1,7 +1,6 @@
 import os
 
 from ddtrace.appsec._iast._stacktrace import get_info_frame
-from tests.appsec.iast_memcheck._stacktrace_py import get_info_frame as get_info_frame_py
 
 
 CWD = os.path.abspath(os.getcwd())
@@ -84,16 +83,6 @@ def func_19(a, b, c):
 
 
 def func_20(a, b, c):
-    if b == "py":
-        func = get_info_frame_py
-    else:
-        func = get_info_frame
-    if a == "empty_byte":
-        frame_info = func(b"")
-    elif a == "empty_string":
-        frame_info = func("")
-    elif a == "random_string":
-        frame_info = func("aaaa")
-    else:
-        frame_info = func(CWD)
+    func = get_info_frame
+    frame_info = func()
     return frame_info
diff --git a/tests/appsec/iast_memcheck/test_iast_mem_check.py b/tests/appsec/iast_memcheck/test_iast_mem_check.py
index 9541bfd2ede..098dbcd8551 100644
--- a/tests/appsec/iast_memcheck/test_iast_mem_check.py
+++ b/tests/appsec/iast_memcheck/test_iast_mem_check.py
@@ -15,7 +15,6 @@
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from tests.appsec.iast.iast_utils import _iast_patched_module
 from tests.appsec.iast.taint_sinks.conftest import _get_span_report
-from tests.appsec.iast_memcheck._stacktrace_py import get_info_frame as get_info_frame_py
 from tests.appsec.iast_memcheck.fixtures.stacktrace import func_1
 
 
@@ -181,7 +180,7 @@ def test_stacktrace_memory_check():
 @pytest.mark.limit_leaks("460 B", filter_fn=IASTFilter())
 def test_stacktrace_memory_check_direct_call():
     for _ in range(LOOPS):
-        frame_info = get_info_frame(CWD)
+        frame_info = get_info_frame()
         if not frame_info:
             pytest.fail("No stacktrace")
 
@@ -190,74 +189,3 @@ def test_stacktrace_memory_check_direct_call():
         assert line_number > 0
         assert method == "test_stacktrace_memory_check_direct_call"
         assert not class_
-
-
-@pytest.mark.limit_leaks("460 KB", filter_fn=IASTFilter())
-def test_stacktrace_memory_check_no_native():
-    for _ in range(LOOPS):
-        frame_info = func_1("", "py", "3")
-        if not frame_info:
-            pytest.fail("No stacktrace")
-
-        file_name, line_number, method, class_ = frame_info
-        assert file_name
-        assert line_number > 0
-        assert method == "func_20"
-        assert not class_
-
-
-@pytest.mark.limit_leaks("24 KB", filter_fn=IASTFilter())
-def test_stacktrace_memory_check_no_native_direct_call():
-    for _ in range(2):
-        frame_info = get_info_frame_py(CWD)
-        if not frame_info:
-            pytest.fail("No stacktrace")
-
-        file_name, line_number, method, class_ = frame_info
-        assert file_name
-        assert line_number > 0
-        assert method == "test_stacktrace_memory_check_no_native_direct_call"
-        assert not class_
-
-
-@pytest.mark.limit_leaks("440 B", filter_fn=IASTFilter())
-def test_stacktrace_memory_empty_byte_check():
-    for _ in range(LOOPS):
-        frame_info = func_1("empty_byte", "2", "3")
-        if not frame_info:
-            pytest.fail("No stacktrace")
-
-        file_name, line_number, method, class_ = frame_info
-        assert file_name
-        assert line_number > 0
-        assert method == "func_20"
-        assert not class_
-
-
-@pytest.mark.limit_leaks("440 B", filter_fn=IASTFilter())
-def test_stacktrace_memory_empty_string_check():
-    for _ in range(LOOPS):
-        frame_info = func_1("empty_string", "2", "3")
-        if not frame_info:
-            pytest.fail("No stacktrace")
-
-        file_name, line_number, method, class_ = frame_info
-        assert file_name
-        assert line_number > 0
-        assert method == "func_20"
-        assert not class_
-
-
-@pytest.mark.limit_leaks("10 KB", filter_fn=IASTFilter())
-def test_stacktrace_memory_random_string_check():
-    """2.1 KB is enough but CI allocates 1.0 MB bytes"""
-    for _ in range(LOOPS):
-        frame_info = func_1("random_string", "2", "3")
-        if not frame_info:
-            pytest.fail("No stacktrace")
-
-        file_name, line_number, method, class_ = frame_info
-        assert file_name == ""
-        assert line_number == -1
-        assert not method
-        assert not class_