diff --git a/.gitlab/tests/appsec.yml b/.gitlab/tests/appsec.yml index ed2acaf7e1a..be67ac46e59 100644 --- a/.gitlab/tests/appsec.yml +++ b/.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 diff --git a/ddtrace/appsec/_iast/__init__.py b/ddtrace/appsec/_iast/__init__.py index e7f73adf313..8b3208baa86 100644 --- a/ddtrace/appsec/_iast/__init__.py +++ b/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", ] diff --git a/ddtrace/appsec/_iast/_ast/ast_patching.py b/ddtrace/appsec/_iast/_ast/ast_patching.py index 9aa7298538d..572fca02ce6 100644 --- a/ddtrace/appsec/_iast/_ast/ast_patching.py +++ b/ddtrace/appsec/_iast/_ast/ast_patching.py @@ -300,11 +300,6 @@ "uvicorn.", "anyio.", "httpcore.", - "pypika.", - "pydantic.", - "pydantic_core.", - "pydantic_settings.", - "tomli.", ) diff --git a/ddtrace/debugging/_function/discovery.py b/ddtrace/debugging/_function/discovery.py index 5e92ad05a13..9cabb4b3a04 100644 --- a/ddtrace/debugging/_function/discovery.py +++ b/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): diff --git a/ddtrace/internal/module.py b/ddtrace/internal/module.py index 2c607990a25..6b5b8192773 100644 --- a/ddtrace/internal/module.py +++ b/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 diff --git a/hatch.toml b/hatch.toml index c582f8f6a9d..ca74e784da6 100644 --- a/hatch.toml +++ b/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 = [ diff --git a/releasenotes/notes/profiling-fix-stack-v2-endpoint-82a1e26366166b8d.yaml b/releasenotes/notes/profiling-fix-stack-v2-endpoint-82a1e26366166b8d.yaml index 0505c26e550..73813f7cd70 100644 --- a/releasenotes/notes/profiling-fix-stack-v2-endpoint-82a1e26366166b8d.yaml +++ b/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. diff --git a/scripts/iast/README b/scripts/iast/README index 1f01828478f..e442e8916cc 100644 --- a/scripts/iast/README +++ b/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 diff --git a/scripts/iast/test_leak_functions.py b/scripts/iast/leak_functions.py similarity index 54% rename from scripts/iast/test_leak_functions.py rename to scripts/iast/leak_functions.py index 1f30210bb8a..d85f2f49486 100644 --- a/scripts/iast/test_leak_functions.py +++ b/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))) diff --git a/scripts/iast/mod_leak_functions.py b/scripts/iast/mod_leak_functions.py index e55480a1d36..afaf79b13fd 100644 --- a/scripts/iast/mod_leak_functions.py +++ b/scripts/iast/mod_leak_functions.py @@ -1,57 +1,231 @@ +from datetime import datetime import os import random import re import subprocess +from typing import Optional +from typing import Tuple +import anyio +from pydantic import BaseModel +from pydantic import Field +from pydantic_core import SchemaValidator import requests -from tests.utils import override_env +from ddtrace.appsec._iast._taint_tracking import OriginType +from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted +from ddtrace.appsec._iast._taint_tracking import taint_pyobject -with override_env({"DD_IAST_ENABLED": "True"}): - from ddtrace.appsec._iast._taint_tracking import OriginType - from ddtrace.appsec._iast._taint_tracking import taint_pyobject +v = SchemaValidator( + { + "type": "typed-dict", + "fields": { + "name": { + "type": "typed-dict-field", + "schema": { + "type": "str", + }, + }, + "age": { + "type": "typed-dict-field", + "schema": { + "type": "int", + "ge": 18, + }, + }, + "is_developer": { + "type": "typed-dict-field", + "schema": { + "type": "default", + "schema": {"type": "bool"}, + "default": True, + }, + }, + }, + } +) -def test_doit(): - origin_string1 = "hiroot" - tainted_string_2 = taint_pyobject( - pyobject="1234", source_name="abcdefghijk", source_value="1234", source_origin=OriginType.PARAMETER +class AspectModel(BaseModel): + foo: str = "bar" + apple: int = 1 + + +class Aspectvalidation(BaseModel): + timestamp: datetime + tuple_strings: Tuple[str, str] + dictionary_strs: dict[str, str] + tag: Optional[str] = None + author: Optional[str] = None + favorited: Optional[str] = None + limit: int = Field(20, ge=1) + offset: int = Field(0, ge=0) + + +class ImSubClassOfAString(str): + def __add__(self, other): + return "ImNotAString.__add__!!" + other + + def __iadd__(self, other): + return "ImNotAString.__iadd__!!" + other + + +class ImNotAString: + def __add__(self, other): + return "ImNotAString.__add__!!" + other + + def __iadd__(self, other): + return "ImNotAString.__iadd__!!" + other + + +def add_variants(string_tainted, string_no_tainted) -> str: + new_string_tainted = string_tainted + string_no_tainted + im_not_a_string = ImNotAString() + # TODO(avara1986): it raises seg fault instead of TypeError: can only concatenate str (not "ImNotAString") to str + # try: + # new_string_no_tainted = string_no_tainted + im_not_a_string + # assert False + # except TypeError: + # pass + new_string_no_tainted = im_not_a_string + string_no_tainted + new_string_no_tainted = ImSubClassOfAString() + new_string_no_tainted + string_no_tainted + new_string_no_tainted += string_no_tainted + # TODO(avara1986): it raises seg fault instead of TypeError: can only concatenate str (not "ImNotAString") to str + # try: + # new_string_no_tainted += ImNotAString() + # assert False + # except TypeError: + # pass + + im_not_a_string += new_string_no_tainted + new_string_no_tainted += ImSubClassOfAString() + new_string_tainted += new_string_no_tainted + new_string_tainted += new_string_no_tainted + new_string_tainted += new_string_tainted + + new_bytes_no_tainted = bytes(string_no_tainted, encoding="utf-8") + bytes(string_no_tainted, encoding="utf-8") + new_bytes_no_tainted += bytes(string_no_tainted, encoding="utf-8") + new_bytes_no_tainted += bytes(string_no_tainted, encoding="utf-8") + new_bytes_tainted = bytes(string_tainted, encoding="utf-8") + bytes(string_tainted, encoding="utf-8") + new_bytes_tainted += bytes(string_no_tainted, encoding="utf-8") + new_bytes_tainted += bytes(string_tainted, encoding="utf-8") + new_bytes_tainted += bytes(string_tainted, encoding="utf-8") + new_bytearray_tainted = bytearray(bytes(string_tainted, encoding="utf-8")) + bytearray( + bytes(string_tainted, encoding="utf-8") ) + new_bytearray_tainted += bytearray(bytes(string_tainted, encoding="utf-8")) + new_bytearray_tainted += bytearray(bytes(string_no_tainted, encoding="utf-8")) - string1 = str(origin_string1) # String with 1 propagation range - string2 = str(tainted_string_2) # String with 1 propagation range + new_string_tainted = ( + new_string_tainted + + new_string_no_tainted + + str(new_bytes_no_tainted, encoding="utf-8") + + str(new_bytes_tainted, encoding="utf-8") + + str(new_bytearray_tainted, encoding="utf-8") + ) + # print(new_string_tainted) + return new_string_tainted - string3 = string1 + string2 # 2 propagation ranges: hiroot1234 - string4 = "-".join([string3, string3, string3]) # 6 propagation ranges: hiroot1234-hiroot1234-hiroot1234 - string4_2 = string1 - string4_2 += " " + " ".join(string_ for string_ in [string4, string4, string4]) - string5 = string4_2[0:20] # 1 propagation range: hiroot1234-hiroot123 - string6 = string5.title() # 1 propagation range: Hiroot1234-Hiroot123 - string7 = string6.upper() # 1 propagation range: HIROOT1234-HIROOT123 - string8 = "%s_notainted" % string7 # 1 propagation range: HIROOT1234-HIROOT123_notainted - string9 = "notainted_{}".format(string8) # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted - string10 = "nottainted\n" + string9 # 2 propagation ranges: notainted\nnotainted_HIROOT1234-HIROOT123_notainted - string11 = string10.splitlines()[1] # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted - string12 = string11 + "_notainted" # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted_notainted - string13 = string12.rsplit("_", 1)[0] # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted + +def format_variants(string_tainted, string_no_tainted) -> str: + string_tainted_2 = "My name is {} and I am {} years old.".format(string_tainted, 30) + string_tainted_3 = "My name is {name} and I am {age} years old.".format(name=string_tainted_2, age=25) + string_tainted_4 = "{0} is {1} years old. {0} lives in {2}.".format(string_tainted_3, 35, string_no_tainted) + string_tainted_5 = "|{:<10}|{:^10}|{:>10}|".format(string_tainted_4, string_no_tainted, string_tainted_4) + string_tainted_6 = "|{:-<10}|{:*^10}|{:.>10}|".format(string_no_tainted, string_tainted_5, string_no_tainted) + string_tainted_7 = "{} is approximately {:.3f}".format(string_tainted_6, 3.1415926535) + string_tainted_8 = "The {} is {:,}".format(string_tainted_7, 1000000) + string_tainted_9 = "{1} Hex: {0:x}, Bin: {0:b}, Oct: {0:o}".format(255, string_tainted_8) + string_tainted_10 = "{} Success rate: {:.2%}".format(string_tainted_9, 0.8765) + string_tainted_11 = "{} {:+d}, {:+d}".format(string_tainted_10, 42, -42) + return string_tainted_11 + + +def modulo_exceptions(string8_4): + # Validate we're not leaking in modulo exceptions + try: + string8_5 = "notainted_%s_" % (string8_4, string8_4) # noqa: F841, F507 + except TypeError: + pass + + try: + string8_5 = "notainted_%d" % string8_4 # noqa: F841 + except TypeError: + pass + + try: + string8_5 = "notainted_%s %s" % "abc", string8_4 # noqa: F841 + except TypeError: + pass + + try: + string8_5 = "notainted_%s %s" % "abc", string8_4 # noqa: F841 + except TypeError: + pass + + try: + string8_5 = "notainted_%s" % (string8_4) # noqa: F841 + except TypeError: + pass + try: + string8_5 = "notainted_%s %(name)s" % string8_4, {"name": string8_4} # noqa: F841, F506 + except TypeError: + pass + + try: + string8_5 = "notainted_%(age)d" % {"age": string8_4} # noqa: F841 + except TypeError: + pass + + +def pydantic_object(tag, string_tainted): + m = Aspectvalidation( + timestamp="2020-01-02T03:04:05Z", + tuple_strings=[string_tainted, string_tainted], + dictionary_strs={ + "wine": string_tainted, + b"cheese": string_tainted, + "cabbage": string_tainted, + }, + tag=string_tainted, + author=string_tainted, + favorited=string_tainted, + limit=20, + offset=0, + ) + + r1 = v.validate_python({"name": m.tuple_strings[0], "age": 35}) + + assert is_pyobject_tainted(m.tuple_strings[0]) + assert is_pyobject_tainted(m.tuple_strings[1]) + assert is_pyobject_tainted(m.dictionary_strs["cabbage"]) + + r2 = v.validate_json('{"name": "' + m.tuple_strings[0] + '", "age": 35}') + + assert r1 == {"name": m.tuple_strings[0], "age": 35, "is_developer": True} + assert r1 == r2 + return m + + +def sink_points(string_tainted): try: # Path traversal vulnerability - m = open("/" + string13 + ".txt") + m = open("/" + string_tainted + ".txt") _ = m.read() except Exception: pass try: # Command Injection vulnerability - _ = subprocess.Popen("ls " + string9) + _ = subprocess.Popen("ls " + string_tainted) except Exception: pass try: # SSRF vulnerability - requests.get("http://" + "foobar") + requests.get("http://" + string_tainted, timeout=1) # urllib3.request("GET", "http://" + "foobar") except Exception: pass @@ -59,16 +233,63 @@ def test_doit(): # Weak Randomness vulnerability _ = random.randint(1, 10) + +async def test_doit(): + sample_str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + + origin_string1 = "".join(random.choices(sample_str, k=5)) + + string_2 = "".join(random.choices(sample_str, k=20)) + tainted_string_2 = taint_pyobject( + pyobject=string_2, source_name="abcdefghijk", source_value=string_2, source_origin=OriginType.PARAMETER + ) + + string1 = str(origin_string1) # String with 1 propagation range + string2 = str(tainted_string_2) # String with 1 propagation range + + string3 = add_variants(string2, string1) + + string4 = "-".join([string3, string3, string3]) + string4_2 = string1 + string4_2 += " " + " ".join(string_ for string_ in [string4, string4, string4]) + string4_2 += " " + " ".join(string_ for string_ in [string1, string1, string1]) + + string4_2 += " " + " ".join([string_ for string_ in [string4_2, string4_2, string4_2]]) + + string5 = string4_2[0:100] + string6 = string5.title() + string7 = string6.upper() + string8 = "%s_notainted" % string7 + string8_2 = "%s_%s_notainted" % (string8, string8) + string8_3 = "notainted_%s_" + string8_2 + string8_4 = string8_3 % "notainted" + + string8_5 = format_variants(string8_4, string1) + await anyio.to_thread.run_sync(modulo_exceptions, string8_5) + + string9 = "notainted#{}".format(string8_5) + string9_2 = f"{string9}_notainted" + string9_3 = f"{string9_2:=^30}_notainted" + string10 = "nottainted\n" + string9_3 + string11 = string10.splitlines()[1] + string12 = string11 + "_notainted" + string13 = string12.rsplit("_", 1)[0] + string13_2 = string13 + " " + string13 + try: + string13_3, string13_5, string13_5 = string13_2.split(" ") + except ValueError: + pass + + sink_points(string13_2) + # os path propagation - string14 = os.path.join(string13, "a") # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted/a - string15 = os.path.split(string14)[0] # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted - string16 = os.path.dirname( - string15 + "/" + "foobar" - ) # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted - string17 = os.path.basename("/foobar/" + string16) # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted - string18 = os.path.splitext(string17 + ".jpg")[0] # 1 propagation range: notainted_HIROOT1234-HIROOT123_notainted - 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 + string14 = os.path.join(string13_2, "a") + string15 = os.path.split(string14)[0] + string16 = os.path.dirname(string15 + "/" + "foobar") + string17 = os.path.basename("/foobar/" + string16) + string18 = os.path.splitext(string17 + ".jpg")[0] + string19 = os.path.normcase(string18) + string20 = os.path.splitdrive(string19)[1] re_slash = re.compile(r"[_.][a-zA-Z]*") string21 = re_slash.findall(string20)[0] # 1 propagation: '_HIROOT @@ -97,4 +318,8 @@ def test_doit(): tmp_str2 = "_extend" string27 += tmp_str2 + # TODO(avara1986): pydantic is in the DENY_LIST, remove from it and uncomment this lines + # result = await anyio.to_thread.run_sync(functools.partial(pydantic_object, string_tainted=string27), string27) + # result = pydantic_object(tag="test2", string_tainted=string27) + # return result.tuple_strings[0] return string27 diff --git a/scripts/iast/requirements.txt b/scripts/iast/requirements.txt index 0fd2996a90d..e6ae243a21d 100644 --- a/scripts/iast/requirements.txt +++ b/scripts/iast/requirements.txt @@ -1 +1,4 @@ -memray==1.12.0 \ No newline at end of file +memray==1.12.0 +anyio +pydantic +pydantic-settings \ No newline at end of file diff --git a/scripts/iast/run_memory.sh b/scripts/iast/run_memory.sh index 0731643b865..a0dd73730eb 100644 --- a/scripts/iast/run_memory.sh +++ b/scripts/iast/run_memory.sh @@ -2,4 +2,6 @@ set -o xtrace PYTHON="${PYTHON_VERSION:-python3.11}" $PYTHON -m pip install -r scripts/iast/requirements.txt -$PYTHON scripts/iast/test_leak_functions.py --iterations 1000000 --print_every 250 \ No newline at end of file +export DD_IAST_ENABLED=true +export _DD_IAST_DEBUG=true +$PYTHON scripts/iast/leak_functions.py --iterations 1000000 --print_every 250 \ No newline at end of file diff --git a/scripts/iast/run_memray.sh b/scripts/iast/run_memray.sh index c785fb520e2..96982a2e513 100644 --- a/scripts/iast/run_memray.sh +++ b/scripts/iast/run_memray.sh @@ -2,5 +2,7 @@ set -o xtrace PYTHON="${PYTHON_VERSION:-python3.11d}" $PYTHON -m pip install -r scripts/iast/requirements.txt -$PYTHON -m memray run --trace-python-allocators --aggregate --native -o lel.bin -f scripts/iast/test_leak_functions.py --iterations 100 +export DD_IAST_ENABLED=true +export _DD_IAST_DEBUG=true +$PYTHON -m memray run --trace-python-allocators --aggregate --native -o lel.bin -f scripts/iast/leak_functions.py --iterations 100 # $PYTHON -m memray flamegraph lel.bin --leaks -f \ No newline at end of file diff --git a/scripts/iast/run_references.sh b/scripts/iast/run_references.sh index 071696b9d94..769771fc632 100644 --- a/scripts/iast/run_references.sh +++ b/scripts/iast/run_references.sh @@ -1,4 +1,7 @@ set -o xtrace PYTHON="${PYTHON_VERSION:-python3.11d}" +$PYTHON -m pip install -r scripts/iast/requirements.txt +export DD_IAST_ENABLED=true +export _DD_IAST_DEBUG=true ${PYTHON} -m ddtrace.commands.ddtrace_run ${PYTHON} scripts/iast/test_references.py \ No newline at end of file diff --git a/scripts/iast/test_references.py b/scripts/iast/test_references.py index e6b91bd3de0..d4cdd2fcc16 100644 --- a/scripts/iast/test_references.py +++ b/scripts/iast/test_references.py @@ -1,3 +1,4 @@ +import asyncio import gc import sys @@ -8,13 +9,13 @@ from ddtrace.appsec._iast._taint_tracking import reset_context -def test_main(): +async def test_main(): for i in range(1000): gc.collect() a = sys.gettotalrefcount() try: create_context() - result = test_doit() # noqa: F841 + result = await test_doit() # noqa: F841 assert is_pyobject_tainted(result) reset_context() except KeyboardInterrupt: @@ -25,4 +26,5 @@ def test_main(): if __name__ == "__main__": - test_main() + loop = asyncio.get_event_loop() + sys.exit(loop.run_until_complete(test_main())) diff --git a/tests/appsec/iast_aggregated_memcheck/test_aggregated_memleaks.py b/tests/appsec/iast_aggregated_memcheck/test_aggregated_memleaks.py index ce213cbe113..c64ca7504d2 100644 --- a/tests/appsec/iast_aggregated_memcheck/test_aggregated_memleaks.py +++ b/tests/appsec/iast_aggregated_memcheck/test_aggregated_memleaks.py @@ -1,8 +1,12 @@ +import pytest + 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 +@pytest.mark.asyncio +async def test_aggregated_leaks(): + with override_env({"DD_IAST_ENABLED": "true"}): + from scripts.iast.leak_functions import iast_leaks - assert test_iast_leaks(75000, 2.0, 100) == 0 + result = await iast_leaks(75000, 1.0, 100) == 0 + assert result