Fix failing tests

Author: manuel-alvarez-alvarez

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java

@@ -204,18 +204,19 @@ private class AppSecConfigChangesListener implements ProductListener {
     @Override
     public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollingRateHinter)
         throws IOException {
-      final String key = configKey.toString();
       if (content == null) {
         remove(configKey, pollingRateHinter);
         return;
       }
+      final String key = configKey.toString();
       Map<String, Object> contentMap =
           ADAPTER.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
       if (contentMap == null || contentMap.isEmpty()) {
         ignoredConfigKeys.add(key);
       } else {
         ignoredConfigKeys.remove(key);
         try {
+          beforeApply(key, contentMap);
           maybeInitializeDefaultConfig();
           handleWafUpdateResultReport(key, contentMap);
         } catch (AppSecModule.AppSecModuleActivationException e) {
@@ -234,6 +235,7 @@ public void remove(ConfigKey configKey, PollingRateHinter pollingRateHinter)
       try {
         maybeInitializeDefaultConfig();
         wafBuilder.removeConfig(key);
+        afterRemove(key);
       } catch (UnclassifiedWafException e) {
         throw new RuntimeException(e);
       }
@@ -243,12 +245,15 @@ public void remove(ConfigKey configKey, PollingRateHinter pollingRateHinter)
     public void commit(PollingRateHinter pollingRateHinter) {
       // no action needed
     }
+
+    protected void beforeApply(final String key, final Map<String, Object> contentMap) {}
+
+    protected void afterRemove(final String key) {}
   }
 
   private class AppSecConfigChangesDDListener extends AppSecConfigChangesListener {
     @Override
-    public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollingRateHinter)
-        throws IOException {
+    protected void beforeApply(final String key, final Map<String, Object> config) {
       if (defaultConfigActivated) { // if we get any config, remove the default one
         log.debug("Removing default config");
         try {
@@ -258,15 +263,12 @@ public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollin
         }
         defaultConfigActivated = false;
       }
-      usedDDWafConfigKeys.add(configKey.toString());
-      super.accept(configKey, content, pollingRateHinter);
+      usedDDWafConfigKeys.add(key);
     }
 
     @Override
-    public void remove(ConfigKey configKey, PollingRateHinter pollingRateHinter)
-        throws IOException {
-      super.remove(configKey, pollingRateHinter);
-      usedDDWafConfigKeys.remove(configKey.toString());
+    protected void afterRemove(final String key) {
+      usedDDWafConfigKeys.remove(key);
     }
   }
 

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy

@@ -1,5 +1,6 @@
 package com.datadog.appsec.ddwaf
 
+import com.datadog.appsec.AppSecSystem
 import com.datadog.appsec.config.AppSecConfigService
 import com.datadog.appsec.config.AppSecConfigServiceImpl
 import com.datadog.appsec.config.AppSecModuleConfigurer
@@ -30,6 +31,8 @@ import com.squareup.moshi.Types
 import datadog.appsec.api.blocking.BlockingContentType
 import datadog.communication.monitor.Monitoring
 import datadog.remoteconfig.ConfigurationPoller
+import datadog.remoteconfig.PollerRequestFactory
+import datadog.remoteconfig.PollingRateHinter
 import datadog.remoteconfig.Product
 import datadog.remoteconfig.state.ConfigKey
 import datadog.remoteconfig.state.ParsedConfigKey
@@ -97,6 +100,7 @@ class WAFModuleSpecification extends DDSpecification {
   void setup() {
     WafMetricCollector.INSTANCE = wafMetricCollector
     AgentTracer.forceRegister(tracer)
+    AppSecSystem.active = true
 
     final configurationPoller = Stub(ConfigurationPoller) {
       addListener(Product.ASM_DD, _ as ProductListener) >> {
@@ -972,7 +976,7 @@ class WAFModuleSpecification extends DDSpecification {
 
   void 'configuration can be given later'() {
     when:
-    initialRuleAddWithMap([waf: null])
+    initialRuleAddWithMap([waf: new BadConfig()]) // empty configs are allowed now
 
     then:
     thrown RuntimeException
@@ -1695,4 +1699,14 @@ class WAFModuleSpecification extends DDSpecification {
       throw new IllegalStateException("Unhandled WafErrorCode: $code")
     }
   }
+
+  private static class BadConfig implements Map<String, Object> {
+    @Delegate
+    private Map<String, Object> delegate
+
+    @Override
+    Set entrySet() {
+      throw new IllegalStateException("You tried to iterate!")
+    }
+  }
 }