chore(iast): add more leak testing to CI

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 "
+                "{percent_increase:.2f}% which is greater than {fail_percent}%"
+            )
+            return 1
+        else:
+            print(
+                f"Success: memory increase is {percent_increase:.2f}% from half-point ({half_iterations} "
+                "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_memcheck/test_iast_mem_check.py

@@ -191,3 +191,10 @@ def test_stacktrace_memory_random_string_check():
         file_name, line_number = frame_info
         assert file_name == ""
         assert line_number == -1
+
+
+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, 500) == 0