Fix crash when remote config replaces the WAF handle during context creation

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

@@ -60,6 +60,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.StampedLock;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import org.slf4j.Logger;
@@ -97,6 +98,13 @@ private static class CtxAndAddresses {
     final Collection<Address<?>> addressesOfInterest;
     final WafHandle ctx;
 
+    /**
+     * Read-held while creating a context from {@link #ctx}; write-held while closing it. Keep the
+     * protected section callback-free: StampedLock is not reentrant. This protects context creation
+     * until libddwaf-java locks and checks the handle itself.
+     */
+    final StampedLock handleLock = new StampedLock();
+
     private CtxAndAddresses(Collection<Address<?>> addressesOfInterest, WafHandle ctx) {
       this.addressesOfInterest = addressesOfInterest;
       this.ctx = ctx;
@@ -208,7 +216,13 @@ private void initOrUpdateWafHandle(AppSecModuleConfigurer.Reconfiguration reconf
     }
 
     if (prevContextAndAddresses != null) {
-      prevContextAndAddresses.ctx.close();
+      // Context creation must finish acquiring native ruleset ownership before retiring the handle.
+      long stamp = prevContextAndAddresses.handleLock.writeLock();
+      try {
+        prevContextAndAddresses.ctx.close();
+      } finally {
+        prevContextAndAddresses.handleLock.unlockWrite(stamp);
+      }
     }
 
     reconf.reloadSubscriptions();
@@ -568,8 +582,30 @@ private Waf.ResultWithData doRunWaf(
         CtxAndAddresses ctxAndAddr,
         GatewayContext gwCtx)
         throws AbstractWafException {
-      WafContext wafContext =
-          reqCtx.getOrCreateWafContext(ctxAndAddr.ctx, wafMetricsEnabled, gwCtx.isRasp);
+      // Existing contexts own their native ruleset, so using one needs no lock.
+      WafContext wafContext = reqCtx.getWafContextIfReady(wafMetricsEnabled, gwCtx.isRasp);
+      if (wafContext == null) {
+        for (; ; ) {
+          long stamp = ctxAndAddr.handleLock.tryReadLock();
+          if (stamp == 0L) {
+            // Avoid waiting for destruction of a retired handle; retry with the current snapshot.
+            ctxAndAddr = ctxAndAddresses.get();
+            continue;
+          }
+          try {
+            // A callback may have captured this snapshot before remote config replaced it; never
+            // create a context from a retired handle.
+            if (ctxAndAddr == ctxAndAddresses.get()) {
+              wafContext =
+                  reqCtx.getOrCreateWafContext(ctxAndAddr.ctx, wafMetricsEnabled, gwCtx.isRasp);
+              break;
+            }
+          } finally {
+            ctxAndAddr.handleLock.unlockRead(stamp);
+          }
+          ctxAndAddr = ctxAndAddresses.get();
+        }
+      }
       if (wafContext == null) {
         // Context closed concurrently with the isWafContextClosed() check in onDataAvailable; skip
         // (APPSEC-69085).

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

@@ -393,7 +393,7 @@ public WafContext getOrCreateWafContext(
       if (wafContextClosed) {
         return null;
       }
-      if (createMetrics) {
+      if (!metricsReady(createMetrics, isRasp)) {
         if (wafMetrics == null) {
           this.wafMetrics = new WafMetrics();
         }
@@ -410,6 +410,24 @@ public WafContext getOrCreateWafContext(
     }
   }
 
+  /**
+   * Returns the request's open {@link WafContext} without locking, or {@code null} when it does not
+   * exist yet, is closed, or still lacks metrics that {@link #getOrCreateWafContext} would create.
+   * On {@code null}, callers fall back to {@link #getOrCreateWafContext}.
+   */
+  public WafContext getWafContextIfReady(boolean createMetrics, boolean isRasp) {
+    WafContext curWafContext = this.wafContext;
+    if (curWafContext == null || wafContextClosed) {
+      return null;
+    }
+    return metricsReady(createMetrics, isRasp) ? curWafContext : null;
+  }
+
+  /** Whether the metrics a WAF call with these flags needs already exist. */
+  private boolean metricsReady(boolean createMetrics, boolean isRasp) {
+    return !createMetrics || (wafMetrics != null && (!isRasp || raspMetrics != null));
+  }
+
   public void closeWafContext() {
     if (wafContextClosed) {
       // Fast path for the common case of redundant close() calls (e.g. the generic fallback