Add comments on missing bits for error handling

Remove unused WafHandle field

Fix numLoadedRules logging

Undo unnecessary changes (fixes memory leaks, race conditions)

spotless

Just capture ASM_DD listener in test via ConfigurationPoller stub

Mock ConfigurationPoller in tests

Author: smola

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

@@ -60,6 +60,8 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+
+import datadog.trace.api.telemetry.LogCollector;
 import okio.Okio;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -97,7 +99,6 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
   private boolean defaultConfigActivated;
   private final Set<String> usedDDWafConfigKeys = new HashSet<>();
   private final String DEFAULT_WAF_CONFIG_RULE = "DEFAULT_WAF_CONFIG";
-  private final AsmDDTypedListener asmDDTypedListener;
   private String currentRuleVersion;
   private List<AppSecModule> modulesToUpdateVersionIn;
 
@@ -112,7 +113,6 @@ public AppSecConfigServiceImpl(
     if (tracerConfig.isAppSecWafMetrics()) {
       traceSegmentPostProcessors.add(statsReporter);
     }
-    asmDDTypedListener = new AsmDDTypedListener();
   }
 
   private void subscribeConfigurationPoller() {
@@ -156,7 +156,7 @@ private void subscribeConfigurationPoller() {
   }
 
   private void subscribeRulesAndData() {
-    this.configurationPoller.addListener(Product.ASM_DD, asmDDTypedListener);
+    this.configurationPoller.addListener(Product.ASM_DD, new AsmDDTypedListener());
     this.configurationPoller.addListener(
         Product.ASM_DATA, new AppSecConfigConfigurationChangesTypedListener());
     this.configurationPoller.addListener(
@@ -233,8 +233,7 @@ public void remove(ConfigKey configKey, PollingRateHinter pollingRateHinter)
     }
   }
 
-  private void handleWafUpdateResultReport(
-      String configKey, Map<String, Object> rawConfig, String... filename)
+  private void handleWafUpdateResultReport(String configKey, Map<String, Object> rawConfig)
       throws AppSecModule.AppSecModuleActivationException {
     wafBuilder = getWafBuilder();
     if (modulesToUpdateVersionIn != null
@@ -244,9 +243,13 @@ private void handleWafUpdateResultReport(
     }
     try {
       WafDiagnostics wafDiagnostics = wafBuilder.addOrUpdateConfig(configKey, rawConfig);
-      if (log.isInfoEnabled() && filename.length > 0) {
-        StandardizedLogging.numLoadedRules(log, filename[0], countRules(rawConfig));
+      if (log.isInfoEnabled()) {
+        StandardizedLogging.numLoadedRules(log, configKey, countRules(rawConfig));
       }
+
+      // TODO: Send diagnostics via telemetry
+      final LogCollector telemetryLogger = LogCollector.get();
+
       initReporter.setReportForPublication(wafDiagnostics);
       if (wafDiagnostics.rulesetVersion != null
           && !wafDiagnostics.rulesetVersion.isEmpty()
@@ -263,6 +266,9 @@ private void handleWafUpdateResultReport(
           "Invalid rule during waf config update for config key {}: {}",
           configKey,
           e.wafDiagnostics);
+
+      // TODO: Propagate diagostics back to remote config apply_error
+
       initReporter.setReportForPublication(e.wafDiagnostics);
       throw new RuntimeException(e);
     } catch (UnclassifiedWafException e) {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java

@@ -82,7 +82,6 @@ public class WAFModule implements AppSecModule {
   private static final Map<String, ActionInfo> DEFAULT_ACTIONS;
 
   private static final String EXPLOIT_DETECTED_MSG = "Exploit detected";
-  private WafHandle wafHandle;
   private boolean init = true;
   private String rulesetVersion;
   private WafBuilder wafBuilder;
@@ -99,9 +98,11 @@ private ActionInfo(String type, Map<String, Object> parameters) {
 
   private static class CtxAndAddresses {
     final Collection<Address<?>> addressesOfInterest;
+    final WafHandle ctx;
 
-    private CtxAndAddresses(Collection<Address<?>> addressesOfInterest) {
+    private CtxAndAddresses(Collection<Address<?>> addressesOfInterest, WafHandle ctx) {
       this.addressesOfInterest = addressesOfInterest;
+      this.ctx = ctx;
     }
   }
 
@@ -137,6 +138,8 @@ static void createLimitsObject() {
             Config.get().getAppSecWafTimeout());
   }
 
+  private final boolean wafMetricsEnabled =
+      Config.get().isAppSecWafMetrics(); // could be static if not for tests
   private final AtomicReference<CtxAndAddresses> ctxAndAddresses = new AtomicReference<>();
   private final RateLimiter rateLimiter;
 
@@ -197,24 +200,27 @@ private void applyConfig(AppSecModuleConfigurer.Reconfiguration reconf)
 
   private void initOrUpdateWafHandle(AppSecModuleConfigurer.Reconfiguration reconf)
       throws AppSecModuleActivationException {
-    WafHandle oldHandle = wafHandle;
+    CtxAndAddresses prevContextAndAddresses = this.ctxAndAddresses.get();
+    WafHandle newHandle;
     try {
-      wafHandle = wafBuilder.buildWafHandleInstance();
-      reconf.reloadSubscriptions();
+      newHandle = wafBuilder.buildWafHandleInstance();
     } catch (AbstractWafException e) {
-      log.info("Could not initialize waf handle, no rules were added!", e);
       throw new AppSecModuleActivationException(
           "Could not initialize waf handle, no rules were added!", e);
-    } finally {
-      if (oldHandle != null && oldHandle.isOnline()) {
-        oldHandle.close();
-      }
     }
 
-    CtxAndAddresses newContextAndAddresses = new CtxAndAddresses(getUsedAddresses(wafHandle));
-    if (!this.ctxAndAddresses.compareAndSet(ctxAndAddresses.get(), newContextAndAddresses)) {
+    Collection<Address<?>> addresses = getUsedAddresses(newHandle);
+    CtxAndAddresses newContextAndAddresses = new CtxAndAddresses(addresses, newHandle);
+    if (!this.ctxAndAddresses.compareAndSet(prevContextAndAddresses, newContextAndAddresses)) {
+      newHandle.close();
       throw new AppSecModuleActivationException("Concurrent update of WAF configuration");
     }
+
+    if (prevContextAndAddresses != null) {
+      prevContextAndAddresses.ctx.close();
+    }
+
+    reconf.reloadSubscriptions();
   }
 
   private static RateLimiter getRateLimiter(Monitoring monitoring) {
@@ -289,13 +295,19 @@ public void onDataAvailable(
         DataBundle newData,
         GatewayContext gwCtx) {
       Waf.ResultWithData resultWithData;
+      CtxAndAddresses ctxAndAddr = ctxAndAddresses.get();
+      if (ctxAndAddr == null) {
+        log.debug("Skipped; the WAF is not configured");
+        return;
+      }
       if (reqCtx.isWafContextClosed()) {
         log.debug("Skipped; the WAF context is closed");
         if (gwCtx.isRasp) {
           WafMetricCollector.get().raspRuleSkipped(gwCtx.raspRuleType);
         }
         return;
       }
+
       StandardizedLogging.executingWAF(log);
       long start = 0L;
       if (log.isDebugEnabled()) {
@@ -307,7 +319,7 @@ public void onDataAvailable(
       }
 
       try {
-        resultWithData = doRunWaf(reqCtx, newData, gwCtx);
+        resultWithData = doRunWaf(reqCtx, newData, ctxAndAddr, gwCtx);
       } catch (TimeoutWafException tpe) {
         if (gwCtx.isRasp) {
           reqCtx.increaseRaspTimeouts();
@@ -497,8 +509,13 @@ private StackTraceEvent createExploitStackTraceEvent(String stackId) {
     }
 
     private Waf.ResultWithData doRunWaf(
-        AppSecRequestContext reqCtx, DataBundle newData, GatewayContext gwCtx)
+        AppSecRequestContext reqCtx,
+        DataBundle newData,
+        CtxAndAddresses ctxAndAddr,
+        GatewayContext gwCtx)
         throws AbstractWafException {
+      WafContext wafContext =
+          reqCtx.getOrCreateWafContext(ctxAndAddr.ctx, wafMetricsEnabled, gwCtx.isRasp);
       WafMetrics metrics;
       if (gwCtx.isRasp) {
         metrics = reqCtx.getRaspMetrics();
@@ -508,18 +525,17 @@ private Waf.ResultWithData doRunWaf(
       }
 
       if (gwCtx.isTransient) {
-        return runWafTransient(
-            reqCtx.getOrCreateWafContext(wafHandle, gwCtx.isRasp), metrics, newData);
+        return runWafTransient(wafContext, metrics, newData, ctxAndAddr);
       } else {
-        return runWafContext(
-            reqCtx.getOrCreateWafContext(wafHandle, gwCtx.isRasp), metrics, newData);
+        return runWafContext(wafContext, metrics, newData, ctxAndAddr);
       }
     }
 
     private Waf.ResultWithData runWafContext(
-        WafContext wafContext, WafMetrics metrics, DataBundle newData) throws AbstractWafException {
+        WafContext wafContext, WafMetrics metrics, DataBundle newData, CtxAndAddresses ctxAndAddr)
+        throws AbstractWafException {
       return wafContext.run(
-          new DataBundleMapWrapper(getUsedAddresses(wafHandle), newData), LIMITS, metrics);
+          new DataBundleMapWrapper(ctxAndAddr.addressesOfInterest, newData), LIMITS, metrics);
     }
   }
 
@@ -534,9 +550,10 @@ private static void incrementErrorCodeMetric(
   }
 
   private Waf.ResultWithData runWafTransient(
-      WafContext wafContext, WafMetrics metrics, DataBundle newData) throws AbstractWafException {
+      WafContext wafContext, WafMetrics metrics, DataBundle newData, CtxAndAddresses ctxAndAddr)
+      throws AbstractWafException {
     return wafContext.runEphemeral(
-        new DataBundleMapWrapper(getUsedAddresses(wafHandle), newData), LIMITS, metrics);
+        new DataBundleMapWrapper(ctxAndAddr.addressesOfInterest, newData), LIMITS, metrics);
   }
 
   private Collection<AppSecEvent> buildEvents(Waf.ResultWithData actionWithData) {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java

@@ -239,8 +239,9 @@ public int getRaspTimeouts() {
     return raspTimeouts;
   }
 
-  public WafContext getOrCreateWafContext(WafHandle wafHandle, boolean isRasp) {
-    if (Config.get().isAppSecWafMetrics()) {
+  public WafContext getOrCreateWafContext(
+      WafHandle wafHandle, boolean createMetrics, boolean isRasp) {
+    if (createMetrics) {
       if (wafMetrics == null) {
         this.wafMetrics = new WafMetrics();
       }