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..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 @@ -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(); 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..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 @@ -353,39 +353,55 @@ 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() { - 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(); 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..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 @@ -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() @@ -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) 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() 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..458781f339c --- /dev/null +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleContextClosedRaceTest.java @@ -0,0 +1,173 @@ +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.datadog.ddwaf.WafContext; +import com.datadog.ddwaf.WafHandle; +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.Collection; +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 countsRaspEvalNotSkippedWhenContextClosedConcurrently() { + 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(); + // 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 = + 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"); + } + + /** + * 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 new file mode 100644 index 00000000000..f9c3e78ddf2 --- /dev/null +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/gateway/AppSecRequestContextWafContextRaceTest.java @@ -0,0 +1,199 @@ +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()); + } + + @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()); + } + + @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 + * 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)"); + } + } +} 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()); + } +}