fix(asm): add global states to ensure patching once [backport 2.15] (#…

…11534)

ddtrace/_monkey.py

@@ -5,11 +5,13 @@
 
 from wrapt.importer import when_imported
 
+from ddtrace.appsec import load_common_appsec_modules
+
+from .appsec._iast._utils import _is_iast_enabled
 from .internal import telemetry
 from .internal.logger import get_logger
 from .internal.utils import formats
 from .settings import _config as config
-from .settings.asm import config as asm_config
 
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -225,17 +227,14 @@ def patch_all(**patch_modules):
     modules.update(patch_modules)
 
     patch(raise_errors=False, **modules)
-    if asm_config._iast_enabled:
+    if _is_iast_enabled():
         from ddtrace.appsec._iast._patch_modules import patch_iast
         from ddtrace.appsec.iast import enable_iast_propagation
 
         patch_iast()
         enable_iast_propagation()
 
-    if asm_config._ep_enabled or asm_config._iast_enabled:
-        from ddtrace.appsec._common_module_patches import patch_common_modules
-
-        patch_common_modules()
+    load_common_appsec_modules()
 
 
 def patch(raise_errors=True, patch_modules_prefix=DEFAULT_MODULES_PREFIX, **patch_modules):

ddtrace/appsec/__init__.py

@@ -1,3 +1,7 @@
+from ddtrace.internal import core
+from ddtrace.settings.asm import config as asm_config
+
+
 _APPSEC_TO_BE_LOADED = True
 
 
@@ -9,3 +13,25 @@ def load_appsec():
     if _APPSEC_TO_BE_LOADED:
         listen()
         _APPSEC_TO_BE_LOADED = False
+
+
+def load_iast():
+    """Lazily load the iast module listeners."""
+    from ddtrace.appsec._iast._iast_request_context import iast_listen
+
+    global _IAST_TO_BE_LOADED
+    if _IAST_TO_BE_LOADED:
+        iast_listen()
+        _IAST_TO_BE_LOADED = False
+
+
+def load_common_appsec_modules():
+    """Lazily load the common module patches."""
+    if (asm_config._ep_enabled and asm_config._asm_enabled) or asm_config._iast_enabled:
+        from ddtrace.appsec._common_module_patches import patch_common_modules
+
+        patch_common_modules()
+
+
+# for tests that needs to load the appsec module later
+core.on("test.config.override", load_common_appsec_modules)

ddtrace/appsec/_common_module_patches.py

@@ -27,8 +27,13 @@
 log = get_logger(__name__)
 _DD_ORIGINAL_ATTRIBUTES: Dict[Any, Any] = {}
 
+_is_patched = False
+
 
 def patch_common_modules():
+    global _is_patched
+    if _is_patched:
+        return
     try_wrap_function_wrapper("builtins", "open", wrapped_open_CFDDB7ABBA9081B6)
     try_wrap_function_wrapper("urllib.request", "OpenerDirector.open", wrapped_open_ED4CF71136E15EBF)
     try_wrap_function_wrapper("_io", "BytesIO.read", wrapped_read_F3E51D71B4EC16EF)
@@ -37,13 +42,18 @@ def patch_common_modules():
     core.on("asm.block.dbapi.execute", execute_4C9BAC8E228EB347)
     if asm_config._iast_enabled:
         _set_metric_iast_instrumented_sink(VULN_PATH_TRAVERSAL)
+    _is_patched = True
 
 
 def unpatch_common_modules():
+    global _is_patched
+    if not _is_patched:
+        return
     try_unwrap("builtins", "open")
     try_unwrap("urllib.request", "OpenerDirector.open")
     try_unwrap("_io", "BytesIO.read")
     try_unwrap("_io", "StringIO.read")
+    _is_patched = False
 
 
 def wrapped_read_F3E51D71B4EC16EF(original_read_callable, instance, args, kwargs):
@@ -69,10 +79,14 @@ def wrapped_open_CFDDB7ABBA9081B6(original_open_callable, instance, args, kwargs
     wrapper for open file function
     """
     if asm_config._iast_enabled:
-        from ddtrace.appsec._iast.taint_sinks.path_traversal import check_and_report_path_traversal
-
-        check_and_report_path_traversal(*args, **kwargs)
+        try:
+            from ddtrace.appsec._iast.taint_sinks.path_traversal import check_and_report_path_traversal
 
+            check_and_report_path_traversal(*args, **kwargs)
+        except ImportError:
+            # open is used during module initialization
+            # and shouldn't be changed at that time
+            return original_open_callable(*args, **kwargs)
     if (
         asm_config._asm_enabled
         and asm_config._ep_enabled

ddtrace/appsec/_iast/__init__.py

@@ -27,6 +27,7 @@ def wrapped_function(wrapped, instance, args, kwargs):
     )
     return wrapped(*args, **kwargs)
 """  # noqa: RST201, RST213, RST210
+
 import inspect
 import os
 import sys
@@ -75,8 +76,14 @@ def ddtrace_iast_flask_patch():
     sys.modules[module_name] = compiled_code
 
 
+_iast_propagation_enabled = False
+
+
 def enable_iast_propagation():
     """Add IAST AST patching in the ModuleWatchdog"""
+    global _iast_propagation_enabled
+    if _iast_propagation_enabled:
+        return
     if asbool(os.getenv(IAST.ENV, "false")):
         from ddtrace.appsec._iast._utils import _is_python_version_supported
 
@@ -86,17 +93,22 @@ def enable_iast_propagation():
 
             log.debug("IAST enabled")
             ModuleWatchdog.register_pre_exec_module_hook(_should_iast_patch, _exec_iast_patched_module)
+            _iast_propagation_enabled = True
 
 
 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
 
+    global _iast_propagation_enabled
+    if not _iast_propagation_enabled:
+        return
     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")
+    _iast_propagation_enabled = False
 
 
 __all__ = [

ddtrace/appsec/_remoteconfiguration.py

@@ -73,6 +73,11 @@ def enable_appsec_rc(test_tracer: Optional[Tracer] = None) -> None:
         remoteconfig_poller.register(PRODUCTS.ASM_DATA, asm_callback)  # IP Blocking
         remoteconfig_poller.register(PRODUCTS.ASM, asm_callback)  # Exclusion Filters & Custom Rules
         remoteconfig_poller.register(PRODUCTS.ASM_DD, asm_callback)  # DD Rules
+    # ensure exploit prevention patches are loaded by one-click activation
+    if asm_config._asm_enabled:
+        from ddtrace.appsec import load_common_appsec_modules
+
+        load_common_appsec_modules()
 
     forksafe.register(_forksafe_appsec_rc)
     telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.APPSEC, True)

releasenotes/notes/exploit_prevention_patch_fix-1bdd7540e1d085d8.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    ASM: This fix ensures that common patches for exploit prevention and sca are only loaded if required, and only loaded once.

tests/appsec/integrations/test_flask_telemetry.py

@@ -1,4 +1,5 @@
 from ddtrace.appsec._handlers import _on_flask_patch
+from ddtrace.appsec._iast.constants import VULN_PATH_TRAVERSAL
 from tests.appsec.appsec_utils import flask_server
 from tests.utils import override_global_config
 
@@ -24,7 +25,8 @@ def test_flask_instrumented_metrics(telemetry_writer):
     metrics_result = telemetry_writer._namespace._metrics_data
     metrics_source_tags_result = [metric._tags[0][1] for metric in metrics_result["generate-metrics"]["iast"].values()]
 
-    assert len(metrics_source_tags_result) == 7
+    assert len(metrics_source_tags_result) == 8
+    assert VULN_PATH_TRAVERSAL in metrics_source_tags_result
     assert origin_to_str(OriginType.HEADER_NAME) in metrics_source_tags_result
     assert origin_to_str(OriginType.HEADER) in metrics_source_tags_result
     assert origin_to_str(OriginType.PARAMETER) in metrics_source_tags_result

tests/utils.py

@@ -22,6 +22,7 @@
 from ddtrace.constants import SPAN_MEASURED_KEY
 from ddtrace.ext import http
 from ddtrace.internal import agent
+from ddtrace.internal import core
 from ddtrace.internal.ci_visibility.writer import CIVisibilityWriter
 from ddtrace.internal.compat import httplib
 from ddtrace.internal.compat import parse
@@ -177,6 +178,7 @@ def override_global_config(values):
         if key in asm_config_keys:
             setattr(ddtrace.settings.asm.config, key, value)
     try:
+        core.dispatch("test.config.override")
         yield
     finally:
         # Reset all to their original values