complete decoupling

ddtrace/appsec/__init__.py

@@ -1,4 +1,4 @@
-# this module must not load any other appsec module
+# this module must not load any other unsafe appsec module directly
 
 from ddtrace.internal import core
 

ddtrace/appsec/_constants.py

@@ -1,3 +1,5 @@
+# this module must not load any other unsafe appsec module directly
+
 import os
 from re import Match
 import sys

ddtrace/appsec/_utils.py

@@ -1,10 +1,13 @@
+# this module must not load any other unsafe appsec module directly
+
 import logging
 import sys
 from typing import Any
 import uuid
 
 from ddtrace.appsec._constants import API_SECURITY
 from ddtrace.appsec._constants import APPSEC
+from ddtrace.appsec._constants import SPAN_DATA_NAMES
 from ddtrace.internal._unpatched import unpatched_json_loads
 from ddtrace.internal.compat import to_unicode
 from ddtrace.internal.logger import get_logger
@@ -21,7 +24,6 @@ def parse_response_body(raw_body):
     import xmltodict
 
     from ddtrace.appsec import _asm_request_context
-    from ddtrace.appsec._constants import SPAN_DATA_NAMES
     from ddtrace.contrib.internal.trace_utils import _get_header_value_case_insensitive
 
     if not raw_body:

ddtrace/contrib/internal/httplib/patch.py

@@ -5,7 +5,6 @@
 import wrapt
 
 from ddtrace import config
-from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_request_asm
 from ddtrace.constants import _ANALYTICS_SAMPLE_RATE_KEY
 from ddtrace.constants import SPAN_KIND
 from ddtrace.contrib import trace_utils
@@ -77,12 +76,14 @@ def _wrap_getresponse(func, instance, args, kwargs):
 
 
 def _call_asm_wrap(func, instance, *args, **kwargs):
+    from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_request_asm
+
     _wrap_request_asm(func, instance, args, kwargs)
 
 
 def _wrap_request(func, instance, args, kwargs):
     # Use any attached tracer if available, otherwise use the global tracer
-    if asm_config._iast_enabled or asm_config._asm_enabled:
+    if asm_config._iast_enabled or (asm_config._asm_enabled and asm_config._ep_enabled):
         func_to_call = functools.partial(_call_asm_wrap, func, instance)
     else:
         func_to_call = func

ddtrace/contrib/internal/mysql/patch.py

@@ -4,8 +4,6 @@
 import wrapt
 
 from ddtrace import config
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
 from ddtrace.contrib.dbapi import TracedConnection
 from ddtrace.contrib.internal.trace_utils import _convert_to_string
 from ddtrace.ext import db
@@ -51,6 +49,9 @@ def patch():
         mysql.connector.Connect = mysql.connector.connect
 
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
+
         _set_metric_iast_instrumented_sink(VULN_SQL_INJECTION)
     mysql.connector._datadog_patch = True
 

ddtrace/contrib/internal/mysqldb/patch.py

@@ -4,8 +4,6 @@
 from wrapt import wrap_function_wrapper as _w
 
 from ddtrace import config
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
 from ddtrace.constants import _SPAN_MEASURED_KEY
 from ddtrace.constants import SPAN_KIND
 from ddtrace.contrib.dbapi import TracedConnection
@@ -67,6 +65,9 @@ def patch():
         _w("MySQLdb", "connect", _connect)
 
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
+
         _set_metric_iast_instrumented_sink(VULN_SQL_INJECTION)
 
 

ddtrace/contrib/internal/requests/patch.py

@@ -4,9 +4,6 @@
 from wrapt import wrap_function_wrapper as _w
 
 from ddtrace import config
-from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_request
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SSRF
 from ddtrace.contrib.internal.trace_utils import unwrap as _u
 from ddtrace.internal.schema import schematize_service_name
 from ddtrace.internal.utils.formats import asbool
@@ -46,10 +43,16 @@ def patch():
 
     _w("requests", "Session.send", _wrap_send)
     # IAST needs to wrap this function because `Session.send` is too late
-    _w("requests", "Session.request", _wrap_request)
+    if asm_config._load_modules:
+        from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_request
+
+        _w("requests", "Session.request", _wrap_request)
     Pin(_config=config.requests).onto(requests.Session)
 
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SSRF
+
         _set_metric_iast_instrumented_sink(VULN_SSRF)
 
 

ddtrace/contrib/internal/sqlalchemy/patch.py

@@ -1,8 +1,6 @@
 import sqlalchemy
 from wrapt import wrap_function_wrapper as _w
 
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
 from ddtrace.contrib.internal.trace_utils import unwrap
 from ddtrace.settings.asm import config as asm_config
 
@@ -24,6 +22,9 @@ def patch():
     _w("sqlalchemy.engine", "create_engine", _wrap_create_engine)
 
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
+
         _set_metric_iast_instrumented_sink(VULN_SQL_INJECTION)
 
 

ddtrace/contrib/internal/sqlite3/patch.py

@@ -5,8 +5,6 @@
 import wrapt
 
 from ddtrace import config
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
 from ddtrace.contrib.dbapi import FetchTracedCursor
 from ddtrace.contrib.dbapi import TracedConnection
 from ddtrace.contrib.dbapi import TracedCursor
@@ -47,6 +45,9 @@ def patch():
     sqlite3.dbapi2.connect = wrapped
 
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
+
         _set_metric_iast_instrumented_sink(VULN_SQL_INJECTION)
 
 

ddtrace/contrib/internal/urllib/patch.py

@@ -2,9 +2,6 @@
 
 from wrapt import wrap_function_wrapper as _w
 
-from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_open
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SSRF
 from ddtrace.contrib.internal.trace_utils import unwrap as _u
 from ddtrace.settings.asm import config as asm_config
 
@@ -20,8 +17,15 @@ def patch():
         return
     urllib.request.__datadog_patch = True
 
-    _w("urllib.request", "urlopen", _wrap_open)
+    if asm_config._load_modules:
+        from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_open
+
+        _w("urllib.request", "urlopen", _wrap_open)
+
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SSRF
+
         _set_metric_iast_instrumented_sink(VULN_SSRF)
 
 

ddtrace/contrib/internal/urllib3/patch.py

@@ -4,9 +4,6 @@
 from wrapt import wrap_function_wrapper as _w
 
 from ddtrace import config
-from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_request
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SSRF
 from ddtrace.constants import _ANALYTICS_SAMPLE_RATE_KEY
 from ddtrace.constants import SPAN_KIND
 from ddtrace.contrib import trace_utils
@@ -54,14 +51,20 @@ def patch():
     urllib3.__datadog_patch = True
 
     _w("urllib3", "connectionpool.HTTPConnectionPool.urlopen", _wrap_urlopen)
-    if hasattr(urllib3, "_request_methods"):
-        _w("urllib3._request_methods", "RequestMethods.request", _wrap_request)
-    else:
-        # Old version before https://github.com/urllib3/urllib3/pull/2398
-        _w("urllib3.request", "RequestMethods.request", _wrap_request)
+    if asm_config._load_modules:
+        from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_request
+
+        if hasattr(urllib3, "_request_methods"):
+            _w("urllib3._request_methods", "RequestMethods.request", _wrap_request)
+        else:
+            # Old version before https://github.com/urllib3/urllib3/pull/2398
+            _w("urllib3.request", "RequestMethods.request", _wrap_request)
     Pin().onto(urllib3.connectionpool.HTTPConnectionPool)
 
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SSRF
+
         _set_metric_iast_instrumented_sink(VULN_SSRF)
 
 

ddtrace/contrib/internal/webbrowser/patch.py

@@ -2,9 +2,6 @@
 
 from wrapt import wrap_function_wrapper as _w
 
-from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_open
-from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-from ddtrace.appsec._iast.constants import VULN_SSRF
 from ddtrace.contrib.internal.trace_utils import unwrap as _u
 from ddtrace.settings.asm import config as asm_config
 
@@ -20,9 +17,15 @@ def patch():
         return
     webbrowser.__datadog_patch = True
 
-    _w("webbrowser", "open", _wrap_open)
+    if asm_config._load_modules:
+        from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_open
+
+        _w("webbrowser", "open", _wrap_open)
 
     if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+        from ddtrace.appsec._iast.constants import VULN_SSRF
+
         _set_metric_iast_instrumented_sink(VULN_SSRF)
 
 

tests/internal/test_serverless.py

@@ -1,4 +1,5 @@
 import mock
+import pytest
 
 from ddtrace.internal.serverless import in_azure_function
 from ddtrace.internal.serverless import in_gcp_function
@@ -78,7 +79,41 @@ def test_not_azure_function_consumption_plan():
     assert in_azure_function() is False
 
 
-def test_slow_imports():
+standard_blocklist = [
+        "ddtrace.appsec._api_security.api_manager",
+        "ddtrace.appsec._iast._ast.ast_patching",
+        "ddtrace.internal.telemetry.telemetry_writer",
+        "email.mime.application",
+        "email.mime.multipart",
+        "logging.handlers",
+        "multiprocessing",
+        "importlib_metadata",
+
+        # These modules must not be imported because their source files are
+        # specifically removed from the serverless python layer.
+        # See https://github.com/DataDog/datadog-lambda-python/blob/main/Dockerfile
+        "ddtrace.appsec._iast._taint_tracking._native",
+        "ddtrace.appsec._iast._stacktrace",
+        "ddtrace.internal.datadog.profiling.libdd_wrapper",
+        "ddtrace.internal.datadog.profiling.ddup._ddup",
+        "ddtrace.internal.datadog.profiling.stack_v2._stack_v2",
+]
+expanded_blocklist = standard_blocklist + [
+        "importlib.metadata",
+]
+
+@pytest.mark.parametrize('package,blocklist', [
+    ('ddtrace', expanded_blocklist),
+    ('ddtrace.contrib.internal.aws_lambda', expanded_blocklist),
+    ('ddtrace.contrib.internal.psycopg', expanded_blocklist),
+    # requests imports urlib3 which imports importlib.metadata
+    pytest.param('ddtrace.contrib.requests', standard_blocklist,
+            # Currently this package will import `ddtrace.appsec._iast._taint_tracking._native`
+            # and so is expected to fail for now. Once that is fixed and this test
+            # begins to XPASS, the xfail should be removed.
+            marks=pytest.mark.xfail(strict=True)),
+])
+def test_slow_imports(package, blocklist):
     # We should lazy load certain modules to avoid slowing down the startup
     # time when running in a serverless environment.  This test will fail if
     # any of those modules are imported during the import of ddtrace.
@@ -93,18 +128,6 @@ def test_slow_imports():
         }
     )
 
-    blocklist = [
-        "ddtrace.appsec._api_security.api_manager",
-        "ddtrace.appsec._iast._ast.ast_patching",
-        "ddtrace.internal.telemetry.telemetry_writer",
-        "email.mime.application",
-        "email.mime.multipart",
-        "logging.handlers",
-        "multiprocessing",
-        "importlib.metadata",
-        "importlib_metadata",
-    ]
-
     class BlockListFinder:
         def find_spec(self, fullname, *args):
             for lib in blocklist:
@@ -113,12 +136,14 @@ def find_spec(self, fullname, *args):
             return None
 
     try:
-        sys.meta_path.insert(0, BlockListFinder())
+        orig_meta_path = sys.meta_path
+        sys.meta_path = [BlockListFinder()] + orig_meta_path
+
+        for mod in sys.modules.copy():
+            if mod in blocklist or mod.startswith("ddtrace"):
+                del sys.modules[mod]
 
-        import ddtrace
-        import ddtrace.contrib.internal.aws_lambda  # noqa:F401
-        import ddtrace.contrib.internal.psycopg  # noqa:F401
+        __import__(package)
 
     finally:
-        if isinstance(sys.meta_path[0], BlockListFinder):
-            sys.meta_path = sys.meta_path[1:]
+        sys.meta_path = orig_meta_path