refactor(appsec): move _stacktrace from _iast to _shared package

ddtrace/appsec/_iast/__init__.py

@@ -199,6 +199,7 @@ def enable_iast_propagation():
         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 initialize_native_state
+        from ddtrace.appsec._shared._stacktrace import get_info_frame  # noqa: F401
 
         global _iast_propagation_enabled
         if _iast_propagation_enabled:

ddtrace/appsec/_patch_utils.py

@@ -15,6 +15,10 @@
 
 log = get_logger(__name__)
 
+# Lazy-cached reference to avoid loading _stacktrace (a C extension used only
+# by IAST) when _patch_utils is imported by non-IAST code paths.
+_get_info_frame = None
+
 # Cached paths for relativizing file paths (computed once at import time).
 _CWD = os.path.abspath(os.getcwd())
 _PURELIB_PATH = sysconfig.get_path("purelib") or ""
@@ -52,12 +56,12 @@ def get_caller_frame_info() -> tuple:
 
     Returns (None, None, None, None) when no relevant frame is found.
     """
-    try:
-        from ddtrace.appsec._iast._stacktrace import get_info_frame
-    except ImportError:
-        return None, None, None, None
+    global _get_info_frame
+    if _get_info_frame is None:
+        from ddtrace.appsec._shared._stacktrace import get_info_frame
 
-    frame_info = get_info_frame()
+        _get_info_frame = get_info_frame
+    frame_info = _get_info_frame()
     if not frame_info or frame_info[0] in ("", -1, None):
         return None, None, None, None
 

ddtrace/appsec/_iast/_stacktrace.c → ddtrace/appsec/_shared/_stacktrace.c

@@ -359,7 +359,7 @@ static PyMethodDef StacktraceMethods[] = { { "get_info_frame",
                                            { NULL, NULL, 0, NULL } };
 
 static struct PyModuleDef stacktrace = { PyModuleDef_HEAD_INIT,
-                                         "ddtrace.appsec._iast._stacktrace",
+                                         "ddtrace.appsec._shared._stacktrace",
                                          "stacktrace module",
                                          -1,
                                          StacktraceMethods };

ddtrace/appsec/_shared/_stacktrace.pyi

@@ -0,0 +1,3 @@
+from typing import Optional
+
+def get_info_frame() -> tuple[Optional[str], Optional[int], Optional[str], Optional[str]]: ...

setup.py

@@ -1275,9 +1275,9 @@ def get_exts_for(name):
     if platform.system() not in ("Windows", ""):
         ext_modules.append(
             Extension(
-                "ddtrace.appsec._iast._stacktrace",
+                "ddtrace.appsec._shared._stacktrace",
                 sources=[
-                    "ddtrace/appsec/_iast/_stacktrace.c",
+                    "ddtrace/appsec/_shared/_stacktrace.c",
                 ],
                 extra_compile_args=extra_compile_args + debug_compile_args + fast_build_args,
             )

tests/appsec/architectures/test_appsec_loading_modules.py

@@ -17,6 +17,7 @@
 MODULE_IAST_ONLY = [
     "ddtrace.appsec._iast",
     "ddtrace.appsec._iast._taint_tracking._native",
+    "ddtrace.appsec._shared._stacktrace",
 ]
 
 

tests/appsec/iast/test_stacktrace.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from ddtrace.appsec._iast._stacktrace import get_info_frame
+from ddtrace.appsec._shared._stacktrace import get_info_frame
 
 
 def test_stacktrace():

tests/appsec/iast_memcheck/fixtures/stacktrace.py

@@ -1,6 +1,6 @@
 import os
 
-from ddtrace.appsec._iast._stacktrace import get_info_frame
+from ddtrace.appsec._shared._stacktrace import get_info_frame
 
 
 CWD = os.path.abspath(os.getcwd())

tests/appsec/iast_memcheck/test_iast_mem_check.py

@@ -8,11 +8,11 @@
 from ddtrace.appsec._iast._iast_request_context_base import _iast_finish_request
 from ddtrace.appsec._iast._iast_request_context_base import _iast_start_request
 from ddtrace.appsec._iast._iast_request_context_base import _num_objects_tainted_in_request
-from ddtrace.appsec._iast._stacktrace import get_info_frame
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking._context import debug_context_array_size
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import get_tainted_ranges
+from ddtrace.appsec._shared._stacktrace import get_info_frame
 from tests.appsec.iast.iast_utils import _iast_patched_module
 from tests.appsec.iast_memcheck.fixtures.stacktrace import func_1
 

tests/appsec/integrations/flask_tests/test_iast_flask.py

@@ -731,7 +731,6 @@ def sqli_14():
     def test_flask_request_body(self):
         @self.app.route("/sqli/body/", methods=("POST",))
         def sqli_10():
-            import json
             import sqlite3
 
             from flask import request
@@ -740,10 +739,7 @@ def sqli_10():
 
             con = sqlite3.connect(":memory:")
             cur = con.cursor()
-            if flask_version > (2, 0):
-                json_data = request.json
-            else:
-                json_data = json.loads(request.data)
+            json_data = request.json
             value = json_data.get("json_body")
             assert value == "master"
 
@@ -807,10 +803,7 @@ def sqli_11():
             con = sqlite3.connect(":memory:")
             cur = con.cursor()
 
-            if flask_version > (2, 0):
-                json_data = request.json
-            else:
-                json_data = json.loads(request.data)
+            json_data = request.json
             value = json_data.get("body").get("body2").get("body3")
             assert value == "master"
             assert is_pyobject_tainted(value)
@@ -873,10 +866,7 @@ def sqli_11():
             con = sqlite3.connect(":memory:")
             cur = con.cursor()
 
-            if flask_version > (2, 0):
-                json_data = request.json
-            else:
-                json_data = json.loads(request.data)
+            json_data = request.json
             value = json_data.get("body").get("body2").get("body3")[3]
             assert value == "master"
             assert is_pyobject_tainted(value)
@@ -939,10 +929,7 @@ def sqli_11():
             con = sqlite3.connect(":memory:")
             cur = con.cursor()
 
-            if flask_version > (2, 0):
-                json_data = request.json
-            else:
-                json_data = json.loads(request.data)
+            json_data = request.json
             value = json_data.get("body").get("body2").get("body3")[3].get("body4")
             assert value == "master"
             assert is_pyobject_tainted(value)
@@ -1018,10 +1005,7 @@ def iterate_json(data, parent_key=""):
                     else:
                         assert not is_pyobject_tainted(data), f"{parent_key}.{data} taint error"
 
-            if flask_version > (2, 0):
-                request_json = request.json
-            else:
-                request_json = json.loads(request.data)
+            request_json = request.json
 
             iterate_json(request_json)
 
@@ -1139,7 +1123,6 @@ def test_flask_request_body_iast_and_appsec(self):
 
         @self.app.route("/sqli/body/", methods=("POST",))
         def sqli_10():
-            import json
             import sqlite3
 
             from flask import request
@@ -1148,10 +1131,7 @@ def sqli_10():
 
             con = sqlite3.connect(":memory:")
             cur = con.cursor()
-            if flask_version > (2, 0):
-                json_data = request.json
-            else:
-                json_data = json.loads(request.data)
+            json_data = request.json
             value = json_data.get("json_body")
             assert value == "master"
 

tests/internal/test_serverless.py

@@ -46,7 +46,7 @@ def test_not_azure_function():
     "ddtrace.appsec._iast._ast.iastpatch",
     "ddtrace.appsec._iast._taint_tracking._native",
     "ddtrace.appsec._iast._taint_tracking._vendor",
-    "ddtrace.appsec._iast._stacktrace",
+    "ddtrace.appsec._shared._stacktrace",
     "ddtrace.internal.datadog.profiling.libdd_wrapper",
     "ddtrace.internal.datadog.profiling.ddup._ddup",
     "ddtrace.internal.datadog.profiling.stack._stack",