diff --git a/ddtrace/appsec/_iast/__init__.py b/ddtrace/appsec/_iast/__init__.py index bf9a4f5a21c..1d14219d042 100644 --- a/ddtrace/appsec/_iast/__init__.py +++ b/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: diff --git a/ddtrace/appsec/_iast/_stacktrace.pyi b/ddtrace/appsec/_iast/_stacktrace.pyi deleted file mode 100644 index 46c1e435ac0..00000000000 --- a/ddtrace/appsec/_iast/_stacktrace.pyi +++ /dev/null @@ -1 +0,0 @@ -def get_info_frame(): ... diff --git a/ddtrace/appsec/_patch_utils.py b/ddtrace/appsec/_patch_utils.py index d767347808c..433fec3ed89 100644 --- a/ddtrace/appsec/_patch_utils.py +++ b/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 diff --git a/ddtrace/appsec/_shared/__init__.py b/ddtrace/appsec/_shared/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/ddtrace/appsec/_iast/_stacktrace.c b/ddtrace/appsec/_shared/_stacktrace.c similarity index 99% rename from ddtrace/appsec/_iast/_stacktrace.c rename to ddtrace/appsec/_shared/_stacktrace.c index 27517e023cf..2db311b88c6 100644 --- a/ddtrace/appsec/_iast/_stacktrace.c +++ b/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 }; diff --git a/ddtrace/appsec/_shared/_stacktrace.pyi b/ddtrace/appsec/_shared/_stacktrace.pyi new file mode 100644 index 00000000000..59674219854 --- /dev/null +++ b/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]]: ... diff --git a/setup.py b/setup.py index 21a63f50b0a..7870eced466 100644 --- a/setup.py +++ b/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, ) diff --git a/tests/appsec/architectures/test_appsec_loading_modules.py b/tests/appsec/architectures/test_appsec_loading_modules.py index ed4c38b9778..a4e64bdc0aa 100644 --- a/tests/appsec/architectures/test_appsec_loading_modules.py +++ b/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", ] diff --git a/tests/appsec/iast/test_stacktrace.py b/tests/appsec/iast/test_stacktrace.py index ab9ff658bae..729c2276f7f 100644 --- a/tests/appsec/iast/test_stacktrace.py +++ b/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(): diff --git a/tests/appsec/iast_memcheck/fixtures/stacktrace.py b/tests/appsec/iast_memcheck/fixtures/stacktrace.py index 6ba133dcf27..bb6b021fe03 100644 --- a/tests/appsec/iast_memcheck/fixtures/stacktrace.py +++ b/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()) diff --git a/tests/appsec/iast_memcheck/test_iast_mem_check.py b/tests/appsec/iast_memcheck/test_iast_mem_check.py index e93fd3e24fd..ea31a106ed9 100644 --- a/tests/appsec/iast_memcheck/test_iast_mem_check.py +++ b/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 diff --git a/tests/appsec/integrations/flask_tests/test_iast_flask.py b/tests/appsec/integrations/flask_tests/test_iast_flask.py index 68d4e3a105e..101b9c7fd00 100644 --- a/tests/appsec/integrations/flask_tests/test_iast_flask.py +++ b/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" diff --git a/tests/internal/test_serverless.py b/tests/internal/test_serverless.py index 5e3ae6535a3..ff1da9ade19 100644 --- a/tests/internal/test_serverless.py +++ b/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",