Fix WafContext TOCTOU race resurrecting orphaned native contexts

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

@@ -312,6 +312,14 @@ public void onDataAvailable(
 
       try {
         resultWithData = doRunWaf(reqCtx, newData, ctxAndAddr, gwCtx);
+        if (resultWithData == null) {
+          // WAF context closed concurrently between the fast-path check and context creation; skip
+          // (APPSEC-69085). raspRuleEval() was already counted above, so don't also count
+          // raspRuleSkipped() here - that counter is reserved for calls that never attempted eval.
+          log.debug("Skipped; the WAF context was closed concurrently");
+          WafMetricCollector.get().wafContextClosedRace();
+          return;
+        }
       } catch (TimeoutWafException tpe) {
         if (gwCtx.isRasp) {
           reqCtx.increaseRaspTimeouts();
@@ -322,10 +330,16 @@ public void onDataAvailable(
         }
         return;
       } catch (UnclassifiedWafException e) {
-        if (!reqCtx.isWafContextClosed()) {
+        if (reqCtx.isWafContextClosed()) {
+          // The context was closed concurrently between getOrCreateWafContext() and this run()
+          // call (APPSEC-69085) - the same benign race already tracked below via
+          // wafContextClosedRace(); avoid double-counting it as a real WAF error.
+          log.debug("Skipped; the WAF context was closed concurrently");
+          WafMetricCollector.get().wafContextClosedRace();
+        } else {
           log.error("Error calling WAF", e);
+          incrementErrorCodeMetric(reqCtx, gwCtx, e.code);
         }
-        incrementErrorCodeMetric(reqCtx, gwCtx, e.code);
         return;
       } catch (AbstractWafException e) {
         incrementErrorCodeMetric(reqCtx, gwCtx, e.code);
@@ -560,6 +574,11 @@ private Waf.ResultWithData doRunWaf(
         throws AbstractWafException {
       WafContext wafContext =
           reqCtx.getOrCreateWafContext(ctxAndAddr.ctx, wafMetricsEnabled, gwCtx.isRasp);
+      if (wafContext == null) {
+        // Context closed concurrently with the isWafContextClosed() check in onDataAvailable; skip
+        // (APPSEC-69085).
+        return null;
+      }
       WafMetrics metrics;
       if (gwCtx.isRasp) {
         metrics = reqCtx.getRaspMetrics();

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

@@ -353,39 +353,55 @@ public void setExtendedDataCollectionMaxHeaders(int extendedDataCollectionMaxHea
     this.extendedDataCollectionMaxHeaders = extendedDataCollectionMaxHeaders;
   }
 
+  /**
+   * Returns the request's {@link WafContext}, creating it on first use.
+   *
+   * <p>Returns {@code null} when the context has already been closed (see {@link
+   * #closeWafContext()}). Callers MUST treat a {@code null} return as "the WAF must not run for
+   * this request" and skip the evaluation. This prevents a late/async data event (e.g. a RASP
+   * callback on a driver or event-loop thread) from resurrecting a brand-new native {@code
+   * ddwaf_context} on an already-finished request, which would never be closed and would leak
+   * off-heap memory (APPSEC-69085).
+   */
   public WafContext getOrCreateWafContext(
       WafHandle wafHandle, boolean createMetrics, boolean isRasp) {
-    if (createMetrics) {
-      if (wafMetrics == null) {
-        this.wafMetrics = new WafMetrics();
+    synchronized (this) {
+      // Atomic with respect to closeWafContext(): both run under this monitor.
+      if (wafContextClosed) {
+        return null;
       }
-      if (isRasp && raspMetrics == null) {
-        this.raspMetrics = new WafMetrics();
+      if (createMetrics) {
+        if (wafMetrics == null) {
+          this.wafMetrics = new WafMetrics();
+        }
+        if (isRasp && raspMetrics == null) {
+          this.raspMetrics = new WafMetrics();
+        }
       }
-    }
-
-    WafContext curWafContext;
-    synchronized (this) {
-      curWafContext = this.wafContext;
-      if (curWafContext != null) {
-        return curWafContext;
+      if (this.wafContext != null) {
+        return this.wafContext;
       }
-      curWafContext = new WafContext(wafHandle);
+      WafContext curWafContext = new WafContext(wafHandle);
       this.wafContext = curWafContext;
+      return curWafContext;
     }
-    return curWafContext;
   }
 
   public void closeWafContext() {
-    if (wafContext != null) {
-      synchronized (this) {
-        if (wafContext != null) {
-          try {
-            wafContextClosed = true;
-            wafContext.close();
-          } finally {
-            wafContext = null;
-          }
+    if (wafContextClosed) {
+      // Fast path for the common case of redundant close() calls (e.g. the generic fallback
+      // close() running after GatewayBridge#onRequestEnded already closed it).
+      return;
+    }
+    synchronized (this) {
+      // Must be set unconditionally, even if the WAF never ran for this request: a late/async
+      // caller of getOrCreateWafContext() must not resurrect a context after close (APPSEC-69085).
+      wafContextClosed = true;
+      if (wafContext != null) {
+        try {
+          wafContext.close();
+        } finally {
+          wafContext = null;
         }
       }
     }
@@ -720,8 +736,10 @@ public void close() {
       if (wafContext != null) {
         log.debug(
             SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)");
-        closeWafContext();
       }
+      // Always close, even if the WAF never ran for this request: wafContextClosed must be set so
+      // a late/async caller of getOrCreateWafContext() cannot resurrect a context (APPSEC-69085).
+      closeWafContext();
       collectedCookies = null;
       requestHeaders.clear();
       responseHeaders.clear();

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

@@ -280,7 +280,7 @@ class WAFModuleSpecification extends DDSpecification {
       rba.statusCode == 403 &&
       rba.blockingContentType == BlockingContentType.AUTO
     })
-    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false)
+    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * tracer.activeSpan()
     1 * ctx.reportEvents(_ as Collection<AppSecEvent>)
     2 * ctx.getWafMetrics()
@@ -304,7 +304,7 @@ class WAFModuleSpecification extends DDSpecification {
       rba.statusCode == 403 &&
       rba.blockingContentType == BlockingContentType.AUTO
     })
-    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false)
+    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * tracer.activeSpan()
     1 * ctx.reportEvents(_ as Collection<AppSecEvent>)
     2 * ctx.getWafMetrics()
@@ -356,7 +356,7 @@ class WAFModuleSpecification extends DDSpecification {
       rba.statusCode == 403 &&
       rba.blockingContentType == BlockingContentType.AUTO
     })
-    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false)
+    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * tracer.activeSpan()
     1 * ctx.reportEvents(_ as Collection<AppSecEvent>)
     2 * ctx.getWafMetrics()
@@ -376,7 +376,7 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then:
-    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false)
+    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics()
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
@@ -428,7 +428,7 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then:
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * tracer.activeSpan()
     1 * ctx.reportEvents(_ as Collection<AppSecEvent>)
     2 * ctx.getWafMetrics()
@@ -450,7 +450,7 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then:
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics()
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
@@ -1074,6 +1074,7 @@ class WAFModuleSpecification extends DDSpecification {
   void 'reloading rules clears waf data and rule toggling'() {
     initialRuleAdd()
     ChangeableFlow flow = Mock()
+    WafContext wafContext
     def ipData = [
       rules_data :
       [
@@ -1116,10 +1117,12 @@ class WAFModuleSpecification extends DDSpecification {
     then: 'no match; rule is disabled'
     1 * wafMetricCollector.wafUpdates(_, true)
     1 * reconf.reloadSubscriptions()
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics()
     1 * ctx.isWafContextClosed() >> false
-    1 * ctx.closeWafContext()
+    1 * ctx.closeWafContext() >> {
+      wafContext.close()
+    }
     _ * ctx.increaseWafTimeouts()
     _ * ctx.increaseRaspTimeouts()
     0 * _
@@ -1132,11 +1135,13 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then: 'no match; data was cleared (though rule is no longer disabled)'
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     1 * ctx.isWafContextClosed() >> false
     1 * wafMetricCollector.wafUpdates(_, true)
     1 * reconf.reloadSubscriptions()
-    1 * ctx.closeWafContext()
+    1 * ctx.closeWafContext() >> {
+      wafContext.close()
+    }
     2 * ctx.getWafMetrics()
     _ * ctx.increaseWafTimeouts()
     _ * ctx.increaseRaspTimeouts()
@@ -1151,13 +1156,15 @@ class WAFModuleSpecification extends DDSpecification {
     then: 'now we have match'
     1 * wafMetricCollector.wafUpdates(_, true)
     1 * reconf.reloadSubscriptions()
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * tracer.activeSpan()
     1 * ctx.reportEvents(_ as Collection<AppSecEvent>)
     2 * ctx.getWafMetrics()
     1 * flow.setAction({ it.blocking })
     1 * ctx.isWafContextClosed() >> false
-    1 * ctx.closeWafContext()
+    1 * ctx.closeWafContext() >> {
+      wafContext.close()
+    }
     1 * flow.isBlocking()
     1 * ctx.isThrottled(null)
     1 * ctx.setManuallyKept(true)
@@ -1174,10 +1181,12 @@ class WAFModuleSpecification extends DDSpecification {
     then: 'nothing again; we disabled the rule'
     1 * wafMetricCollector.wafUpdates(_, true)
     1 * reconf.reloadSubscriptions()
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics()
     1 * ctx.isWafContextClosed() >> false
-    1 * ctx.closeWafContext()
+    1 * ctx.closeWafContext() >> {
+      wafContext.close()
+    }
     _ * ctx.increaseWafTimeouts()
     _ * ctx.increaseRaspTimeouts()
     0 * _
@@ -1490,7 +1499,7 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then:
-    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false)
+    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics()
     1 * ctx.isThrottled(null)
     1 * ctx.setManuallyKept(true)
@@ -1523,7 +1532,7 @@ class WAFModuleSpecification extends DDSpecification {
     })
     1 * flow.isBlocking()
     1 * ctx.isWafContextClosed() >> false
-    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false)
+    1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics()
     1 * ctx.isThrottled(null)
     1 * ctx.setManuallyKept(true)
@@ -1538,7 +1547,6 @@ class WAFModuleSpecification extends DDSpecification {
     final flow = Mock(ChangeableFlow)
     final fingerprint = '_dd.appsec.fp.http.endpoint'
     initialRuleAdd 'fingerprint_config.json'
-    ctx.closeWafContext()
     final bundle = MapDataBundle.ofDelegate([
       (KnownAddresses.WAF_CONTEXT_PROCESSOR): [fingerprint: true],
       (KnownAddresses.REQUEST_METHOD): 'GET',
@@ -1567,7 +1575,6 @@ class WAFModuleSpecification extends DDSpecification {
     final sessionId = UUID.randomUUID().toString()
     initialRuleAdd 'fingerprint_config.json'
     wafModule.applyConfig(reconf)
-    ctx.closeWafContext()
     final bundle = MapDataBundle.ofDelegate([
       (KnownAddresses.WAF_CONTEXT_PROCESSOR): [fingerprint: true],
       (KnownAddresses.REQUEST_COOKIES): [JSESSIONID: [sessionId]],
@@ -1905,7 +1912,7 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then:
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics() >> metrics
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
@@ -1923,7 +1930,7 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then:
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics() >> metrics
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
@@ -1942,7 +1949,7 @@ class WAFModuleSpecification extends DDSpecification {
     ctx.closeWafContext()
 
     then:
-    1 * ctx.getOrCreateWafContext(_, true, false)
+    1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) }
     2 * ctx.getWafMetrics() >> metrics
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()