Merge branch 'main' into munir/internalize-some-tracing-config-stuff-2

.gitlab/tests/appsec.yml

@@ -61,7 +61,7 @@ appsec aggregated leak testing:
   variables:
     SUITE_NAME: "appsec_aggregated_leak_testing"
   retry: 2
-  timeout: 25m
+  timeout: 35m
 
 appsec iast native:
   extends: .test_base_hatch

ddtrace/appsec/_iast/__init__.py

@@ -35,6 +35,7 @@ def wrapped_function(wrapped, instance, args, kwargs):
 from ddtrace.internal.module import ModuleWatchdog
 from ddtrace.internal.utils.formats import asbool
 
+from .._constants import IAST
 from ._overhead_control_engine import OverheadControl
 from ._utils import _is_iast_enabled
 
@@ -71,7 +72,8 @@ def ddtrace_iast_flask_patch():
 
 
 def enable_iast_propagation():
-    if asbool(os.getenv("DD_IAST_ENABLED", False)):
+    """Add IAST AST patching in the ModuleWatchdog"""
+    if asbool(os.getenv(IAST.ENV, "false")):
         from ddtrace.appsec._iast._utils import _is_python_version_supported
 
         if _is_python_version_supported():
@@ -82,8 +84,20 @@ def enable_iast_propagation():
             ModuleWatchdog.register_pre_exec_module_hook(_should_iast_patch, _exec_iast_patched_module)
 
 
+def disable_iast_propagation():
+    """Remove IAST AST patching from the ModuleWatchdog. Only for testing proposes"""
+    from ddtrace.appsec._iast._ast.ast_patching import _should_iast_patch
+    from ddtrace.appsec._iast._loader import _exec_iast_patched_module
+
+    try:
+        ModuleWatchdog.remove_pre_exec_module_hook(_should_iast_patch, _exec_iast_patched_module)
+    except KeyError:
+        log.warning("IAST is already disabled and it's not in the ModuleWatchdog")
+
+
 __all__ = [
     "oce",
     "ddtrace_iast_flask_patch",
     "enable_iast_propagation",
+    "disable_iast_propagation",
 ]

ddtrace/appsec/_iast/_ast/ast_patching.py

@@ -300,11 +300,6 @@
     "uvicorn.",
     "anyio.",
     "httpcore.",
-    "pypika.",
-    "pydantic.",
-    "pydantic_core.",
-    "pydantic_settings.",
-    "tomli.",
 )
 
 

ddtrace/debugging/_function/discovery.py

@@ -66,7 +66,9 @@ def __init__(
         origin: Optional[Union[Tuple["ContainerIterator", ContainerKey], Tuple[FullyNamedFunction, str]]] = None,
     ) -> None:
         if isinstance(container, (type, ModuleType)):
-            self._iter = iter(container.__dict__.items())
+            # DEV: A module object could be partially initialised, therefore
+            # __dict__ can mutate.
+            self._iter = iter(container.__dict__.copy().items())
             self.__name__ = container.__name__
 
         elif isinstance(container, tuple):

ddtrace/internal/module.py

@@ -670,6 +670,14 @@ def register_pre_exec_module_hook(
         instance = t.cast(ModuleWatchdog, cls._instance)
         instance._pre_exec_module_hooks.add((cond, hook))
 
+    @classmethod
+    def remove_pre_exec_module_hook(
+        cls: t.Type["ModuleWatchdog"], cond: PreExecHookCond, hook: PreExecHookType
+    ) -> None:
+        """Register a hook to execute before/instead of exec_module. Only for testing proposes"""
+        instance = t.cast(ModuleWatchdog, cls._instance)
+        instance._pre_exec_module_hooks.remove((cond, hook))
+
     @classmethod
     def register_import_exception_hook(
         cls: t.Type["ModuleWatchdog"], cond: ImportExceptionHookCond, hook: ImportExceptionHookType

hatch.toml

@@ -339,10 +339,15 @@ dependencies = [
     "pytest-cov",
     "hypothesis",
     "requests",
+    "pytest-asyncio",
+    "anyio",
+    "pydantic",
+    "pydantic-settings",
 ]
 
 [envs.appsec_aggregated_leak_testing.env-vars]
 CMAKE_BUILD_PARALLEL_LEVEL = "12"
+DD_IAST_ENABLED = "true"
 
 [envs.appsec_aggregated_leak_testing.scripts]
 test = [

releasenotes/notes/profiling-fix-stack-v2-endpoint-82a1e26366166b8d.yaml

@@ -1,6 +1,6 @@
 ---
 fixes:
   - |
-    profiling: enables endpoint profiling for stack v2, ``DD_PROFILING_STACK_V2_ENABLED``
+    profiling: fixes endpoint profiling for stack v2, when ``DD_PROFILING_STACK_V2_ENABLED``
     is set.
 

scripts/iast/README

@@ -53,7 +53,7 @@ sh scripts/iast/run_memory.sh"
 docker run --rm -it -v ${PWD}:/ddtrace python_311_debug /bin/bash -c "cd /ddtrace && source scripts/iast/.env && \
 valgrind --tool=memcheck --leak-check=full  --log-file=scripts/iast/valgrind_bench_overload.out --track-origins=yes \
 --suppressions=scripts/iast/valgrind-python.supp  --show-leak-kinds=all \
-python3.11 scripts/iast/test_leak_functions.py --iterations 100"
+python3.11 scripts/iast/leak_functions.py  --iterations 100"
 
 ##### Understanding results of memcheck
 

scripts/iast/test_leak_functions.py → scripts/iast/leak_functions.py

@@ -1,15 +1,19 @@
 import argparse
+import asyncio
+import dis
+import io
 import resource
 import sys
 
-from tests.appsec.iast.aspects.conftest import _iast_patched_module
-from tests.utils import override_env
-
+import pytest
 
-with override_env({"DD_IAST_ENABLED": "True"}):
-    from ddtrace.appsec._iast._taint_tracking import create_context
-    from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
-    from ddtrace.appsec._iast._taint_tracking import reset_context
+from ddtrace.appsec._iast import disable_iast_propagation
+from ddtrace.appsec._iast import enable_iast_propagation
+from ddtrace.appsec._iast._taint_tracking import active_map_addreses_size
+from ddtrace.appsec._iast._taint_tracking import create_context
+from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
+from ddtrace.appsec._iast._taint_tracking import reset_context
+from tests.utils import override_env
 
 
 def parse_arguments():
@@ -22,26 +26,41 @@ def parse_arguments():
     return parser.parse_args()
 
 
-def test_iast_leaks(iterations: int, fail_percent: float, print_every: int):
-    if iterations < 60000:
+def _pre_checks(module, aspect_to_check="add_aspect"):
+    """Ensure the module code is replaced by IAST patching. To do that, this function inspects the bytecode"""
+    dis_output = io.StringIO()
+    dis.dis(module, file=dis_output)
+    str_output = dis_output.getvalue()
+    # Should have replaced the binary op with the aspect in add_test:
+    assert f"({aspect_to_check})" in str_output
+
+
+@pytest.mark.asyncio
+async def iast_leaks(iterations: int, fail_percent: float, print_every: int):
+    mem_reference_iterations = 50000
+    if iterations < mem_reference_iterations:
         print(
             "Error: not running with %d iterations. At least 60.000 are needed to stabilize the RSS info" % iterations
         )
         sys.exit(1)
 
     try:
-        mem_reference_iterations = 50000
         print("Test %d iterations" % iterations)
         current_rss = 0
         half_rss = 0
+        enable_iast_propagation()
+        from scripts.iast.mod_leak_functions import test_doit
+
+        # TODO(avara1986): pydantic is in the DENY_LIST, remove from it and uncomment this lines
+        #  from pydantic import main
+        #  _pre_checks(main, "index_aspect")
 
-        mod = _iast_patched_module("scripts.iast.mod_leak_functions")
-        test_doit = mod.test_doit
+        _pre_checks(test_doit)
 
         for i in range(iterations):
             create_context()
-            result = test_doit()  # noqa: F841
-            assert result == "DDD_III_extend", f"result is {result}"  # noqa: F841
+            result = await test_doit()
+            assert result == "DDD_III_extend", f"result is {result}"
             assert is_pyobject_tainted(result)
             reset_context()
 
@@ -52,11 +71,13 @@ def test_iast_leaks(iterations: int, fail_percent: float, print_every: int):
             current_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
 
             if i % print_every == 0:
-                print(f"Round {i} Max RSS: {current_rss}")
+                print(
+                    f"Round {i} Max RSS: {current_rss}, Number of active maps addresses: {active_map_addreses_size()}"
+                )
 
         final_rss = current_rss
 
-        print(f"Round {iterations} Max RSS: {final_rss}")
+        print(f"Round {iterations} Max RSS: {final_rss}, Number of active maps addresses: {active_map_addreses_size()}")
 
         percent_increase = ((final_rss - half_rss) / half_rss) * 100
         if percent_increase > fail_percent:
@@ -74,9 +95,12 @@ def test_iast_leaks(iterations: int, fail_percent: float, print_every: int):
 
     except KeyboardInterrupt:
         print("Test interrupted.")
+    finally:
+        disable_iast_propagation()
 
 
 if __name__ == "__main__":
+    loop = asyncio.get_event_loop()
     args = parse_arguments()
     with override_env({"DD_IAST_ENABLED": "True"}):
-        sys.exit(test_iast_leaks(args.iterations, args.fail_percent, args.print_every))
+        sys.exit(loop.run_until_complete(iast_leaks(args.iterations, args.fail_percent, args.print_every)))