From a265dd03dcf5e8112265e0e441ffc80348172f85 Mon Sep 17 00:00:00 2001 From: Christophe Papazian Date: Wed, 5 Feb 2025 15:22:32 +0100 Subject: [PATCH] complete decoupling --- ddtrace/appsec/__init__.py | 2 +- ddtrace/appsec/_constants.py | 2 + ddtrace/appsec/_utils.py | 4 +- ddtrace/contrib/internal/httplib/patch.py | 5 +- ddtrace/contrib/internal/mysql/patch.py | 5 +- ddtrace/contrib/internal/mysqldb/patch.py | 5 +- ddtrace/contrib/internal/requests/patch.py | 11 ++-- ddtrace/contrib/internal/sqlalchemy/patch.py | 5 +- ddtrace/contrib/internal/sqlite3/patch.py | 5 +- ddtrace/contrib/internal/urllib/patch.py | 12 ++-- ddtrace/contrib/internal/urllib3/patch.py | 19 +++--- ddtrace/contrib/internal/webbrowser/patch.py | 11 ++-- tests/internal/test_serverless.py | 63 ++++++++++++++------ 13 files changed, 98 insertions(+), 51 deletions(-) diff --git a/ddtrace/appsec/__init__.py b/ddtrace/appsec/__init__.py index a32442ccc0a..6b5758a95c2 100644 --- a/ddtrace/appsec/__init__.py +++ b/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 diff --git a/ddtrace/appsec/_constants.py b/ddtrace/appsec/_constants.py index 2172548205b..454483fcf17 100644 --- a/ddtrace/appsec/_constants.py +++ b/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 diff --git a/ddtrace/appsec/_utils.py b/ddtrace/appsec/_utils.py index 79f8f8b5311..e4dbae7a27f 100644 --- a/ddtrace/appsec/_utils.py +++ b/ddtrace/appsec/_utils.py @@ -1,3 +1,5 @@ +# this module must not load any other unsafe appsec module directly + import logging import sys from typing import Any @@ -5,6 +7,7 @@ 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: diff --git a/ddtrace/contrib/internal/httplib/patch.py b/ddtrace/contrib/internal/httplib/patch.py index 79a8ea2816f..7db5d59d31c 100644 --- a/ddtrace/contrib/internal/httplib/patch.py +++ b/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 diff --git a/ddtrace/contrib/internal/mysql/patch.py b/ddtrace/contrib/internal/mysql/patch.py index 6425bd33766..0aad999e546 100644 --- a/ddtrace/contrib/internal/mysql/patch.py +++ b/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 diff --git a/ddtrace/contrib/internal/mysqldb/patch.py b/ddtrace/contrib/internal/mysqldb/patch.py index cde0f58629f..fe9c62bbd5e 100644 --- a/ddtrace/contrib/internal/mysqldb/patch.py +++ b/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) diff --git a/ddtrace/contrib/internal/requests/patch.py b/ddtrace/contrib/internal/requests/patch.py index eab51c2c0a4..47518c6341f 100644 --- a/ddtrace/contrib/internal/requests/patch.py +++ b/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) diff --git a/ddtrace/contrib/internal/sqlalchemy/patch.py b/ddtrace/contrib/internal/sqlalchemy/patch.py index 916cc53daa4..c6d6df476c1 100644 --- a/ddtrace/contrib/internal/sqlalchemy/patch.py +++ b/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) diff --git a/ddtrace/contrib/internal/sqlite3/patch.py b/ddtrace/contrib/internal/sqlite3/patch.py index 03c79789661..68ee5779983 100644 --- a/ddtrace/contrib/internal/sqlite3/patch.py +++ b/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) diff --git a/ddtrace/contrib/internal/urllib/patch.py b/ddtrace/contrib/internal/urllib/patch.py index ed5e2891f06..1ba279fb20a 100644 --- a/ddtrace/contrib/internal/urllib/patch.py +++ b/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) diff --git a/ddtrace/contrib/internal/urllib3/patch.py b/ddtrace/contrib/internal/urllib3/patch.py index 6c10526c125..7c5d6adc28d 100644 --- a/ddtrace/contrib/internal/urllib3/patch.py +++ b/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) diff --git a/ddtrace/contrib/internal/webbrowser/patch.py b/ddtrace/contrib/internal/webbrowser/patch.py index 1387df37ac9..1f90e9cf9aa 100644 --- a/ddtrace/contrib/internal/webbrowser/patch.py +++ b/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) diff --git a/tests/internal/test_serverless.py b/tests/internal/test_serverless.py index 48d7586d25e..2f0dad15087 100644 --- a/tests/internal/test_serverless.py +++ b/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