chore(iast): optimize iast by migrating should_iast_patch function to c

.gitlab/benchmarks/microbenchmarks.yml

@@ -16,7 +16,7 @@ variables:
   tags: ["runner:apm-k8s-tweaked-metal"]
   image: $MICROBENCHMARKS_CI_IMAGE
   interruptible: true
-  timeout: 1h
+  timeout: 2h
   variables:
     CMAKE_BUILD_PARALLEL_LEVEL: 12
     CARGO_BUILD_JOBS: 12
@@ -144,8 +144,9 @@ microbenchmarks:
         - "flask_sqli"
         - "core_api"
         - "otel_span"
-        - "appsec_iast_propagation"
         - "appsec_iast_aspects"
+        - "appsec_iast_django_startup"
+        - "appsec_iast_propagation"
         # Flaky. Timeout errors
         # - "encoder"
         - "http_propagation_extract"

benchmarks/appsec_iast_django_startup/config.yaml

@@ -0,0 +1,22 @@
+baseline: &baseline
+  tracer_enabled: false
+  tracer_debug: false
+  appsec_enabled: false
+  iast_enabled: false
+tracer: &tracer
+  <<: *baseline
+  tracer_enabled: true
+appsec:
+  <<: *tracer
+  appsec_enabled: true
+iast:
+  <<: *tracer
+  iast_enabled: true
+iast_debug_mode:
+  <<: *tracer
+  tracer_debug: true
+  iast_enabled: true
+appsec_iast:
+  <<: *tracer
+  appsec_enabled: true
+  iast_enabled: true

benchmarks/appsec_iast_django_startup/install.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+apt-get update && apt-get install --no-install-recommends -y \
+  sqlite3 libgl1 ffmpeg libsm6 libxext6 python3-opencv \
+  && apt-get purge -y --auto-remove -o APT::AutoRemove::RecommendsImportant=false \
+  && rm -rf /var/lib/apt/lists/*

benchmarks/appsec_iast_django_startup/manage.py

@@ -0,0 +1,33 @@
+import os
+import sys
+
+# TODO: CV2 fails in the CI with ImportError: libGL.so.1: cannot open shared object file: No such file or directory
+# from cv2 import *  # noqa: F401, F403
+from django import *  # noqa: F401, F403
+from scipy import *  # noqa: F401, F403
+
+
+# TODO: we want to check django first
+# from kombu import *  # noqa: F401, F403
+# from matplotlib import *  # noqa: F401, F403
+# from nibabel import *  # noqa: F401, F403
+# from pandas import *  # noqa: F401, F403
+# from PIL import *  # noqa: F401, F403
+
+
+def main():
+    """Run administrative tasks."""
+    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings")
+    try:
+        from django.core.management import execute_from_command_line
+    except ImportError as exc:
+        raise ImportError(
+            "Couldn't import Django. Are you sure it's installed and "
+            "available on your PYTHONPATH environment variable? Did you "
+            "forget to activate a virtual environment?"
+        ) from exc
+    execute_from_command_line(sys.argv)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/appsec_iast_django_startup/requirements_scenario.txt

@@ -0,0 +1,10 @@
+django==4.2.5
+kombu==5.4.2
+matplotlib==3.8.1
+nibabel==5.3.2
+numpy==1.26.4
+# opencv-python==4.10.0.82
+pandas==2.2.2
+pillow==10.4.0
+scipy==1.13.1
+requests==2.31.0

benchmarks/appsec_iast_django_startup/scenario.py

@@ -0,0 +1,70 @@
+import os
+import subprocess
+
+import bm
+import requests
+import tenacity
+
+
+SERVER_URL = "http://0.0.0.0:8000"
+
+
+@tenacity.retry(
+    wait=tenacity.wait_fixed(0.2),
+    stop=tenacity.stop_after_attempt(1570),
+)
+def _get_response(path="/"):
+    r = requests.get(SERVER_URL + path)
+    r.raise_for_status()
+
+
+def server(scenario):
+    env = {
+        "DD_TRACE_DEBUG": str(scenario.tracer_debug),
+        "DD_ENV": "prod",
+        "_DD_IAST_PATCH_MODULES": "benchmarks.,tests.appsec.",
+        "DD_VERSION": "1.0",
+        "DD_APPSEC_ENABLED": str(scenario.appsec_enabled),
+        "DD_IAST_ENABLED": str(scenario.iast_enabled),
+    }
+    # copy over current environ
+    env.update(os.environ)
+    if scenario.tracer_enabled:
+        cmd = ["python", "-m", "ddtrace.commands.ddtrace_run", "python", "manage.py", "runserver", "--noreload"]
+    else:
+        cmd = ["python", "manage.py", "runserver", "--noreload"]
+    proc = subprocess.Popen(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, start_new_session=True, text=True, env=env
+    )
+    # make sure process has been started
+    try:
+        # print("Starting server....")
+        _get_response()
+        # print("Shutdown server....")
+        _get_response(path="/shutdown")
+    except Exception as e:
+        proc.terminate()
+        proc.wait()
+        raise AssertionError(
+            "Server failed to start, see stdout and stderr logs\n."
+            "Exception %s\n."
+            "\n=== Captured STDOUT ===\n%s=== End of captured STDOUT ==="
+            "\n=== Captured STDERR ===\n%s=== End of captured STDERR ===" % (e, proc.stdout.read(), proc.stderr.read())
+        )
+
+    proc.terminate()
+    proc.wait()
+
+
+class IASTDjangoStartup(bm.Scenario):
+    tracer_enabled: bool
+    tracer_debug: bool
+    appsec_enabled: bool
+    iast_enabled: bool
+
+    def run(self):
+        def _(loops):
+            for _ in range(loops):
+                server(self)
+
+        yield _

benchmarks/appsec_iast_django_startup/settings.py

@@ -0,0 +1,22 @@
+import os
+
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+DEBUG = False
+ROOT_URLCONF = "urls"
+DATABASES = {
+    "default": {
+        "ENGINE": "django.db.backends.sqlite3",
+        "NAME": ":memory:",
+    }
+}
+TEMPLATES = [
+    {
+        "BACKEND": "django.template.backends.django.DjangoTemplates",
+        "DIRS": [
+            BASE_DIR,
+        ],
+    }
+]
+SECRET_KEY = ("SECRET",)
+ALLOWED_HOSTS = ["*"]

benchmarks/appsec_iast_django_startup/urls.py

@@ -0,0 +1,6 @@
+from django.urls import re_path
+from views import index
+from views import shutdown_view
+
+
+urlpatterns = [re_path(r"^$", index), re_path(r"^shutdown", shutdown_view, name="response-header")]

benchmarks/appsec_iast_django_startup/views.py

@@ -0,0 +1,28 @@
+from django.http import HttpResponse
+from django.template import Context
+from django.template import Template
+
+from ddtrace import tracer
+
+
+def index(request):
+    # render a large table template
+    index = Template(
+        """
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+    <title>Django Simple</title>
+  </head>
+  <body>
+    <p>Hello</p>
+  </body>
+</html>
+    """
+    )
+    return HttpResponse(index.render(Context({})))
+
+
+def shutdown_view(request):
+    tracer._writer.flush_queue()
+    return HttpResponse("SHUTDOWN")

ddtrace/appsec/_iast/_ast/ast_patching.py

@@ -149,7 +149,6 @@
     "cycler.",
     "cython.",
     "dateutil.",
-    "dateutil.",
     "ddsketch.",
     "ddtrace.",
     "defusedxml.",
@@ -471,7 +470,7 @@ def _is_first_party(module_name: str):
 
 def _should_iast_patch(module_name: Text) -> bool:
     """
-    select if module_name should be patch from the longest prefix that match in allow or deny list.
+    select if module_name should be patched from the longest prefix that match in allow or deny list.
     if a prefix is in both list, deny is selected.
     """
     # TODO: A better solution would be to migrate the original algorithm to C++:
@@ -510,7 +509,7 @@ def _should_iast_patch(module_name: Text) -> bool:
 
 
 def visit_ast(
-    source_text: Text,
+    source_text: bytes,
     module_path: Text,
     module_name: Text = "",
 ) -> Optional[ast.Module]:
@@ -555,7 +554,7 @@ def {_PREFIX}set_dir_filter():
 {_PREFIX}set_dir_filter()
 
     """
-)
+).encode()
 
 
 def astpatch_module(module: ModuleType) -> Tuple[str, Optional[ast.Module]]:
@@ -586,7 +585,7 @@ def astpatch_module(module: ModuleType) -> Tuple[str, Optional[ast.Module]]:
         iast_compiling_debug_log(f"Extension not supported: {module_ext} for: {module_path}")
         return "", None
 
-    with open(module_path, "r", encoding=get_encoding(module_path)) as source_file:
+    with open(module_path, "rb") as source_file:
         try:
             source_text = source_file.read()
         except UnicodeDecodeError:

ddtrace/appsec/_iast/_ast/visitor.py

@@ -474,8 +474,7 @@ def visit_Module(self, module_node: ast.Module) -> Any:
         module_node.body.insert(insert_position, replacements_import)
         # Must be called here instead of the start so the line offset is already
         # processed
-        self.generic_visit(module_node)
-        return module_node
+        return self.generic_visit(module_node)
 
     def visit_FunctionDef(self, def_node: ast.FunctionDef) -> Any:
         """

ddtrace/appsec/_iast/_handlers.py

@@ -238,9 +238,7 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
                         v, source_name=k, source_value=v, source_origin=OriginType.PATH_PARAMETER
                     )
             except Exception:
-                iast_propagation_listener_log_log(
-                    "IAST: Unexpected exception while tainting path parameters", exc_info=True
-                )
+                iast_propagation_listener_log_log("Unexpected exception while tainting path parameters", exc_info=True)
 
 
 def _custom_protobuf_getattribute(self, name):

ddtrace/appsec/_iast/_loader.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+import ctypes
 
 from ddtrace.appsec._iast._logs import iast_compiling_debug_log
 from ddtrace.internal.logger import get_logger
@@ -7,6 +7,10 @@
 from ._ast.ast_patching import astpatch_module
 
 
+PyEval_EvalCode = ctypes.pythonapi.PyEval_EvalCode
+PyEval_EvalCode.argtypes = [ctypes.py_object, ctypes.py_object, ctypes.py_object]
+PyEval_EvalCode.restype = ctypes.py_object
+
 log = get_logger(__name__)
 
 IS_IAST_ENABLED = asm_config._iast_enabled
@@ -33,7 +37,10 @@ def _exec_iast_patched_module(module_watchdog, module):
     if compiled_code:
         iast_compiling_debug_log(f"INSTRUMENTED CODE. executing {module_path}")
         # Patched source is executed instead of original module
-        exec(compiled_code, module.__dict__)  # nosec B102
+        # exec(compiled_code, module.__dict__)  # nosec B102
+        if "__builtins__" not in module.__dict__:
+            module.__dict__["__builtins__"] = __builtins__
+        PyEval_EvalCode(compiled_code, module.__dict__, module.__dict__)
     elif module_watchdog.loader is not None:
         try:
             iast_compiling_debug_log(f"DEFAULT CODE. executing {module}")

tests/appsec/iast/_ast/test_ast_patching.py

@@ -70,6 +70,7 @@ def test_visit_ast_changed(source_text, module_path, module_name):
     [
         ("tests.appsec.iast.fixtures.ast.str.class_str"),
         ("tests.appsec.iast.fixtures.ast.str.function_str"),
+        ("tests.appsec.iast.fixtures.ast.str.non_utf8_content"),  # EUC-JP file content
     ],
 )
 def test_astpatch_module_changed(module_name):
@@ -152,7 +153,6 @@ def test_astpatch_source_changed_with_future_imports(module_name):
         ("tests.appsec.iast.fixtures.ast.str.class_no_str"),
         ("tests.appsec.iast.fixtures.ast.str.function_no_str"),
         ("tests.appsec.iast.fixtures.ast.str.__init__"),  # Empty __init__.py
-        ("tests.appsec.iast.fixtures.ast.str.non_utf8_content"),  # EUC-JP file content
         ("tests.appsec.iast.fixtures.ast.str.empty_file"),
         ("tests.appsec.iast.fixtures.ast.subscript.store_context"),
     ],

tests/appsec/iast/test_loader.py

@@ -30,15 +30,12 @@ def test_patching_error():
 
         ddtrace.appsec._iast._loader.IS_IAST_ENABLED = True
 
-        with mock.patch("ddtrace.appsec._iast._loader.compile", side_effect=ValueError) as loader_compile, mock.patch(
-            "ddtrace.appsec._iast._loader.exec"
-        ) as loader_exec:
+        with mock.patch("ddtrace.appsec._iast._loader.compile", side_effect=ValueError) as loader_compile:
             importlib.reload(ddtrace.bootstrap.preload)
             imported_fixture_module = importlib.import_module(fixture_module)
 
             imported_fixture_module.add(2, 1)
             loader_compile.assert_called_once()
-            loader_exec.assert_not_called()
             assert ASPECTS_MODULE not in sys.modules
 
     finally: