chore(iast): add more leak testing to CI

.gitlab/tests/appsec.yml

@@ -52,6 +52,14 @@ appsec threats fastapi:
     SUITE_NAME: "appsec_threats_fastapi"
   retry: 2
 
+appsec aggregated leak testing:
+  extends: .test_base_hatch
+  parallel: 6
+  variables:
+    SUITE_NAME: "appsec_aggregated_leak_testing"
+  retry: 2
+  timeout: 25m
+
 appsec iast native:
   extends: .test_base_hatch
   parallel: 6

hatch.toml

@@ -309,6 +309,31 @@ test = [
     "DD_IAST_ENABLED=true DD_IAST_REQUEST_SAMPLING=100 python -m pytest tests/appsec/contrib_appsec/test_fastapi.py"
 ]
 
+# ASM Appsec Aggregated Leak Testing
+
+[envs.appsec_aggregated_leak_testing]
+template = "appsec_aggregated_leak_testing"
+dependencies = [
+    "pytest",
+    "pytest-cov",
+    "hypothesis",
+    "requests",
+]
+
+[envs.appsec_aggregated_leak_testing.env-vars]
+CMAKE_BUILD_PARALLEL_LEVEL = "12"
+
+[envs.appsec_aggregated_leak_testing.scripts]
+test = [
+    "uname -a",
+    "pip freeze",
+    "python -m pytest tests/appsec/iast_aggregated_memcheck/test_aggregated_memleaks.py",
+]
+
+[[envs.appsec_aggregated_leak_testing.matrix]]
+python = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
+
+
 # if you add or remove a version here, please also update the parallelism parameter
 # in .circleci/config.templ.yml
 [[envs.appsec_threats_fastapi.matrix]]

scripts/iast/mod_leak_functions.py

@@ -4,23 +4,19 @@
 
 import requests
 
-from ddtrace.appsec._iast._utils import _is_iast_enabled
+from tests.utils import override_env
 
 
-if _is_iast_enabled():
+with override_env({"DD_IAST_ENABLED": "True"}):
     from ddtrace.appsec._iast._taint_tracking import OriginType
     from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 
 
 def test_doit():
     origin_string1 = "hiroot"
-
-    if _is_iast_enabled():
-        tainted_string_2 = taint_pyobject(
-            pyobject="1234", source_name="abcdefghijk", source_value="1234", source_origin=OriginType.PARAMETER
-        )
-    else:
-        tainted_string_2 = "1234"
+    tainted_string_2 = taint_pyobject(
+        pyobject="1234", source_name="abcdefghijk", source_value="1234", source_origin=OriginType.PARAMETER
+    )
 
     string1 = str(origin_string1)  # String with 1 propagation range
     string2 = str(tainted_string_2)  # String with 1 propagation range
@@ -71,6 +67,5 @@ def test_doit():
     string19 = os.path.normcase(string18)  # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted
     string20 = os.path.splitdrive(string19)[1]  # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted
 
-    expected = "notainted_HIROOT1234-HIROOT123_notainted"  # noqa: F841
-    # assert string20 == expected
+    # expected = "notainted_HIROOT1234-HIROOT123_notainted"  # noqa: F841
     return string20

scripts/iast/test_leak_functions.py

@@ -1,35 +1,80 @@
-import ddtrace.auto  # noqa: F401  # isort: skip
+import argparse
 import resource
 import sys
 
-from mod_leak_functions import test_doit
+from tests.appsec.iast.aspects.conftest import _iast_patched_module
+from tests.utils import override_env
 
-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
 
+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
 
-def test_main():
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Memory leak test script.")
+    parser.add_argument("--iterations", type=int, default=100000, help="Number of iterations.")
+    parser.add_argument(
+        "--fail_percent", type=float, default=2.0, help="Failure threshold for memory increase percentage."
+    )
+    parser.add_argument("--print_every", type=int, default=250, help="Print status every N iterations.")
+    return parser.parse_args()
+
+
+def test_iast_leaks(iterations: int, fail_percent: float, print_every: int):
+    if iterations < 100000:
+        print(
+            "Warning: running with %d iterations. At least 100.000 are recommended to stabilize the RSS info"
+            % iterations
+        )
     try:
-        rounds = int(sys.argv[1])
-    except ValueError:
-        rounds = 1
-    print("Test %d rounds" % rounds)
-    for i in range(rounds):
-        try:
+        half_iterations = iterations // 2
+        print("Test %d iterations" % iterations)
+        current_rss = 0
+        half_rss = 0
+
+        mod = _iast_patched_module("scripts.iast.mod_leak_functions")
+        test_doit = mod.test_doit
+
+        for i in range(iterations):
             create_context()
             result = test_doit()  # noqa: F841
+            assert result == "notainted_HIROOT1234-HIROOT123_notainted"  # noqa: F841
             assert is_pyobject_tainted(result)
             reset_context()
-        except KeyboardInterrupt:
-            print("Control-C stopped at %d rounds" % i)
-            break
-        if i % 250 == 0:
-            print("Round %d Max RSS: " % i, end="")
-            print(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024)
-    print("Round %d Max RSS: " % rounds, end="")
-    print(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024)
+
+            if i == half_iterations:
+                half_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
+
+            current_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
+
+            if i % print_every == 0:
+                print(f"Round {i} Max RSS: {current_rss}")
+
+        final_rss = current_rss
+
+        print(f"Round {iterations} Max RSS: {final_rss}")
+
+        percent_increase = ((final_rss - half_rss) / half_rss) * 100
+        if percent_increase > fail_percent:
+            print(
+                f"Failed: memory increase from half-point ({half_iterations} iterations) is "
+                f"{percent_increase:.2f}% which is greater than {fail_percent}%"
+            )
+            return 1
+
+        print(
+            f"Success: memory increase is {percent_increase:.2f}% from half-point ({half_iterations} "
+            f"iterations) which is less than {fail_percent}%"
+        )
+        return 0
+
+    except KeyboardInterrupt:
+        print("Test interrupted.")
 
 
 if __name__ == "__main__":
-    test_main()
+    args = parse_arguments()
+    with override_env({"DD_IAST_ENABLED": "True"}):
+        sys.exit(test_iast_leaks(args.iterations, args.fail_percent, args.print_every))

tests/appsec/iast_aggregated_memcheck/test_aggregated_memleaks.py

@@ -0,0 +1,8 @@
+from tests.utils import override_env
+
+
+def test_aggregated_leaks():
+    with override_env({"DD_IAST_ENABLED": "True"}):
+        from scripts.iast.test_leak_functions import test_iast_leaks
+
+        assert test_iast_leaks(100000, 2.0, 100) == 0