From ce22800c262e60caa40b94d8095a0ad5e44f2485 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Tue, 21 Jul 2026 15:45:55 +0200 Subject: [PATCH 1/4] Fix WafContext TOCTOU race resurrecting orphaned native contexts (APPSEC-69085) getOrCreateWafContext() created a new native WafContext whenever the wafContext field was null, without checking wafContextClosed. A late/async RASP callback (e.g. reactive JDBC/r2dbc, async HTTP client) could race with closeWafContext() and resurrect a brand-new context on an already-finished request, which was never closed and leaked native memory. Move the wafContextClosed check to the top of the existing synchronized block in getOrCreateWafContext, returning null if the context is already closed. WAFModule.doRunWaf and WAFDataCallback.onDataAvailable now treat a null WafContext as a skip instead of dereferencing it. --- .../com/datadog/appsec/ddwaf/WAFModule.java | 14 ++ .../appsec/gateway/AppSecRequestContext.java | 39 +++-- .../ddwaf/WAFModuleSpecification.groovy | 30 ++-- ...ppSecRequestContextWafContextRaceTest.java | 159 ++++++++++++++++++ 4 files changed, 213 insertions(+), 29 deletions(-) create mode 100644 dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java index a7aac0a6f98..6425bb95e06 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java @@ -312,6 +312,15 @@ 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). + log.debug("Skipped; the WAF context was closed concurrently"); + if (gwCtx.isRasp) { + WafMetricCollector.get().raspRuleSkipped(gwCtx.raspRuleType); + } + return; + } } catch (TimeoutWafException tpe) { if (gwCtx.isRasp) { reqCtx.increaseRaspTimeouts(); @@ -560,6 +569,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(); diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index 3639ce80c5e..28c347cfafc 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java @@ -353,27 +353,38 @@ public void setExtendedDataCollectionMaxHeaders(int extendedDataCollectionMaxHea this.extendedDataCollectionMaxHeaders = extendedDataCollectionMaxHeaders; } + /** + * Returns the request's {@link WafContext}, creating it on first use. + * + *

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() { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy index 9baa6dfa8f7..23a9f2db428 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy +++ b/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) 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) 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) 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) 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() @@ -1116,7 +1116,7 @@ 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() @@ -1132,7 +1132,7 @@ 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() @@ -1151,7 +1151,7 @@ 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) 2 * ctx.getWafMetrics() @@ -1174,7 +1174,7 @@ 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() @@ -1490,7 +1490,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 +1523,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) @@ -1905,7 +1905,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 +1923,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 +1942,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() diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java new file mode 100644 index 00000000000..752f2f20fb3 --- /dev/null +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java @@ -0,0 +1,159 @@ +package com.datadog.appsec.gateway; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.datadog.appsec.ddwaf.WafInitialization; +import com.datadog.ddwaf.Waf; +import com.datadog.ddwaf.WafBuilder; +import com.datadog.ddwaf.WafContext; +import com.datadog.ddwaf.WafHandle; +import com.squareup.moshi.JsonAdapter; +import com.squareup.moshi.Moshi; +import com.squareup.moshi.Types; +import java.io.InputStream; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import okio.Okio; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Contract and concurrency tests for the {@code getOrCreateWafContext}/{@code closeWafContext} + * lifecycle, guarding against the TOCTOU race that used to resurrect an orphaned native {@link + * WafContext} on already-closed requests (APPSEC-69085). + */ +class AppSecRequestContextWafContextRaceTest { + + private static final JsonAdapter> ADAPTER = + new Moshi.Builder() + .build() + .adapter(Types.newParameterizedType(Map.class, String.class, Object.class)); + + private WafBuilder wafBuilder; + private WafHandle wafHandle; + + @BeforeEach + void setup() throws Exception { + // Force the native library to load and initialize. + assertTrue(WafInitialization.ONLINE, "libddwaf must be available for this test"); + Waf.initialize(false); + wafBuilder = new WafBuilder(); + try (InputStream stream = + getClass().getClassLoader().getResourceAsStream("test_multi_config.json")) { + wafBuilder.addOrUpdateConfig("test", ADAPTER.fromJson(Okio.buffer(Okio.source(stream)))); + } + wafHandle = wafBuilder.buildWafHandleInstance(); + } + + @AfterEach + void tearDown() { + if (wafBuilder != null) { + wafBuilder.close(); + } + } + + @Test + void firstUseCreatesContext() { + AppSecRequestContext ctx = new AppSecRequestContext(); + try { + WafContext created = ctx.getOrCreateWafContext(wafHandle, false, false); + assertNotNull(created); + assertTrue(created.isOnline()); + assertFalse(ctx.isWafContextClosed()); + } finally { + ctx.closeWafContext(); + } + } + + @Test + void reuseBeforeCloseReturnsSameInstance() { + AppSecRequestContext ctx = new AppSecRequestContext(); + try { + WafContext first = ctx.getOrCreateWafContext(wafHandle, false, false); + WafContext second = ctx.getOrCreateWafContext(wafHandle, false, false); + assertNotNull(first); + assertSame(first, second); + } finally { + ctx.closeWafContext(); + } + } + + @Test + void rejectsCreationAfterClose() { + AppSecRequestContext ctx = new AppSecRequestContext(); + WafContext created = ctx.getOrCreateWafContext(wafHandle, false, false); + assertNotNull(created); + + ctx.closeWafContext(); + assertTrue(ctx.isWafContextClosed()); + assertFalse(created.isOnline()); + + // A late/async caller must not resurrect a brand-new orphaned context. + WafContext afterClose = ctx.getOrCreateWafContext(wafHandle, false, false); + assertNull(afterClose); + assertTrue(ctx.isWafContextClosed()); + } + + /** + * Drives {@code closeWafContext()} and {@code getOrCreateWafContext()} concurrently via a barrier + * across many iterations. The invariant: whatever {@code getOrCreateWafContext} returns after the + * race is either the ORIGINAL context (same instance, still online) or {@code null} - never a + * second/different instance. At the end, no observed context may still be online (all created + * contexts were eventually closed - no orphans). + */ + @Test + void concurrentCloseNeverCreatesOrphan() throws Exception { + final int iterations = 5_000; + ExecutorService pool = Executors.newFixedThreadPool(2); + List observed = new ArrayList<>(iterations); + try { + for (int i = 0; i < iterations; i++) { + final AppSecRequestContext ctx = new AppSecRequestContext(); + WafContext original = ctx.getOrCreateWafContext(wafHandle, false, false); + assertNotNull(original); + observed.add(original); + + final CyclicBarrier barrier = new CyclicBarrier(2); + Future closer = + pool.submit( + () -> { + barrier.await(); + ctx.closeWafContext(); + return null; + }); + Future creator = + pool.submit( + () -> { + barrier.await(); + return ctx.getOrCreateWafContext(wafHandle, false, false); + }); + + closer.get(); + WafContext late = creator.get(); + + // Never a freshly created second context: only the original or null. + if (late != null) { + assertSame(original, late, "getOrCreateWafContext resurrected a new orphan context"); + } + assertTrue(ctx.isWafContextClosed()); + } + } finally { + pool.shutdownNow(); + } + + // Global invariant: every created context was eventually closed - no orphan left online. + for (WafContext context : observed) { + assertFalse(context.isOnline(), "orphaned WafContext left online (never closed)"); + } + } +} From 562351a83bbb4a330813c6699871ee0f6d6a3272 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Wed, 22 Jul 2026 10:45:33 +0200 Subject: [PATCH 2/4] review: pre-PR checks - Set wafContextClosed unconditionally in closeWafContext(), even if no WafContext was ever created, preventing a late/async caller from resurrecting an orphaned native context on an already-finished request - Add regression test for close-before-first-use ordering - Add WAFModule-level test covering the doRunWaf null-skip path when the context is closed concurrently - Fix native WafContext leaks in 4 stub sites of 'reloading rules clears waf data and rule toggling' by wiring ctx.closeWafContext() to actually close the ephemeral context, and shadow the class-level wafContext field with a local to avoid a cleanup() double-close - Remove premature closeWafContext() calls in two fingerprint tests that relied on the pre-fix silent no-op behavior - Move the WAF-context-closed-race metric test from Groovy to JUnit5 Java --- .../com/datadog/appsec/ddwaf/WAFModule.java | 1 + .../appsec/gateway/AppSecRequestContext.java | 18 +-- .../ddwaf/WAFModuleSpecification.groovy | 19 ++- .../ddwaf/WAFModuleContextClosedRaceTest.java | 123 ++++++++++++++++++ ...ppSecRequestContextWafContextRaceTest.java | 14 ++ .../api/telemetry/WafMetricCollector.java | 24 ++++ ...fMetricCollectorContextClosedRaceTest.java | 29 +++++ 7 files changed, 213 insertions(+), 15 deletions(-) create mode 100644 dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java create mode 100644 internal-api/src/test/java/datadog/trace/api/telemetry/WafMetricCollectorContextClosedRaceTest.java diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java index 6425bb95e06..dc811d905b4 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java @@ -316,6 +316,7 @@ public void onDataAvailable( // WAF context closed concurrently between the fast-path check and context creation; skip // (APPSEC-69085). log.debug("Skipped; the WAF context was closed concurrently"); + WafMetricCollector.get().wafContextClosedRace(); if (gwCtx.isRasp) { WafMetricCollector.get().raspRuleSkipped(gwCtx.raspRuleType); } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index 28c347cfafc..c4679fd84ff 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java @@ -388,15 +388,15 @@ public WafContext getOrCreateWafContext( } public void closeWafContext() { - if (wafContext != null) { - synchronized (this) { - if (wafContext != null) { - try { - wafContextClosed = true; - wafContext.close(); - } finally { - wafContext = null; - } + 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; } } } diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy index 23a9f2db428..93ea2f56bf5 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy @@ -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 : [ @@ -1119,7 +1120,9 @@ class WAFModuleSpecification extends DDSpecification { 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 * _ @@ -1136,7 +1139,9 @@ class WAFModuleSpecification extends DDSpecification { 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() @@ -1157,7 +1162,9 @@ class WAFModuleSpecification extends DDSpecification { 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) @@ -1177,7 +1184,9 @@ class WAFModuleSpecification extends DDSpecification { 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 * _ @@ -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]], diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java new file mode 100644 index 00000000000..f069f975151 --- /dev/null +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java @@ -0,0 +1,123 @@ +package com.datadog.appsec.ddwaf; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.datadog.appsec.config.AppSecModuleConfigurer; +import com.datadog.appsec.config.TraceSegmentPostProcessor; +import com.datadog.appsec.event.ChangeableFlow; +import com.datadog.appsec.event.DataListener; +import com.datadog.appsec.event.data.MapDataBundle; +import com.datadog.appsec.gateway.AppSecRequestContext; +import com.datadog.appsec.gateway.GatewayContext; +import com.datadog.ddwaf.Waf; +import com.datadog.ddwaf.WafBuilder; +import com.squareup.moshi.JsonAdapter; +import com.squareup.moshi.Moshi; +import com.squareup.moshi.Types; +import datadog.trace.api.telemetry.RuleType; +import datadog.trace.api.telemetry.WafMetricCollector; +import java.io.InputStream; +import java.util.Collections; +import java.util.Map; +import okio.Okio; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Covers the {@code WAFModule.WAFDataCallback.onDataAvailable} null-skip branch exercised when + * {@code doRunWaf} returns {@code null} because the {@code WafContext} was closed concurrently + * between the {@code isWafContextClosed()} fast-path check and context creation (APPSEC-69085). + */ +class WAFModuleContextClosedRaceTest { + + private static final JsonAdapter> ADAPTER = + new Moshi.Builder() + .build() + .adapter(Types.newParameterizedType(Map.class, String.class, Object.class)); + + private WafBuilder wafBuilder; + private WAFModule wafModule; + private DataListener dataListener; + + @BeforeEach + void setup() throws Exception { + assertTrue(WafInitialization.ONLINE, "libddwaf must be available for this test"); + Waf.initialize(false); + wafBuilder = new WafBuilder(); + try (InputStream stream = + getClass().getClassLoader().getResourceAsStream("test_multi_config.json")) { + wafBuilder.addOrUpdateConfig("test", ADAPTER.fromJson(Okio.buffer(Okio.source(stream)))); + } + + wafModule = new WAFModule(); + wafModule.setWafBuilder(wafBuilder); + AppSecModuleConfigurer.SubconfigListener[] captured = + new AppSecModuleConfigurer.SubconfigListener[1]; + wafModule.config( + new AppSecModuleConfigurer() { + @Override + public void addSubConfigListener( + String key, AppSecModuleConfigurer.SubconfigListener listener) { + captured[0] = listener; + } + + @Override + public void addTraceSegmentPostProcessor(TraceSegmentPostProcessor interceptor) {} + }); + captured[0].onNewSubconfig(null, AppSecModuleConfigurer.Reconfiguration.NOOP); + dataListener = wafModule.getDataSubscriptions().iterator().next(); + } + + @AfterEach + void tearDown() { + if (wafBuilder != null) { + wafBuilder.close(); + } + } + + @Test + void skipsAndIncrementsCounterWhenContextClosedConcurrently() { + AppSecRequestContext reqCtx = mock(AppSecRequestContext.class); + when(reqCtx.isWafContextClosed()).thenReturn(false); + when(reqCtx.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())).thenReturn(null); + + ChangeableFlow flow = new ChangeableFlow(); + GatewayContext gwCtx = new GatewayContext(false); + + dataListener.onDataAvailable( + flow, reqCtx, MapDataBundle.ofDelegate(Collections.emptyMap()), gwCtx); + + assertFalse(flow.isBlocking()); + WafMetricCollector.get().prepareMetrics(); + boolean sawContextClosedRace = + WafMetricCollector.get().drain().stream() + .anyMatch(m -> "waf.context_closed_race".equals(m.metricName)); + assertTrue(sawContextClosedRace, "expected waf.context_closed_race to be reported"); + } + + @Test + void skipsRaspRuleWhenContextClosedConcurrently() { + AppSecRequestContext reqCtx = mock(AppSecRequestContext.class); + when(reqCtx.isWafContextClosed()).thenReturn(false); + when(reqCtx.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())).thenReturn(null); + + ChangeableFlow flow = new ChangeableFlow(); + GatewayContext gwCtx = new GatewayContext(false, RuleType.LFI); + + dataListener.onDataAvailable( + flow, reqCtx, MapDataBundle.ofDelegate(Collections.emptyMap()), gwCtx); + + assertFalse(flow.isBlocking()); + WafMetricCollector.get().prepareMetrics(); + boolean sawRaspSkipped = + WafMetricCollector.get().drain().stream() + .anyMatch(m -> "rasp.rule.skipped".equals(m.metricName)); + assertTrue(sawRaspSkipped, "expected rasp.rule.skipped to be reported"); + } +} diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java index 752f2f20fb3..347e3986c20 100644 --- a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java @@ -104,6 +104,20 @@ void rejectsCreationAfterClose() { assertTrue(ctx.isWafContextClosed()); } + @Test + void closeBeforeFirstUsePreventsLaterCreation() { + AppSecRequestContext ctx = new AppSecRequestContext(); + + // Close before the WAF ever ran for this request (e.g. an early-blocked request). + ctx.closeWafContext(); + assertTrue(ctx.isWafContextClosed()); + + // A late/async caller must not create a brand-new orphaned context after close. + WafContext afterClose = ctx.getOrCreateWafContext(wafHandle, false, false); + assertNull(afterClose); + assertTrue(ctx.isWafContextClosed()); + } + /** * Drives {@code closeWafContext()} and {@code getOrCreateWafContext()} concurrently via a barrier * across many iterations. The invariant: whatever {@code getOrCreateWafContext} returns after the diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java index d4792a5bd69..61a3af09896 100644 --- a/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java +++ b/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java @@ -58,6 +58,7 @@ private WafMetricCollector() { private static final AtomicLongArray appSecSdkEventQueue = new AtomicLongArray(LoginEvent.getNumValues() * LoginVersion.getNumValues()); private static final AtomicInteger wafConfigErrorCounter = new AtomicInteger(); + private static final AtomicInteger contextClosedRaceCounter = new AtomicInteger(); private static final AtomicLongArray aiGuardRequests = new AtomicLongArray(AIGuard.Action.values().length * 2); // 3 actions * block private static final AtomicInteger aiGuardErrors = new AtomicInteger(); @@ -391,6 +392,14 @@ public void prepareMetrics() { } } + // WafContext closed-concurrently race (APPSEC-69085) + int contextClosedRace = contextClosedRaceCounter.getAndSet(0); + if (contextClosedRace > 0) { + if (!rawMetricsQueue.offer(new ContextClosedRace(contextClosedRace))) { + return; + } + } + // AI Guard successful requests for (final AIGuard.Action action : AIGuard.Action.values()) { final long blocked = aiGuardRequests.getAndSet(action.ordinal() * 2 + 1, 0); @@ -525,6 +534,21 @@ public void addWafConfigError(int nbErrors) { wafConfigErrorCounter.addAndGet(nbErrors); } + /** + * Records that {@code getOrCreateWafContext} rejected a run because the request's {@code + * WafContext} was already closed concurrently (APPSEC-69085). Used to measure the frequency of + * this race in production. + */ + public void wafContextClosedRace() { + contextClosedRaceCounter.incrementAndGet(); + } + + public static class ContextClosedRace extends WafMetric { + public ContextClosedRace(final long counter) { + super("waf.context_closed_race", counter); + } + } + public static class WafConfigError extends WafMetric { public WafConfigError(final long counter, final String wafVersion, final String rulesVersion) { super( diff --git a/internal-api/src/test/java/datadog/trace/api/telemetry/WafMetricCollectorContextClosedRaceTest.java b/internal-api/src/test/java/datadog/trace/api/telemetry/WafMetricCollectorContextClosedRaceTest.java new file mode 100644 index 00000000000..c1b8a06a99c --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/telemetry/WafMetricCollectorContextClosedRaceTest.java @@ -0,0 +1,29 @@ +package datadog.trace.api.telemetry; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Collection; +import java.util.Optional; +import org.junit.jupiter.api.Test; + +class WafMetricCollectorContextClosedRaceTest { + + @Test + void reportsContextClosedRaceCount() { + WafMetricCollector collector = WafMetricCollector.get(); + + collector.wafContextClosedRace(); + collector.wafContextClosedRace(); + collector.prepareMetrics(); + + Collection metrics = collector.drain(); + Optional raceMetric = + metrics.stream().filter(m -> "waf.context_closed_race".equals(m.metricName)).findFirst(); + + assertTrue(raceMetric.isPresent(), "expected waf.context_closed_race to be reported"); + assertEquals("count", raceMetric.get().type); + assertEquals("appsec", raceMetric.get().namespace); + assertEquals(2L, raceMetric.get().value.longValue()); + } +} From 326621000991931d8145238795891d2459dea459 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Wed, 22 Jul 2026 12:07:21 +0200 Subject: [PATCH 3/4] Address review comments: fallback close path and RASP eval double-count - AppSecRequestContext.close() now always calls closeWafContext(), even when no WafContext was ever created for the request. The previous wafContext != null guard meant the fallback path (missed request-end event) never set wafContextClosed for requests that hadn't run the WAF yet, leaving the same TOCTOU window this PR fixes. - WAFModule's null-skip branch no longer also reports rasp.rule.skipped: raspRuleEval() is already counted before doRunWaf() runs, so counting raspRuleSkipped() too double-counts the same callback as both evaluated and skipped, unlike the isWafContextClosed() fast path which only attempted eval when this branch is *not* taken. --- .../java/com/datadog/appsec/ddwaf/WAFModule.java | 6 ++---- .../appsec/gateway/AppSecRequestContext.java | 4 +++- .../ddwaf/WAFModuleContextClosedRaceTest.java | 15 +++++++++++---- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java index dc811d905b4..7afc65eae14 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java @@ -314,12 +314,10 @@ public void onDataAvailable( resultWithData = doRunWaf(reqCtx, newData, ctxAndAddr, gwCtx); if (resultWithData == null) { // WAF context closed concurrently between the fast-path check and context creation; skip - // (APPSEC-69085). + // (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(); - if (gwCtx.isRasp) { - WafMetricCollector.get().raspRuleSkipped(gwCtx.raspRuleType); - } return; } } catch (TimeoutWafException tpe) { diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index c4679fd84ff..81f158cee2b 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java @@ -731,8 +731,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(); diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java index f069f975151..86bba656a08 100644 --- a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java @@ -22,6 +22,7 @@ import datadog.trace.api.telemetry.RuleType; import datadog.trace.api.telemetry.WafMetricCollector; import java.io.InputStream; +import java.util.Collection; import java.util.Collections; import java.util.Map; import okio.Okio; @@ -102,7 +103,7 @@ void skipsAndIncrementsCounterWhenContextClosedConcurrently() { } @Test - void skipsRaspRuleWhenContextClosedConcurrently() { + void countsRaspEvalNotSkippedWhenContextClosedConcurrently() { AppSecRequestContext reqCtx = mock(AppSecRequestContext.class); when(reqCtx.isWafContextClosed()).thenReturn(false); when(reqCtx.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())).thenReturn(null); @@ -115,9 +116,15 @@ void skipsRaspRuleWhenContextClosedConcurrently() { assertFalse(flow.isBlocking()); WafMetricCollector.get().prepareMetrics(); + // The eval attempt was already counted before the race was detected; rasp.rule.skipped is + // reserved for calls that never attempted eval (e.g. the isWafContextClosed() fast path), so + // it must not also be reported here - otherwise the same callback is double-counted. + Collection metrics = WafMetricCollector.get().drain(); + boolean sawRaspEval = metrics.stream().anyMatch(m -> "rasp.rule.eval".equals(m.metricName)); boolean sawRaspSkipped = - WafMetricCollector.get().drain().stream() - .anyMatch(m -> "rasp.rule.skipped".equals(m.metricName)); - assertTrue(sawRaspSkipped, "expected rasp.rule.skipped to be reported"); + metrics.stream().anyMatch(m -> "rasp.rule.skipped".equals(m.metricName)); + assertTrue(sawRaspEval, "expected rasp.rule.eval to be reported"); + assertFalse( + sawRaspSkipped, "rasp.rule.skipped must not double-count an already-evaluated call"); } } From f68de91d4645a4068a686703f5e013653b8e5639 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Tue, 28 Jul 2026 11:01:18 +0200 Subject: [PATCH 4/4] Skip fast-path in closeWafContext and fix telemetry for closed-context race Add a fast-path return in closeWafContext() for redundant close() calls, and stop misclassifying a WafContext closed concurrently between getOrCreateWafContext() and run() as a real WAF error - it's the same benign race already tracked via wafContextClosedRace(). --- .../com/datadog/appsec/ddwaf/WAFModule.java | 10 ++++- .../appsec/gateway/AppSecRequestContext.java | 5 +++ .../ddwaf/WAFModuleContextClosedRaceTest.java | 43 +++++++++++++++++++ ...ppSecRequestContextWafContextRaceTest.java | 26 +++++++++++ 4 files changed, 82 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java index 7afc65eae14..d8a984cc2c0 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java @@ -330,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); diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index 81f158cee2b..3a3a65002fe 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java @@ -388,6 +388,11 @@ public WafContext getOrCreateWafContext( } public void closeWafContext() { + 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). diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java index 86bba656a08..458781f339c 100644 --- a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java @@ -16,6 +16,8 @@ import com.datadog.appsec.gateway.GatewayContext; import com.datadog.ddwaf.Waf; import com.datadog.ddwaf.WafBuilder; +import com.datadog.ddwaf.WafContext; +import com.datadog.ddwaf.WafHandle; import com.squareup.moshi.JsonAdapter; import com.squareup.moshi.Moshi; import com.squareup.moshi.Types; @@ -127,4 +129,45 @@ void countsRaspEvalNotSkippedWhenContextClosedConcurrently() { assertFalse( sawRaspSkipped, "rasp.rule.skipped must not double-count an already-evaluated call"); } + + /** + * Covers the structurally distinct race pointed out in review: {@code getOrCreateWafContext()} + * can return a non-null {@link WafContext} that is closed by another thread between the fetch and + * the actual {@code run()} call. In practice {@code closeWafContext()} flips {@code + * wafContextClosed} atomically with closing the native context, so by the time {@code run()} + * observes the closed context, {@code isWafContextClosed()} is already {@code true} - this + * mirrors that ordering rather than closing the context independently of the flag. {@code run()} + * then throws because the native context is no longer online, which {@code WAFModule} rewraps as + * {@code UnclassifiedWafException}. That case must be treated the same as the null-return race: + * no error log, no error-code metric, only {@code wafContextClosedRace()}. + */ + @Test + void skipsAndIncrementsCounterWhenContextClosedDuringRun() throws Exception { + WafHandle wafHandle = wafBuilder.buildWafHandleInstance(); + WafContext closedContext = new WafContext(wafHandle); + closedContext.close(); + + AppSecRequestContext reqCtx = mock(AppSecRequestContext.class); + // First call is the fast-path check before doRunWaf() runs (must be false to reach run()); + // subsequent calls mimic closeWafContext() flipping the flag concurrently while run() fails. + when(reqCtx.isWafContextClosed()).thenReturn(false, true); + when(reqCtx.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())).thenReturn(closedContext); + + ChangeableFlow flow = new ChangeableFlow(); + GatewayContext gwCtx = new GatewayContext(false); + + dataListener.onDataAvailable( + flow, reqCtx, MapDataBundle.ofDelegate(Collections.emptyMap()), gwCtx); + + assertFalse(flow.isBlocking()); + WafMetricCollector.get().prepareMetrics(); + Collection metrics = WafMetricCollector.get().drain(); + boolean sawContextClosedRace = + metrics.stream().anyMatch(m -> "waf.context_closed_race".equals(m.metricName)); + boolean sawErrorCode = metrics.stream().anyMatch(m -> "waf.error".equals(m.metricName)); + assertTrue(sawContextClosedRace, "expected waf.context_closed_race to be reported"); + assertFalse( + sawErrorCode, + "a benign context-closed race must not be double-counted as a real WAF error"); + } } diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java index 347e3986c20..f9c3e78ddf2 100644 --- a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java @@ -118,6 +118,32 @@ void closeBeforeFirstUsePreventsLaterCreation() { assertTrue(ctx.isWafContextClosed()); } + @Test + void fallbackCloseDoesNotClosePendingApiSecurityWafContext() { + AppSecRequestContext ctx = new AppSecRequestContext(); + WafContext created = ctx.getOrCreateWafContext(wafHandle, false, false); + assertNotNull(created); + + // API Security has sampled this request and asked to keep the context open for later + // schema-extraction post-processing (see ApiSecuritySamplerImpl#preSampleRequest). + ctx.setKeepOpenForApiSecurityPostProcessing(true); + + // The generic fallback close (CoreTracer#onRootSpanPublished) must not tear down the WAF + // context while API Security post-processing is still pending, or schema extraction would + // run against an already-closed context. + ctx.close(); + assertFalse(ctx.isWafContextClosed()); + assertTrue(created.isOnline()); + + // AppSecSpanPostProcessor#process's finally block: post-processing is done, now really close. + ctx.setKeepOpenForApiSecurityPostProcessing(false); + ctx.closeWafContext(); + ctx.close(); + + assertTrue(ctx.isWafContextClosed()); + assertFalse(created.isOnline()); + } + /** * Drives {@code closeWafContext()} and {@code getOrCreateWafContext()} concurrently via a barrier * across many iterations. The invariant: whatever {@code getOrCreateWafContext} returns after the