Fix NPE in AppSecConfigServiceImpl (#9165)

Author: manuel-alvarez-alvarez

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

@@ -75,8 +75,7 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
   private final ConfigurationPoller configurationPoller;
   private WafBuilder wafBuilder;
 
-  private MergedAsmFeatures mergedAsmFeatures;
-  private volatile boolean initialized;
+  private final MergedAsmFeatures mergedAsmFeatures = new MergedAsmFeatures();
 
   private final ConcurrentHashMap<String, SubconfigListener> subconfigListeners =
       new ConcurrentHashMap<>();
@@ -173,9 +172,7 @@ private class AppSecConfigChangesListener implements ProductListener {
     @Override
     public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollingRateHinter)
         throws IOException {
-      if (!initialized) {
-        throw new IllegalStateException();
-      }
+      maybeInitializeDefaultConfig();
 
       if (content == null) {
         try {
@@ -219,8 +216,8 @@ public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollin
         }
         defaultConfigActivated = false;
       }
-      super.accept(configKey, content, pollingRateHinter);
       usedDDWafConfigKeys.add(configKey.toString());
+      super.accept(configKey, content, pollingRateHinter);
     }
 
     @Override
@@ -282,13 +279,7 @@ private void subscribeAsmFeatures() {
         Product.ASM_FEATURES,
         AppSecFeaturesDeserializer.INSTANCE,
         (configKey, newConfig, hinter) -> {
-          if (!hasUserWafConfig && !defaultConfigActivated) {
-            // features activated in runtime
-            init();
-          }
-          if (!initialized) {
-            throw new IllegalStateException();
-          }
+          maybeInitializeDefaultConfig();
           if (newConfig == null) {
             mergedAsmFeatures.removeConfig(configKey);
           } else {
@@ -305,10 +296,7 @@ private void subscribeAsmFeatures() {
 
   private void distributeSubConfigurations(
       String key, AppSecModuleConfigurer.Reconfiguration reconfiguration) {
-    if (usedDDWafConfigKeys.isEmpty() && !defaultConfigActivated && !hasUserWafConfig) {
-      // no config left in the WAF builder, add the default config
-      init();
-    }
+    maybeInitializeDefaultConfig();
     for (Map.Entry<String, SubconfigListener> entry : subconfigListeners.entrySet()) {
       SubconfigListener listener = entry.getValue();
       try {
@@ -320,6 +308,13 @@ private void distributeSubConfigurations(
     }
   }
 
+  private void maybeInitializeDefaultConfig() {
+    if (usedDDWafConfigKeys.isEmpty() && !hasUserWafConfig && !defaultConfigActivated) {
+      // no config left in the WAF builder, add the default config
+      init();
+    }
+  }
+
   @Override
   public void init() {
     Map<String, Object> wafConfig;
@@ -341,8 +336,8 @@ public void init() {
     } else {
       hasUserWafConfig = true;
     }
-    this.mergedAsmFeatures = new MergedAsmFeatures();
-    this.initialized = true;
+    this.mergedAsmFeatures.clear();
+    this.usedDDWafConfigKeys.clear();
 
     if (wafConfig.isEmpty()) {
       throw new IllegalStateException("Expected default waf config to be available");

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

@@ -49,4 +49,9 @@ private void mergeAutoUserInstrum(
     }
     target.autoUserInstrum = newValue;
   }
+
+  public void clear() {
+    mergedData = null;
+    configs.clear();
+  }
 }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy

@@ -759,6 +759,67 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
     p.toFile().delete()
   }
 
+  // https://github.com/DataDog/dd-trace-java/issues/9159
+  void 'test initialization issues while applying remote config'() {
+    setup:
+    final key = new ParsedConfigKey('Test', '1234', 1, 'ASM_DD', 'ID')
+    final service = new AppSecConfigServiceImpl(config, poller, reconf)
+    config.getAppSecActivation() >> ProductActivation.ENABLED_INACTIVE
+
+    when:
+    service.maybeSubscribeConfigPolling()
+
+    then:
+    1 * poller.addListener(Product.ASM_DD, _) >> {
+      listeners.savedWafDataChangesListener = it[1]
+    }
+
+    when:
+    listeners.savedWafDataChangesListener.accept(key, '''{"rules_override": [{"rules_target": [{"rule_id": "foo"}], "enabled": false}]}'''.getBytes(), NOOP)
+
+    then:
+    noExceptionThrown()
+  }
+
+  void 'config keys are added and removed to the set when receiving ASM_DD payloads'() {
+    setup:
+    final key = new ParsedConfigKey('Test', '1234', 1, 'ASM_DD', 'ID')
+    final service = new AppSecConfigServiceImpl(config, poller, reconf)
+    config.getAppSecActivation() >> ProductActivation.ENABLED_INACTIVE
+
+    when:
+    service.maybeSubscribeConfigPolling()
+
+    then:
+    1 * poller.addListener(Product.ASM_DD, _) >> {
+      listeners.savedWafDataChangesListener = it[1]
+    }
+    1 * poller.addListener(Product.ASM_FEATURES, _, _) >> {
+      listeners.savedFeaturesDeserializer = it[1]
+      listeners.savedFeaturesListener = it[2]
+    }
+
+    when:
+    listeners.savedFeaturesListener.accept('asm_features conf',
+      listeners.savedFeaturesDeserializer.deserialize('{"asm":{"enabled": true}}'.bytes),
+      NOOP)
+
+    then:
+    service.usedDDWafConfigKeys.empty
+
+    when:
+    listeners.savedWafDataChangesListener.accept(key, '''{"rules_override": [{"rules_target": [{"rule_id": "foo"}], "enabled": false}]}'''.getBytes(), NOOP)
+
+    then:
+    service.usedDDWafConfigKeys.toList() == [key.toString()]
+
+    when:
+    listeners.savedWafDataChangesListener.remove(key, NOOP)
+
+    then:
+    service.usedDDWafConfigKeys.empty
+  }
+
   private static AppSecFeatures autoUserInstrum(String mode) {
     return new AppSecFeatures().tap { features ->
       features.autoUserInstrum = new AppSecFeatures.AutoUserInstrum().tap { instrum ->