chore(iast): improve iast scripts

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():
+    """Add iast AST patching in 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.register_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/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,14 @@ dependencies = [
     "pytest-cov",
     "hypothesis",
     "requests",
+    "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 = [

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 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,7 +26,17 @@ def parse_arguments():
     return parser.parse_args()
 
 
-def test_iast_leaks(iterations: int, fail_percent: float, print_every: int):
+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):
     if iterations < 60000:
         print(
             "Error: not running with %d iterations. At least 60.000 are needed to stabilize the RSS info" % iterations
@@ -34,13 +48,18 @@ def test_iast_leaks(iterations: int, fail_percent: float, print_every: int):
         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
+        #  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
+            result = await test_doit()  # noqa: F841
             assert result == "DDD_III_extend", f"result is {result}"  # noqa: F841
             assert is_pyobject_tainted(result)
             reset_context()
@@ -52,11 +71,11 @@ 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}, {active_map_addreses_size()}")
 
         final_rss = current_rss
 
-        print(f"Round {iterations} Max RSS: {final_rss}")
+        print(f"Round {iterations} Max RSS: {final_rss}, {active_map_addreses_size()}")
 
         percent_increase = ((final_rss - half_rss) / half_rss) * 100
         if percent_increase > fail_percent:
@@ -74,9 +93,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)))