fix(rcm): store payloads target files content

ddtrace/appsec/_remoteconfiguration.py

@@ -20,6 +20,7 @@
 
 if TYPE_CHECKING:  # pragma: no cover
     from typing import Any
+    from typing import Dict
 
     try:
         from typing import Literal
@@ -31,7 +32,6 @@
     from typing import Union
 
     from ddtrace import Tracer
-    from ddtrace.internal.remoteconfig.client import ConfigMetadata
 
 log = get_logger(__name__)
 
@@ -45,28 +45,31 @@ def enable_appsec_rc(test_tracer=None):
     else:
         tracer = test_tracer
 
-    appsec_features_callback = RCAppSecFeaturesCallBack(tracer)
-    appsec_callback = RCAppSecCallBack(tracer)
+    asm_features_callback = RCAppSecFeaturesCallBack(tracer)
+    asm_dd_callback = RCASMDDCallBack(tracer)
+    asm_callback = RCAppSecCallBack(tracer)
 
     if _appsec_rc_features_is_enabled():
         from ddtrace.internal.remoteconfig import RemoteConfig
 
-        RemoteConfig.register(PRODUCTS.ASM_FEATURES, appsec_features_callback)
+        RemoteConfig.register(PRODUCTS.ASM_FEATURES, asm_features_callback)
 
     if tracer._appsec_enabled:
         from ddtrace.internal.remoteconfig import RemoteConfig
 
-        RemoteConfig.register(PRODUCTS.ASM_DATA, appsec_callback)  # IP Blocking
-        RemoteConfig.register(PRODUCTS.ASM, appsec_callback)  # Exclusion Filters & Custom Rules
-        RemoteConfig.register(PRODUCTS.ASM_DD, appsec_callback)  # DD Rules
+        RemoteConfig.register(PRODUCTS.ASM_DATA, asm_callback)  # IP Blocking
+        RemoteConfig.register(PRODUCTS.ASM, asm_callback)  # Exclusion Filters & Custom Rules
+        RemoteConfig.register(PRODUCTS.ASM_DD, asm_dd_callback)  # DD Rules
 
 
-def _add_rules_to_list(features, feature, message, rule_list):
-    # type: (Mapping[str, Any], str, str, list[Any]) -> None
-    rules = features.get(feature, [])
-    if rules:
+def _add_rules_to_list(features, feature, message, ruleset):
+    # type: (Mapping[str, Any], str, str, Dict[str, Any]) -> None
+    rules = features.get(feature, None)
+    if rules is not None:
         try:
-            rule_list += rules
+            if ruleset.get(feature) is None:
+                ruleset[feature] = []
+            ruleset[feature] += rules
             log.debug("Reloading Appsec %s: %s", message, rules)
         except JSONDecodeError:
             log.error("ERROR Appsec %s: invalid JSON content from remote configuration", message)
@@ -75,18 +78,32 @@ def _add_rules_to_list(features, feature, message, rule_list):
 def _appsec_rules_data(tracer, features):
     # type: (Tracer, Mapping[str, Any]) -> bool
     if features and tracer._appsec_processor:
-        ruleset = {"rules": [], "rules_data": [], "exclusions": [], "rules_override": []}  # type: dict[str, list[Any]]
-        _add_rules_to_list(features, "rules_data", "rules data", ruleset["rules_data"])
-        _add_rules_to_list(features, "custom_rules", "custom rules", ruleset["rules"])
-        _add_rules_to_list(features, "rules", "Datadog rules", ruleset["rules"])
-        _add_rules_to_list(features, "exclusions", "exclusion filters", ruleset["exclusions"])
-        _add_rules_to_list(features, "rules_override", "rules override", ruleset["rules_override"])
-        if any(ruleset.values()):
-            return tracer._appsec_processor._update_rules({k: v for k, v in ruleset.items() if v})
+        ruleset = {
+            "rules": None,
+            "rules_data": None,
+            "exclusions": None,
+            "rules_override": None,
+        }  # type: dict[str, Optional[list[Any]]]
+        _add_rules_to_list(features, "rules_data", "rules data", ruleset)
+        _add_rules_to_list(features, "custom_rules", "custom rules", ruleset)
+        _add_rules_to_list(features, "rules", "Datadog rules", ruleset)
+        _add_rules_to_list(features, "exclusions", "exclusion filters", ruleset)
+        _add_rules_to_list(features, "rules_override", "rules override", ruleset)
+        return tracer._appsec_processor._update_rules({k: v for k, v in ruleset.items() if v is not None})
 
     return False
 
 
+class RCASMDDCallBack(RemoteConfigCallBack):
+    def __init__(self, tracer):
+        # type: (Tracer) -> None
+        self.tracer = tracer
+
+    def __call__(self, metadata, features):
+        if features is not None:
+            _appsec_rules_data(self.tracer, features)
+
+
 class RCAppSecFeaturesCallBack(RemoteConfigCallBack):
     def __init__(self, tracer):
         # type: (Tracer) -> None
@@ -126,10 +143,11 @@ def _appsec_1click_activation(self, features):
             log.debug("Updating ASM Remote Configuration ASM_FEATURES: %s", rc_appsec_enabled)
 
             if rc_appsec_enabled:
-                appsec_callback = RCAppSecCallBack(self.tracer)
-                RemoteConfig.register(PRODUCTS.ASM_DATA, appsec_callback)  # IP Blocking
-                RemoteConfig.register(PRODUCTS.ASM, appsec_callback)  # Exclusion Filters & Custom Rules
-                RemoteConfig.register(PRODUCTS.ASM_DD, appsec_callback)  # DD Rules
+                asm_dd_callback = RCASMDDCallBack(self.tracer)
+                asm_callback = RCAppSecCallBack(self.tracer)
+                RemoteConfig.register(PRODUCTS.ASM_DATA, asm_callback)  # IP Blocking
+                RemoteConfig.register(PRODUCTS.ASM, asm_callback)  # Exclusion Filters & Custom Rules
+                RemoteConfig.register(PRODUCTS.ASM_DD, asm_dd_callback)  # DD Rules
                 if not self.tracer._appsec_enabled:
                     self.tracer.configure(appsec_enabled=True)
                 else:
@@ -146,13 +164,12 @@ def _appsec_1click_activation(self, features):
 
 
 class RCAppSecCallBack(RemoteConfigCallBackAfterMerge):
-    configs = {}
-
     def __init__(self, tracer):
         # type: (Tracer) -> None
+        super(RCAppSecCallBack, self).__init__()
         self.tracer = tracer
 
-    def __call__(self, metadata, features):
-        # type: (Optional[ConfigMetadata], Any) -> None
+    def __call__(self, target, features):
+        # type: (str, Any) -> None
         if features is not None:
             _appsec_rules_data(self.tracer, features)

ddtrace/internal/remoteconfig/client.py

@@ -154,17 +154,39 @@ def __call__(self, metadata, config):
         pass
 
 
-class RemoteConfigCallBackAfterMerge(RemoteConfigCallBack):
+class RemoteConfigCallBackAfterMerge(six.with_metaclass(abc.ABCMeta)):
     configs = {}  # type: Dict[str, Any]
 
-    def append(self, config):
-        self.configs.update(config)
+    @abc.abstractmethod
+    def __call__(self, target, config):
+        # type: (str, Any) -> None
+        pass
+
+    def append(self, target, config):
+        if not self.configs.get(target):
+            self.configs[target] = {}
+        if config is False:
+            # Remove old config from the configs dict. _remove_previously_applied_configurations function should
+            # call to this method
+            del self.configs[target]
+        elif config is not None:
+            # Append the new config to the configs dict. _load_new_configurations function should
+            # call to this method
+            if isinstance(config, dict):
+                self.configs[target].update(config)
+            else:
+                raise ValueError("target %s config %s has type of %s" % (target, config, type(config)))
 
     def dispatch(self):
-        try:
-            self.__call__(None, self.configs)
-        finally:
-            self.configs = {}
+        config_result = {}
+        for target, config in self.configs.items():
+            for key, value in config.items():
+                if isinstance(value, list):
+                    config_result[key] = config_result.get(key, []) + value
+                else:
+                    raise ValueError("target %s key %s has type of %s" % (target, key, type(value)))
+        if config_result:
+            self.__call__("", config_result)
 
 
 class RemoteConfigClient(object):
@@ -351,29 +373,47 @@ def _process_targets(self, payload):
         backend_state = signed.custom.get("opaque_backend_state")
         return signed.version, backend_state, targets
 
+    @staticmethod
+    def _apply_callback(list_callbacks, callback, config_content, target, config):
+        # type: (List[RemoteConfigCallBackAfterMerge], Any, Any, str, ConfigMetadata) -> None
+        if isinstance(callback, RemoteConfigCallBackAfterMerge):
+            callback.append(target, config_content)
+            if callback not in list_callbacks and not any(
+                filter(lambda x: isinstance(x, type(callback)), list_callbacks)
+            ):
+                list_callbacks.append(callback)
+        else:
+            callback(config, config_content)
+
     def _remove_previously_applied_configurations(self, applied_configs, client_configs, targets):
         # type: (AppliedConfigType, TargetsType, TargetsType) -> None
+        list_callbacks = []  # type: List[RemoteConfigCallBackAfterMerge]
         for target, config in self._applied_configs.items():
-            callback_action = None
             if target in client_configs and targets.get(target) == config:
                 # The configuration has not changed.
                 applied_configs[target] = config
                 continue
             elif target not in client_configs:
                 log.debug("Disable configuration: %s", target)
                 callback_action = False
+            else:
+                continue
 
             callback = self._products.get(config.product_name)
             if callback:
                 try:
-                    callback(config, callback_action)
+                    log.debug("Disable configuration 2: %s. ", target)
+                    self._apply_callback(list_callbacks, callback, callback_action, target, config)
                 except Exception:
                     log.debug("error while removing product %s config %r", config.product_name, config)
                     continue
 
+        for callback_to_dispach in list_callbacks:
+            callback_to_dispach.dispatch()
+
     def _load_new_configurations(self, applied_configs, client_configs, payload):
         # type: (AppliedConfigType, TargetsType, AgentPayload) -> None
-        list_callbacks = []
+        list_callbacks = []  # type: List[RemoteConfigCallBackAfterMerge]
         for target, config in client_configs.items():
             callback = self._products.get(config.product_name)
             if callback:
@@ -387,21 +427,17 @@ def _load_new_configurations(self, applied_configs, client_configs, payload):
 
                 try:
                     log.debug("Load new configuration: %s. content ", target)
-                    if isinstance(callback, RemoteConfigCallBackAfterMerge):
-                        callback.append(config_content)
-                        if callback not in list_callbacks:
-                            list_callbacks.append(callback)
-                    else:
-                        callback(config, config_content)
+                    self._apply_callback(list_callbacks, callback, config_content, target, config)
                 except Exception:
                     log.debug(
                         "Failed to load configuration %s for product %r", config, config.product_name, exc_info=True
                     )
                     continue
                 else:
                     applied_configs[target] = config
-        for callback in list_callbacks:
-            callback.dispatch()
+
+        for callback_to_dispach in list_callbacks:
+            callback_to_dispach.dispatch()
 
     def _add_apply_config_to_cache(self):
         if self._applied_configs: