chore(aap): remove waf_mockup + prevent concurrent access to ASM_Environment after finalization

ddtrace/appsec/_api_security/api_manager.py

@@ -8,6 +8,7 @@
 
 from ddtrace._trace._limits import MAX_SPAN_META_VALUE_LEN
 from ddtrace._trace.processor.resource_renaming import SimplifiedEndpointComputer
+from ddtrace.appsec._asm_request_context import _WAF_CALL
 from ddtrace.appsec._asm_request_context import ASM_Environment
 from ddtrace.appsec._constants import API_SECURITY
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
@@ -196,7 +197,8 @@ def _schema_callback(self, env):
                 value = transform(value)
             waf_payload[address] = value
 
-        result = self._asm_context.call_waf_callback(waf_payload)
+        callback = env.callbacks[_WAF_CALL]
+        result = callback(waf_payload)
         if result is None:
             return
         nb_schemas = 0

ddtrace/appsec/_asm_request_context.py

@@ -283,6 +283,7 @@ def finalize_asm_env(env: ASM_Environment) -> None:
     flush_waf_triggers(env)
     for function in env.callbacks[_CONTEXT_CALL]:
         function(env)
+    env.callbacks.clear()
     entry_span = env.entry_span
     if entry_span:
         if env.waf_info:
@@ -332,11 +333,17 @@ def set_body_response(body_response):
     # local import to avoid circular import
     from ddtrace.appsec._utils import parse_response_body
 
+    env = _get_asm_context()
+    if env is None:
+        extra = {"product": "appsec", "more_info": "::set_body_response", "stack_limit": 4}
+        logger.debug("asm_context::set_body_response::no_active_context", extra=extra, stack_info=True)
+        return
+
     set_waf_address(
         SPAN_DATA_NAMES.RESPONSE_BODY,
         lambda: parse_response_body(
             body_response,
-            get_waf_address(SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES),
+            env.waf_addresses.get(SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES, None),
         ),
     )
 
@@ -355,7 +362,7 @@ def set_waf_address(address: str, value: Any) -> None:
 
 def get_value(category: str, address: str, default: Any = None) -> Any:
     env = _get_asm_context()
-    if env is None:
+    if env is None or env.finalized:
         extra = {"product": "appsec", "more_info": f"::{category}::{address}", "stack_limit": 4}
         logger.debug("asm_context::get_value::no_active_context", extra=extra, stack_info=True)
         return default

ddtrace/appsec/_ddwaf/__init__.py

@@ -1,28 +1,28 @@
+from typing import Optional
+
 from ddtrace.appsec._ddwaf.waf_stubs import WAF
 from ddtrace.appsec._ddwaf.waf_stubs import DDWafRulesType
 from ddtrace.appsec._utils import DDWaf_info
 from ddtrace.appsec._utils import DDWaf_result
 from ddtrace.internal.logger import get_logger
-from ddtrace.internal.settings.asm import config as asm_config
 
 
-__all__ = ["DDWaf", "DDWaf_info", "DDWaf_result", "version", "DDWafRulesType"]
+__all__ = ["WAF", "DDWaf_info", "DDWaf_result", "version", "DDWafRulesType"]
 
 LOGGER = get_logger(__name__)
 
 _DDWAF_LOADED: bool = False
+version: str = "unloaded"
+
 
-if asm_config._asm_libddwaf_available:
+def waf_module() -> Optional[type[WAF]]:
     try:
         import ddtrace.appsec._ddwaf.waf as waf_module
 
+        global _DDWAF_LOADED, version
         _DDWAF_LOADED = True
+        version = waf_module.version()
+        return waf_module.DDWaf
     except Exception:
-        import ddtrace.appsec._ddwaf.waf_mock as waf_module  # type: ignore[no-redef]
-
         LOGGER.warning("DDWaf features disabled. WARNING: Dynamic Library not loaded", exc_info=True)
-else:
-    import ddtrace.appsec._ddwaf.waf_mock as waf_module  # type: ignore[no-redef]
-
-DDWaf: type[WAF] = waf_module.DDWaf
-version = waf_module.version
+        return None

ddtrace/appsec/_processor.py

@@ -2,21 +2,12 @@
 import errno
 from json.decoder import JSONDecodeError
 import os
-from typing import TYPE_CHECKING
 from typing import Any
 from typing import ClassVar
 from typing import Optional
 from typing import Sequence
 from typing import Union
 
-from ddtrace.ext import SpanTypes
-from ddtrace.internal import core
-
-
-if TYPE_CHECKING:
-    import ddtrace.appsec._ddwaf as ddwaf
-
-
 from ddtrace._trace.processor import SpanProcessor
 from ddtrace._trace.span import Span
 from ddtrace.appsec import _asm_request_context
@@ -27,6 +18,8 @@
 from ddtrace.appsec._constants import STACK_TRACE
 from ddtrace.appsec._constants import WAF_ACTIONS
 from ddtrace.appsec._constants import WAF_DATA_NAMES
+import ddtrace.appsec._ddwaf.ddwaf_types as ddwaf_types
+from ddtrace.appsec._ddwaf.waf_stubs import WAF
 from ddtrace.appsec._exploit_prevention.stack_traces import report_stack
 from ddtrace.appsec._trace_utils import _asm_manual_keep
 from ddtrace.appsec._utils import Binding_error
@@ -35,6 +28,8 @@
 from ddtrace.appsec._utils import is_inferred_span
 from ddtrace.constants import _ORIGIN_KEY
 from ddtrace.constants import _RUNTIME_FAMILY
+from ddtrace.ext import SpanTypes
+from ddtrace.internal import core
 from ddtrace.internal._unpatched import unpatched_open as open  # noqa: A004
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.rate_limiter import RateLimiter
@@ -134,21 +129,30 @@ def __post_init__(self) -> None:
     def delayed_init(self) -> None:
         try:
             if self._rules is not None and not hasattr(self, "_ddwaf"):
-                from ddtrace.appsec._ddwaf import DDWaf  # noqa: E402
+                from ddtrace.appsec._ddwaf import waf_module  # noqa: E402
                 import ddtrace.appsec._metrics as metrics  # noqa: E402
 
+                DDWaf = waf_module()
+                if DDWaf is None:
+                    log.warning("DDWaf features disabled. WARNING: Dynamic Library not loaded")
+                    self._ddwaf = None
+                    return
                 self.metrics = metrics
-                self._ddwaf = DDWaf(
+                self._ddwaf: Optional[WAF] = DDWaf(
                     self._rules, self.obfuscation_parameter_key_regexp, self.obfuscation_parameter_value_regexp, metrics
                 )
-                self.metrics._set_waf_init_metric(self._ddwaf.info, self._ddwaf.initialized)
+                if self._ddwaf:
+                    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", exc_info=True)
+            self._ddwaf = None
 
         self._update_required()
 
     def _update_required(self):
+        if self._ddwaf is None:
+            return
         self._addresses_to_keep.clear()
         for address in self._ddwaf.required_data:
             self._addresses_to_keep.add(address)
@@ -162,6 +166,8 @@ def _update_rules(
     ) -> bool:
         if not hasattr(self, "_ddwaf"):
             self.delayed_init()
+        if self._ddwaf is None:
+            return False
         result = False
         if asm_config._asm_static_rule_file is not None:
             return result
@@ -195,6 +201,8 @@ def on_span_start(self, span: Span) -> None:
 
         if not hasattr(self, "_ddwaf"):
             self.delayed_init()
+        if self._ddwaf is None:
+            return
 
         if span.span_type not in asm_config._asm_processed_span_types:
             return
@@ -246,7 +254,7 @@ def waf_callable(custom_data=None, **kwargs):
     def _waf_action(
         self,
         entry_span: Span,
-        ctx: "ddwaf.ddwaf_types.ddwaf_context_capsule",
+        ctx: ddwaf_types.ddwaf_context_capsule,
         custom_data: Optional[dict[str, Any]] = None,
         crop_trace: Optional[str] = None,
         rule_type: Optional[str] = None,
@@ -263,6 +271,9 @@ def _waf_action(
         be retrieved from the `core`. This can be used when you don't want to store
         the value in the `core` before checking the `WAF`.
         """
+        if not hasattr(self, "_ddwaf") or self._ddwaf is None:
+            return None
+
         if _asm_request_context.get_blocked():
             # We still must run the waf if we need to extract schemas for API SECURITY
             if not custom_data or not custom_data.get("PROCESSOR_SETTINGS", {}).get("extract-schema", False):
@@ -310,7 +321,7 @@ def _waf_action(
         except Exception:
             log.debug("appsec::processor::waf::run", exc_info=True)
             waf_results = Binding_error
-        _asm_request_context.set_waf_info(lambda: self._ddwaf.info)
+        _asm_request_context.set_waf_info(lambda: self._ddwaf.info)  # type: ignore
         if waf_results.return_code < 0:
             error_tag = APPSEC.RASP_ERROR if rule_type else APPSEC.WAF_ERROR
             previous = entry_span.get_tag(error_tag)
@@ -390,9 +401,12 @@ def _is_needed(self, address: str) -> bool:
         return address in self._addresses_to_keep
 
     def on_span_finish(self, span: Span) -> None:
+        ddwaf = getattr(self, "_ddwaf", None)
+        if ddwaf is None:
+            return
         if span.span_type in asm_config._asm_processed_span_types:
             _asm_request_context.call_waf_callback_no_instrumentation()
-            self._ddwaf._at_request_end()
+            ddwaf._at_request_end()
             _asm_request_context.end_context(span)
 
     @classmethod