chore(asm): more tags for appsec telemetry (#13005)

APPSEC-56592

- small type refactor, moving some classes from inside _ddwaf to
appsec._utils
- small telemetry refactor with proper definition of a class to
accumulate results instead of just using dictionaries.
- add 2 missing telemetry tags `waf_error` and `rate_limited` and
telemetry config_error metric counter.
- update some tests to handle the new tags and metrics

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: christophe-papazian

ddtrace/appsec/_asm_request_context.py

@@ -1,21 +1,21 @@
 import functools
 import json
 import re
-from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
-from typing import Literal  # noqa:F401
+from typing import Literal
 from typing import Optional
 from typing import Set
 from typing import Union
 from urllib import parse
 
 from ddtrace.appsec._constants import APPSEC
-from ddtrace.appsec._constants import EXPLOIT_PREVENTION
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
-from ddtrace.appsec._utils import _observator
+from ddtrace.appsec._utils import DDWaf_info
+from ddtrace.appsec._utils import DDWaf_result
+from ddtrace.appsec._utils import Telemetry_result
 from ddtrace.appsec._utils import add_context_log
 from ddtrace.appsec._utils import get_triggers
 from ddtrace.internal import core
@@ -26,10 +26,6 @@
 from ddtrace.trace import Span
 
 
-if TYPE_CHECKING:
-    from ddtrace.appsec._ddwaf import DDWaf_info
-    from ddtrace.appsec._ddwaf import DDWaf_result
-
 log = get_logger(__name__)
 
 # Stopgap module for providing ASM context for the blocking features wrapping some contextvars.
@@ -42,7 +38,6 @@
 _CONTEXT_CALL: Literal["context"] = "context"
 _WAF_CALL: Literal["waf_run"] = "waf_run"
 _BLOCK_CALL: Literal["block"] = "block"
-_TELEMETRY_WAF_RESULTS: Literal["t_waf_results"] = "t_waf_results"
 
 
 GLOBAL_CALLBACKS: Dict[str, List[Callable]] = {_CONTEXT_CALL: []}
@@ -73,28 +68,10 @@ def __init__(self, span: Optional[Span] = None):
         else:
             self.framework = self.span.name
         self.framework = self.framework.lower().replace(" ", "_")
-        self.waf_info: Optional[Callable[[], "DDWaf_info"]] = None
+        self.waf_info: Optional[Callable[[], DDWaf_info]] = None
         self.waf_addresses: Dict[str, Any] = {}
         self.callbacks: Dict[str, Any] = {_CONTEXT_CALL: []}
-        self.telemetry: Dict[str, Any] = {
-            _TELEMETRY_WAF_RESULTS: {
-                "blocked": False,
-                "triggered": False,
-                "timeout": 0,
-                "version": None,
-                "duration": 0.0,
-                "total_duration": 0.0,
-                "rasp": {
-                    "sum_eval": 0,
-                    "duration": 0.0,
-                    "total_duration": 0.0,
-                    "eval": {t: 0 for _, t in EXPLOIT_PREVENTION.TYPE},
-                    "match": {t: 0 for _, t in EXPLOIT_PREVENTION.TYPE},
-                    "timeout": {t: 0 for _, t in EXPLOIT_PREVENTION.TYPE},
-                },
-                "truncation": {"string_length": [], "container_size": [], "container_depth": []},
-            }
-        }
+        self.telemetry: Telemetry_result = Telemetry_result()
         self.addresses_sent: Set[str] = set()
         self.waf_triggers: List[Dict[str, Any]] = []
         self.blocked: Optional[Dict[str, Any]] = None
@@ -183,6 +160,8 @@ def update_span_metrics(span: Span, name: str, value: Union[float, int]) -> None
 
 
 def flush_waf_triggers(env: ASM_Environment) -> None:
+    from ddtrace.appsec._metrics import DDWAF_VERSION
+
     # Make sure we find a root span to attach the triggers to
     if env.span is None:
         from ddtrace.trace import tracer
@@ -204,35 +183,29 @@ def flush_waf_triggers(env: ASM_Environment) -> None:
         else:
             root_span.set_tag(APPSEC.JSON, json.dumps({"triggers": report_list}, separators=(",", ":")))
         env.waf_triggers = []
-    telemetry_results = get_value(_TELEMETRY, _TELEMETRY_WAF_RESULTS)
-    if telemetry_results:
-        from ddtrace.appsec._metrics import DDWAF_VERSION
-
-        root_span.set_tag_str(APPSEC.WAF_VERSION, DDWAF_VERSION)
-        if telemetry_results["total_duration"]:
-            update_span_metrics(root_span, APPSEC.WAF_DURATION, telemetry_results["duration"])
-            telemetry_results["duration"] = 0.0
-            update_span_metrics(root_span, APPSEC.WAF_DURATION_EXT, telemetry_results["total_duration"])
-            telemetry_results["total_duration"] = 0.0
-        if telemetry_results["timeout"]:
-            update_span_metrics(root_span, APPSEC.WAF_TIMEOUTS, telemetry_results["timeout"])
-        rasp_timeouts = sum(telemetry_results["rasp"]["timeout"].values())
-        if rasp_timeouts:
-            update_span_metrics(root_span, APPSEC.RASP_TIMEOUTS, rasp_timeouts)
-        if telemetry_results["rasp"]["sum_eval"]:
-            update_span_metrics(root_span, APPSEC.RASP_DURATION, telemetry_results["rasp"]["duration"])
-            update_span_metrics(root_span, APPSEC.RASP_DURATION_EXT, telemetry_results["rasp"]["total_duration"])
-            update_span_metrics(root_span, APPSEC.RASP_RULE_EVAL, telemetry_results["rasp"]["sum_eval"])
-        if telemetry_results["truncation"]["string_length"]:
-            root_span.set_metric(APPSEC.TRUNCATION_STRING_LENGTH, max(telemetry_results["truncation"]["string_length"]))
-        if telemetry_results["truncation"]["container_size"]:
-            root_span.set_metric(
-                APPSEC.TRUNCATION_CONTAINER_SIZE, max(telemetry_results["truncation"]["container_size"])
-            )
-        if telemetry_results["truncation"]["container_depth"]:
-            root_span.set_metric(
-                APPSEC.TRUNCATION_CONTAINER_DEPTH, max(telemetry_results["truncation"]["container_depth"])
-            )
+    telemetry_results: Telemetry_result = env.telemetry
+
+    root_span.set_tag_str(APPSEC.WAF_VERSION, DDWAF_VERSION)
+    if telemetry_results.total_duration:
+        update_span_metrics(root_span, APPSEC.WAF_DURATION, telemetry_results.duration)
+        telemetry_results.duration = 0.0
+        update_span_metrics(root_span, APPSEC.WAF_DURATION_EXT, telemetry_results.total_duration)
+        telemetry_results.total_duration = 0.0
+    if telemetry_results.timeout:
+        update_span_metrics(root_span, APPSEC.WAF_TIMEOUTS, telemetry_results.timeout)
+    rasp_timeouts = sum(telemetry_results.rasp.timeout.values())
+    if rasp_timeouts:
+        update_span_metrics(root_span, APPSEC.RASP_TIMEOUTS, rasp_timeouts)
+    if telemetry_results.rasp.sum_eval:
+        update_span_metrics(root_span, APPSEC.RASP_DURATION, telemetry_results.rasp.duration)
+        update_span_metrics(root_span, APPSEC.RASP_DURATION_EXT, telemetry_results.rasp.total_duration)
+        update_span_metrics(root_span, APPSEC.RASP_RULE_EVAL, telemetry_results.rasp.sum_eval)
+    if telemetry_results.truncation.string_length:
+        root_span.set_metric(APPSEC.TRUNCATION_STRING_LENGTH, max(telemetry_results.truncation.string_length))
+    if telemetry_results.truncation.container_size:
+        root_span.set_metric(APPSEC.TRUNCATION_CONTAINER_SIZE, max(telemetry_results.truncation.container_size))
+    if telemetry_results.truncation.container_depth:
+        root_span.set_metric(APPSEC.TRUNCATION_CONTAINER_DEPTH, max(telemetry_results.truncation.container_depth))
 
 
 def finalize_asm_env(env: ASM_Environment) -> None:
@@ -338,7 +311,7 @@ def set_waf_callback(value) -> None:
     set_value(_CALLBACKS, _WAF_CALL, value)
 
 
-def set_waf_info(info: Callable[[], "DDWaf_info"]) -> None:
+def set_waf_info(info: Callable[[], DDWaf_info]) -> None:
     env = _get_asm_context()
     if env is None:
         info_str = add_context_log(log, "appsec.asm_context.warning::set_waf_info::no_active_context")
@@ -347,7 +320,7 @@ def set_waf_info(info: Callable[[], "DDWaf_info"]) -> None:
     env.waf_info = info
 
 
-def call_waf_callback(custom_data: Optional[Dict[str, Any]] = None, **kwargs) -> Optional["DDWaf_result"]:
+def call_waf_callback(custom_data: Optional[Dict[str, Any]] = None, **kwargs) -> Optional[DDWaf_result]:
     if not asm_config._asm_enabled:
         return None
     callback = get_value(_CALLBACKS, _WAF_CALL)
@@ -435,44 +408,53 @@ def asm_request_context_set(
 
 def set_waf_telemetry_results(
     rules_version: Optional[str],
-    is_triggered: bool,
     is_blocked: bool,
-    is_timeout: bool,
+    waf_results: DDWaf_result,
     rule_type: Optional[str],
-    duration: float,
-    total_duration: float,
-    truncation: _observator,
+    is_sampled: bool,
 ) -> None:
-    result = get_value(_TELEMETRY, _TELEMETRY_WAF_RESULTS)
+    env = _get_asm_context()
+    if env is None:
+        return
+    result: Telemetry_result = env.telemetry
+    is_triggered = bool(waf_results.data)
     from ddtrace.appsec._metrics import _report_waf_truncations
 
-    _report_waf_truncations(truncation)
-    if result is not None:
-        for key in ["container_size", "container_depth", "string_length"]:
-            res = getattr(truncation, key)
-            if isinstance(res, int):
-                result["truncation"][key].append(res)
-        if rule_type is None:
-            # Request Blocking telemetry
-            result["triggered"] |= is_triggered
-            result["blocked"] |= is_blocked
-            result["timeout"] += is_timeout
-            if rules_version is not None:
-                result["version"] = rules_version
-            result["duration"] += duration
-            result["total_duration"] += total_duration
+    result.rate_limited |= is_sampled
+    if waf_results.return_code:
+        if result.error:
+            result.error = max(result.error, waf_results.return_code)
         else:
-            # Exploit Prevention telemetry
-            result["rasp"]["sum_eval"] += 1
-            result["rasp"]["eval"][rule_type] += 1
-            result["rasp"]["match"][rule_type] += int(is_triggered)
-            result["rasp"]["timeout"][rule_type] += int(is_timeout)
-            result["rasp"]["duration"] += duration
-            result["rasp"]["total_duration"] += total_duration
+            result.error = waf_results.return_code
+    _report_waf_truncations(waf_results.truncation)
+    for key in ["container_size", "container_depth", "string_length"]:
+        res = getattr(waf_results.truncation, key)
+        if isinstance(res, int):
+            getattr(result.truncation, key).append(res)
+    if rule_type is None:
+        # Request Blocking telemetry
+        result.triggered |= is_triggered
+        result.blocked |= is_blocked
+        result.timeout += waf_results.timeout
+        if rules_version is not None:
+            result.version = rules_version
+        result.duration += waf_results.runtime
+        result.total_duration += waf_results.total_runtime
+    else:
+        # Exploit Prevention telemetry
+        result.rasp.sum_eval += 1
+        result.rasp.eval[rule_type] += 1
+        result.rasp.match[rule_type] += int(is_triggered)
+        result.rasp.timeout[rule_type] += int(waf_results.timeout)
+        result.rasp.duration += waf_results.runtime
+        result.rasp.total_duration += waf_results.total_runtime
 
 
-def get_waf_telemetry_results() -> Optional[Dict[str, Any]]:
-    return get_value(_TELEMETRY, _TELEMETRY_WAF_RESULTS)
+def get_waf_telemetry_results() -> Optional[Telemetry_result]:
+    env = _get_asm_context()
+    if env:
+        return env.telemetry
+    return None
 
 
 def store_waf_results_data(data) -> None:

ddtrace/appsec/_ddwaf/__init__.py

@@ -1,9 +1,9 @@
 from typing import Type
 
 from ddtrace.appsec._ddwaf.waf_stubs import WAF
-from ddtrace.appsec._ddwaf.waf_stubs import DDWaf_info
-from ddtrace.appsec._ddwaf.waf_stubs import DDWaf_result
 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.settings.asm import config as asm_config
 

ddtrace/appsec/_ddwaf/waf_stubs.py

@@ -13,7 +13,8 @@
 from typing import Union
 
 from ddtrace.appsec._constants import DEFAULT
-from ddtrace.appsec._utils import _observator
+from ddtrace.appsec._utils import DDWaf_info
+from ddtrace.appsec._utils import DDWaf_result
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.remoteconfig import PayloadType
 
@@ -27,56 +28,6 @@
 DDWafRulesType = Union[None, int, str, List[Any], Dict[str, Any]]
 
 
-class DDWaf_result:
-    __slots__ = ["return_code", "data", "actions", "runtime", "total_runtime", "timeout", "truncation", "derivatives"]
-
-    def __init__(
-        self,
-        return_code: int,
-        data: List[Dict[str, Any]],
-        actions: Dict[str, Any],
-        runtime: float,
-        total_runtime: float,
-        timeout: bool,
-        truncation: _observator,
-        derivatives: Dict[str, Any],
-    ):
-        self.return_code = return_code
-        self.data = data
-        self.actions = actions
-        self.runtime = runtime
-        self.total_runtime = total_runtime
-        self.timeout = timeout
-        self.truncation = truncation
-        self.derivatives = derivatives
-
-    def __repr__(self):
-        return (
-            f"DDWaf_result(return_code: {self.return_code} data: {self.data},"
-            f" actions: {self.actions}, runtime: {self.runtime},"
-            f" total_runtime: {self.total_runtime}, timeout: {self.timeout},"
-            f" truncation: {self.truncation}, derivatives: {self.derivatives})"
-        )
-
-
-class DDWaf_info:
-    __slots__ = ["loaded", "failed", "errors", "version"]
-
-    def __init__(self, loaded: int, failed: int, errors: str, version: str):
-        self.loaded = loaded
-        self.failed = failed
-        self.errors = errors
-        self.version = version
-
-    def __repr__(self):
-        return "{loaded: %d, failed: %d, errors: %s, version: %s}" % (
-            self.loaded,
-            self.failed,
-            self.errors,
-            self.version,
-        )
-
-
 class ddwaf_handle_capsule(Generic[T]):
     def __init__(self, handle: Type[T], free_fn: Callable[[Type[T]], None]) -> None:
         self.handle = handle

ddtrace/appsec/_metrics.py

@@ -20,13 +20,18 @@
 
 @deduplication
 def _set_waf_error_log(msg: str, version: str, error_level: bool = True) -> None:
-    tags = {
+    log_tags = {
         "waf_version": DDWAF_VERSION,
         "event_rules_version": version or UNKNOWN_VERSION,
         "lib_language": "python",
     }
     level = TELEMETRY_LOG_LEVEL.ERROR if error_level else TELEMETRY_LOG_LEVEL.WARNING
-    telemetry.telemetry_writer.add_log(level, msg, tags=tags)
+    telemetry.telemetry_writer.add_log(level, msg, tags=log_tags)
+    tags = (
+        ("waf_version", DDWAF_VERSION),
+        ("event_rules_version", version or UNKNOWN_VERSION),
+    )
+    telemetry.telemetry_writer.add_count_metric(TELEMETRY_NAMESPACE.APPSEC, "waf.config_errors", 1, tags=tags)
 
 
 def _set_waf_updates_metric(info, success: bool):
@@ -110,26 +115,26 @@ def _set_waf_request_metrics(*_args):
             # TODO: enable it when Telemetry intake accepts this tag
             # is_truncation = any((result.truncation for result in list_results))
 
-            truncation = result["truncation"]
-            input_truncated = bool(
-                truncation["string_length"] or truncation["container_size"] or truncation["container_depth"]
-            )
+            truncation = result.truncation
+            input_truncated = bool(truncation.string_length or truncation.container_size or truncation.container_depth)
             tags_request = (
-                ("event_rules_version", result["version"] or UNKNOWN_VERSION),
+                ("event_rules_version", result.version or UNKNOWN_VERSION),
                 ("waf_version", DDWAF_VERSION),
-                ("rule_triggered", bool_str[result["triggered"]]),
-                ("request_blocked", bool_str[result["blocked"]]),
-                ("waf_timeout", bool_str[bool(result["timeout"])]),
+                ("rule_triggered", bool_str[result.triggered]),
+                ("request_blocked", bool_str[result.blocked]),
+                ("waf_timeout", bool_str[bool(result.timeout)]),
                 ("input_truncated", bool_str[input_truncated]),
+                ("waf_error", str(result.error)),
+                ("rate_limited", bool_str[result.rate_limited]),
             )
 
             telemetry.telemetry_writer.add_count_metric(
                 TELEMETRY_NAMESPACE.APPSEC, "waf.requests", 1, tags=tags_request
             )
-            rasp = result["rasp"]
-            if rasp["sum_eval"]:
+            rasp = result.rasp
+            if rasp.sum_eval:
                 for t, n in [("eval", "rasp.rule.eval"), ("match", "rasp.rule.match"), ("timeout", "rasp.timeout")]:
-                    for rule_type, value in rasp[t].items():
+                    for rule_type, value in getattr(rasp, t).items():
                         if value:
                             telemetry.telemetry_writer.add_count_metric(
                                 TELEMETRY_NAMESPACE.APPSEC,
@@ -138,7 +143,7 @@ def _set_waf_request_metrics(*_args):
                                 tags=_TYPES_AND_TAGS.get(rule_type, ())
                                 + (
                                     ("waf_version", DDWAF_VERSION),
-                                    ("event_rules_version", result["version"] or UNKNOWN_VERSION),
+                                    ("event_rules_version", result.version or UNKNOWN_VERSION),
                                 ),
                             )