chore(debugging): in-product enablement

ddtrace/appsec/_ddwaf/ddwaf_types.py

@@ -12,6 +12,7 @@
 from ddtrace.appsec._ddwaf.waf_stubs import ddwaf_context_capsule
 from ddtrace.appsec._ddwaf.waf_stubs import ddwaf_handle_capsule
 from ddtrace.appsec._utils import _observator
+from ddtrace.internal._unpatched import unpatched_Popen
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
@@ -27,11 +28,16 @@
 if system() == "Linux":
     try:
         asm_config._bypass_instrumentation_for_waf = True
+        import subprocess
+
+        current_Popen = subprocess.Popen
+        subprocess.Popen = unpatched_Popen  # type: ignore[misc]
         ctypes.CDLL(ctypes.util.find_library("rt"), mode=ctypes.RTLD_GLOBAL)
     except Exception:  # nosec
         pass
     finally:
         asm_config._bypass_instrumentation_for_waf = False
+        subprocess.Popen = current_Popen  # type: ignore[misc]
 
 ARCHI = machine().lower()
 

ddtrace/appsec/_processor.py

@@ -128,7 +128,7 @@ def delayed_init(self) -> None:
                 self.metrics._set_waf_init_metric(self._ddwaf.info, self._ddwaf.initialized)
         except Exception:
             # Partial of DDAS-0005-00
-            log.warning("[DDAS-0005-00] WAF initialization failed")
+            log.warning("[DDAS-0005-00] WAF initialization failed", exc_info=True)
 
         self._update_required()
 

ddtrace/bootstrap/sitecustomize.py

@@ -86,6 +86,7 @@ def drop(module_name):
             "google",
             "google.protobuf",  # the upb backend in >= 4.21 does not like being unloaded
             "wrapt",
+            "bytecode",  # needed by before-fork hooks
         ]
     )
     for m in list(_ for _ in sys.modules if _ not in LOADED_MODULES):

ddtrace/debugging/_debugger.py

@@ -7,7 +7,6 @@
 import sys
 import threading
 import time
-from types import CodeType
 from types import FunctionType
 from types import ModuleType
 from types import TracebackType
@@ -27,6 +26,7 @@
 from ddtrace.debugging._function.discovery import FunctionDiscovery
 from ddtrace.debugging._function.store import FullyNamedContextWrappedFunction
 from ddtrace.debugging._function.store import FunctionStore
+from ddtrace.debugging._import import DebuggerModuleWatchdog
 from ddtrace.debugging._metrics import metrics
 from ddtrace.debugging._probe.model import FunctionLocationMixin
 from ddtrace.debugging._probe.model import FunctionProbe
@@ -45,8 +45,6 @@
 from ddtrace.debugging._uploader import UploaderProduct
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.metrics import Metrics
-from ddtrace.internal.module import ModuleHookType
-from ddtrace.internal.module import ModuleWatchdog
 from ddtrace.internal.module import origin
 from ddtrace.internal.module import register_post_run_module_hook
 from ddtrace.internal.module import unregister_post_run_module_hook
@@ -71,70 +69,6 @@ class DebuggerError(Exception):
     pass
 
 
-class DebuggerModuleWatchdog(ModuleWatchdog):
-    _locations: Set[str] = set()
-
-    def transform(self, code: CodeType, module: ModuleType) -> CodeType:
-        return FunctionDiscovery.transformer(code, module)
-
-    @classmethod
-    def register_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None:
-        if origin in cls._locations:
-            # We already have a hook for this origin, don't register a new one
-            # but invoke it directly instead, if the module was already loaded.
-            module = cls.get_by_origin(origin)
-            if module is not None:
-                hook(module)
-
-            return
-
-        cls._locations.add(str(origin))
-
-        super().register_origin_hook(origin, hook)
-
-    @classmethod
-    def unregister_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None:
-        try:
-            cls._locations.remove(str(origin))
-        except KeyError:
-            # Nothing to unregister.
-            return
-
-        return super().unregister_origin_hook(origin, hook)
-
-    @classmethod
-    def register_module_hook(cls, module: str, hook: ModuleHookType) -> None:
-        if module in cls._locations:
-            # We already have a hook for this origin, don't register a new one
-            # but invoke it directly instead, if the module was already loaded.
-            mod = sys.modules.get(module)
-            if mod is not None:
-                hook(mod)
-
-            return
-
-        cls._locations.add(module)
-
-        super().register_module_hook(module, hook)
-
-    @classmethod
-    def unregister_module_hook(cls, module: str, hook: ModuleHookType) -> None:
-        try:
-            cls._locations.remove(module)
-        except KeyError:
-            # Nothing to unregister.
-            return
-
-        return super().unregister_module_hook(module, hook)
-
-    @classmethod
-    def on_run_module(cls, module: ModuleType) -> None:
-        if cls._instance is not None:
-            # Treat run module as an import to trigger import hooks and register
-            # the module's origin.
-            cls._instance.after_import(module)
-
-
 class DebuggerWrappingContext(WrappingContext):
     __priority__ = 99  # Execute after all other contexts
 
@@ -275,8 +209,6 @@ def enable(cls) -> None:
 
         di_config.enabled = True
 
-        cls.__watchdog__.install()
-
         if di_config.metrics:
             metrics.enable()
 
@@ -308,7 +240,6 @@ def disable(cls, join: bool = True) -> None:
         cls._instance.stop(join=join)
         cls._instance = None
 
-        cls.__watchdog__.uninstall()
         if di_config.metrics:
             metrics.disable()
 

ddtrace/debugging/_import.py

@@ -0,0 +1,73 @@
+from pathlib import Path
+import sys
+from types import CodeType
+from types import ModuleType
+from typing import Set
+
+from ddtrace.debugging._function.discovery import FunctionDiscovery
+from ddtrace.internal.module import ModuleHookType
+from ddtrace.internal.module import ModuleWatchdog
+
+
+class DebuggerModuleWatchdog(ModuleWatchdog):
+    _locations: Set[str] = set()
+
+    def transform(self, code: CodeType, module: ModuleType) -> CodeType:
+        return FunctionDiscovery.transformer(code, module)
+
+    @classmethod
+    def register_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None:
+        if origin in cls._locations:
+            # We already have a hook for this origin, don't register a new one
+            # but invoke it directly instead, if the module was already loaded.
+            module = cls.get_by_origin(origin)
+            if module is not None:
+                hook(module)
+
+            return
+
+        cls._locations.add(str(origin))
+
+        super().register_origin_hook(origin, hook)
+
+    @classmethod
+    def unregister_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None:
+        try:
+            cls._locations.remove(str(origin))
+        except KeyError:
+            # Nothing to unregister.
+            return
+
+        return super().unregister_origin_hook(origin, hook)
+
+    @classmethod
+    def register_module_hook(cls, module: str, hook: ModuleHookType) -> None:
+        if module in cls._locations:
+            # We already have a hook for this origin, don't register a new one
+            # but invoke it directly instead, if the module was already loaded.
+            mod = sys.modules.get(module)
+            if mod is not None:
+                hook(mod)
+
+            return
+
+        cls._locations.add(module)
+
+        super().register_module_hook(module, hook)
+
+    @classmethod
+    def unregister_module_hook(cls, module: str, hook: ModuleHookType) -> None:
+        try:
+            cls._locations.remove(module)
+        except KeyError:
+            # Nothing to unregister.
+            return
+
+        return super().unregister_module_hook(module, hook)
+
+    @classmethod
+    def on_run_module(cls, module: ModuleType) -> None:
+        if cls._instance is not None:
+            # Treat run module as an import to trigger import hooks and register
+            # the module's origin.
+            cls._instance.after_import(module)

ddtrace/debugging/_products/dynamic_instrumentation.py

@@ -1,3 +1,5 @@
+import enum
+
 from ddtrace.settings.dynamic_instrumentation import config
 
 
@@ -8,11 +10,27 @@ def post_preload():
     pass
 
 
+def _start():
+    from ddtrace.debugging import DynamicInstrumentation
+
+    DynamicInstrumentation.enable()
+
+
 def start():
+    from ddtrace.debugging._import import DebuggerModuleWatchdog
+
+    # We need to install this on start-up because if DI gets enabled remotely
+    # we won't be able to capture many of the code objects from the modules
+    # that are already loaded.
+    DebuggerModuleWatchdog.install()
+
     if config.enabled:
-        from ddtrace.debugging import DynamicInstrumentation
+        _start()
+
 
-        DynamicInstrumentation.enable()
+def before_fork():
+    # We need to make sure that each process shares the same RC data connector
+    import ddtrace.debugging._probe.remoteconfig  # noqa
 
 
 def restart(join=False):
@@ -29,3 +47,13 @@ def stop(join=False):
 
 def at_exit(join=False):
     stop(join=join)
+
+
+class APMCapabilities(enum.IntFlag):
+    APM_TRACING_ENABLE_DYNAMIC_INSTRUMENTATION = 1 << 38
+
+
+def apm_tracing_rc(lib_config, _config):
+    if (enabled := lib_config.get("dynamic_instrumentation_enabled")) is not None:
+        should_start = (config.spec.enabled.full_name not in config.source or config.enabled) and enabled
+        _start() if should_start else stop()

ddtrace/internal/README.md

@@ -41,3 +41,4 @@ gets extended to add support for additional features.
 | `config: DDConfig` | A configuration object; when an instance of `DDConfig`, configuration telemetry is automatically reported |
 | `APMCapabilities: Type[enum.IntFlag]` | A set of capabilities that the product provides |
 | `apm_tracing_rc: (dict, ddtrace.settings._core.Config) -> None` | Product-specific remote configuration handler (e.g. remote enablement) |
+| `before_fork() -> None` | A function with the logic required to prepare the product for a fork |

ddtrace/internal/_unpatched.py

@@ -10,3 +10,8 @@
 # to get a reference to the right threading module.
 import threading as _threading  # noqa
 import gc as _gc  # noqa
+from subprocess import Popen as unpatched_Popen  # noqa
+
+import sys
+
+del sys.modules["subprocess"]

ddtrace/internal/products.py

@@ -136,6 +136,16 @@ def start_products(self) -> None:
                 log.exception("Failed to start product '%s'", name)
                 failed.add(name)
 
+    def before_fork(self) -> None:
+        for name, product in self.products:
+            try:
+                if (hook := getattr(product, "before_fork", None)) is None:
+                    continue
+                hook()
+                log.debug("Before-fork hook for product '%s' executed", name)
+            except Exception:
+                log.exception("Failed to execute before-fork hook for product '%s'", name)
+
     def restart_products(self, join: bool = False) -> None:
         failed: t.Set[str] = set()
 
@@ -185,6 +195,9 @@ def _do_products(self) -> None:
         # Start all products
         self.start_products()
 
+        # Execute before fork hooks
+        forksafe.register_before_fork(self.before_fork)
+
         # Restart products on fork
         forksafe.register(self.restart_products)
 

tests/debugging/exploration/debugger.py

@@ -13,9 +13,9 @@
 from ddtrace.debugging._config import di_config
 import ddtrace.debugging._debugger as _debugger
 from ddtrace.debugging._debugger import Debugger
-from ddtrace.debugging._debugger import DebuggerModuleWatchdog
 from ddtrace.debugging._encoding import LogSignalJsonEncoder
 from ddtrace.debugging._function.discovery import FunctionDiscovery
+from ddtrace.debugging._import import DebuggerModuleWatchdog
 from ddtrace.debugging._probe.model import Probe
 from ddtrace.debugging._probe.remoteconfig import ProbePollerEvent
 from ddtrace.debugging._signal.collector import SignalCollector

tests/debugging/mocking.py

@@ -202,7 +202,11 @@ def _debugger(config_to_override: DDConfig, config_overrides: Any) -> Generator[
 def debugger(**config_overrides: Any) -> Generator[TestDebugger, None, None]:
     """Test with the debugger enabled."""
     with _debugger(di_config, config_overrides) as debugger:
-        yield debugger
+        debugger.__watchdog__.install()
+        try:
+            yield debugger
+        finally:
+            debugger.__watchdog__.uninstall()
 
 
 class MockSpanExceptionHandler(SpanExceptionHandler):