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 16affc8ac59..02ea658406b 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 @@ -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> 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> 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). 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 af817f7eca7..3cfea27eb66 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 @@ -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 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 ec4904c4882..6dd7c4ec92a 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 @@ -209,6 +209,7 @@ class WAFModuleSpecification extends DDSpecification { rba.blockingContentType == BlockingContentType.JSON }) + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -242,6 +243,7 @@ class WAFModuleSpecification extends DDSpecification { rba.statusCode == 403 && rba.blockingContentType == BlockingContentType.AUTO }) + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) } @@ -280,6 +282,7 @@ class WAFModuleSpecification extends DDSpecification { rba.statusCode == 403 && rba.blockingContentType == BlockingContentType.AUTO }) + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) } 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -304,6 +307,7 @@ class WAFModuleSpecification extends DDSpecification { rba.statusCode == 403 && rba.blockingContentType == BlockingContentType.AUTO }) + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) } 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -356,6 +360,7 @@ class WAFModuleSpecification extends DDSpecification { rba.statusCode == 403 && rba.blockingContentType == BlockingContentType.AUTO }) + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) } 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -376,6 +381,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() 1 * ctx.isWafContextClosed() >> false @@ -428,6 +434,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -450,6 +457,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() 1 * ctx.isWafContextClosed() >> false @@ -505,6 +513,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * tracer.activeSpan() 1 * ctx.reportEvents({ it.size() == 1 }) @@ -572,6 +581,7 @@ class WAFModuleSpecification extends DDSpecification { rba.statusCode == 401 && rba.blockingContentType == BlockingContentType.AUTO }) + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -595,6 +605,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -665,6 +676,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -693,6 +705,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(false, false) 1 * ctx.getOrCreateWafContext(_, false, false) 2 * ctx.getWafMetrics() >> null 1 * ctx.isWafContextClosed() >> false @@ -719,6 +732,7 @@ class WAFModuleSpecification extends DDSpecification { pp.processTraceSegment(segment, ctx, []) then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 1 * ctx.closeWafContext() 3 * ctx.getWafMetrics() >> { @@ -748,6 +762,7 @@ class WAFModuleSpecification extends DDSpecification { dataListener.onDataAvailable(flow, ctx, ATTACK_BUNDLE, gwCtx) then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() >> metrics 1 * ctx.reportEvents(*_) @@ -773,6 +788,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, _) 1 * ctx.getOrCreateWafContext(_, true, _) 1 * ctx.reportEvents(_ as Collection) >> { event = it[0].iterator().next() } 1 * ctx.reportStackTrace(_ as StackTraceEvent) >> { stackTrace = it[0] } @@ -808,6 +824,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, _) 1 * ctx.getOrCreateWafContext(_, true, _) ctx.reportEvents(_ as Collection) >> { event = it[0].iterator().next() } @@ -849,6 +866,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, _) 1 * ctx.getOrCreateWafContext(_, true, _) ctx.reportEvents(_ as Collection) >> { event = it[0].iterator().next() } @@ -868,6 +886,7 @@ class WAFModuleSpecification extends DDSpecification { dataListener.onDataAvailable(flow, ctx, db, gwCtx) then: + 1 * ctx.getWafContextIfReady(true, _) 1 * ctx.getOrCreateWafContext(_, true, _) !flow.blocking } @@ -895,6 +914,7 @@ class WAFModuleSpecification extends DDSpecification { dataListener.onDataAvailable(flow, ctx, db, gwCtx) then: + 1 * ctx.getWafContextIfReady(true, _) 1 * ctx.getOrCreateWafContext(_, true, _) !flow.blocking } @@ -908,6 +928,7 @@ class WAFModuleSpecification extends DDSpecification { dataListener.onDataAvailable(flow, ctx, db, gwCtx) then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() 1 * ctx.setWafErrors() @@ -934,6 +955,7 @@ class WAFModuleSpecification extends DDSpecification { then: assert !flow.blocking 1 * ctx.isWafContextClosed() + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() 1 * ctx.increaseWafTimeouts() @@ -970,6 +992,7 @@ class WAFModuleSpecification extends DDSpecification { then: assert !flow.blocking + 1 * ctx.getWafContextIfReady(true, true) 1 * ctx.getOrCreateWafContext(_, true, true) 1 * ctx.isWafContextClosed() 1 * ctx.getRaspMetrics() @@ -1007,6 +1030,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 1 * ctx.reportEvents(_ as Collection) 1 * ctx.isWafContextClosed() @@ -1045,6 +1069,7 @@ class WAFModuleSpecification extends DDSpecification { then: 1 * wafMetricCollector.wafUpdates(_, true) 1 * reconf.reloadSubscriptions() + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -1117,6 +1142,7 @@ class WAFModuleSpecification extends DDSpecification { then: 'no match; rule is disabled' 1 * wafMetricCollector.wafUpdates(_, true) 1 * reconf.reloadSubscriptions() + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() 1 * ctx.isWafContextClosed() >> false @@ -1135,6 +1161,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: 'no match; data was cleared (though rule is no longer disabled)' + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 1 * ctx.isWafContextClosed() >> false 1 * wafMetricCollector.wafUpdates(_, true) @@ -1156,6 +1183,7 @@ class WAFModuleSpecification extends DDSpecification { then: 'now we have match' 1 * wafMetricCollector.wafUpdates(_, true) 1 * reconf.reloadSubscriptions() + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) @@ -1181,6 +1209,7 @@ class WAFModuleSpecification extends DDSpecification { then: 'nothing again; we disabled the rule' 1 * wafMetricCollector.wafUpdates(_, true) 1 * reconf.reloadSubscriptions() + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() 1 * ctx.isWafContextClosed() >> false @@ -1207,6 +1236,7 @@ class WAFModuleSpecification extends DDSpecification { 1 * wafMetricCollector.wafUpdates(null, true) 1 * reconf.reloadSubscriptions() // no attack + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { WafHandle wafHandle = it[0] as WafHandle wafContext = new WafContext(wafHandle) @@ -1230,6 +1260,7 @@ class WAFModuleSpecification extends DDSpecification { 1 * wafMetricCollector.wafUpdates(_, true) 1 * reconf.reloadSubscriptions() // no attack + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { WafHandle wafHandle = it[0] as WafHandle wafContext = new WafContext(wafHandle) @@ -1254,6 +1285,7 @@ class WAFModuleSpecification extends DDSpecification { 1 * wafMetricCollector.wafUpdates(_, true) 1 * reconf.reloadSubscriptions() // attack found + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { WafHandle wafHandle = it[0] as WafHandle wafContext = new WafContext(wafHandle) @@ -1285,6 +1317,7 @@ class WAFModuleSpecification extends DDSpecification { 1 * wafMetricCollector.wafUpdates(_, true) 1 * reconf.reloadSubscriptions() // no attack + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { WafHandle wafHandle = it[0] as WafHandle wafContext = new WafContext(wafHandle) @@ -1403,6 +1436,7 @@ class WAFModuleSpecification extends DDSpecification { dataListener.onDataAvailable(flow, ctx, transientBundle, gwCtx) then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) >> { @@ -1420,7 +1454,9 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: - 1 * ctx.getOrCreateWafContext(_, true, false) + // the context created by the first call is reused without going through creation + 1 * ctx.getWafContextIfReady(true, false) + 0 * ctx.getOrCreateWafContext(_, _, _) 1 * flow.setAction({ it.blocking }) 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) >> { @@ -1499,6 +1535,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() 1 * ctx.isThrottled(null) @@ -1532,6 +1569,7 @@ class WAFModuleSpecification extends DDSpecification { }) 1 * flow.isBlocking() 1 * ctx.isWafContextClosed() >> false + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_ as WafHandle, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() 1 * ctx.isThrottled(null) @@ -1760,6 +1798,7 @@ class WAFModuleSpecification extends DDSpecification { then: (1..2) * ctx.isWafContextClosed() >> false // if UnclassifiedWafException it's called twice + 1 * ctx.getWafContextIfReady(true, true) 1 * ctx.getOrCreateWafContext(_, true, true) >> wafContext 1 * wafMetricCollector.raspRuleEval(RuleType.SQL_INJECTION) 1 * wafContext.run(_, _, _) >> { throw createWafException(wafErrorCode as WafErrorCode) } @@ -1789,6 +1828,7 @@ class WAFModuleSpecification extends DDSpecification { then: (1..2) * ctx.isWafContextClosed() >> false // if UnclassifiedWafException it's called twice + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> wafContext 1 * wafContext.run(_, _, _) >> { throw createWafException(wafErrorCode as WafErrorCode) } 1 * wafMetricCollector.wafInit(Waf.LIB_VERSION, _, true) @@ -2001,6 +2041,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -2019,6 +2060,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -2039,6 +2081,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -2056,6 +2099,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) >> { wafContext = new WafContext(it[0]) } 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -2159,6 +2203,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -2224,6 +2269,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false @@ -2436,6 +2482,7 @@ class WAFModuleSpecification extends DDSpecification { ctx.closeWafContext() then: + 1 * ctx.getWafContextIfReady(true, false) 1 * ctx.getOrCreateWafContext(_, true, false) 2 * ctx.getWafMetrics() >> metrics 1 * ctx.isWafContextClosed() >> false diff --git a/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleHandleReloadRaceTest.java b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleHandleReloadRaceTest.java new file mode 100644 index 00000000000..c3c6f6bf2fc --- /dev/null +++ b/dd-java-agent/appsec/src/test/java/com/datadog/appsec/ddwaf/WAFModuleHandleReloadRaceTest.java @@ -0,0 +1,370 @@ +package com.datadog.appsec.ddwaf; + +import static java.util.Collections.emptyMap; +import static org.junit.jupiter.api.Assertions.assertEquals; +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 static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +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.Moshi; +import com.squareup.moshi.Types; +import datadog.trace.api.telemetry.RuleType; +import java.io.InputStream; +import java.lang.reflect.Field; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.StampedLock; +import okio.Okio; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests handle replacement against context creation without accessing retired native pointers. */ +class WAFModuleHandleReloadRaceTest { + private WafHandle initialHandle; + private WafBuilder builder; + private WAFModule module; + private final AppSecRequestContext request = mock(AppSecRequestContext.class); + private final WafContext context = mock(WafContext.class); + private AppSecModuleConfigurer.SubconfigListener listener; + private DataListener callback; + + @BeforeEach + void setup() throws Exception { + assertTrue(WafInitialization.ONLINE); + builder = new WafBuilder(); + try (InputStream stream = getClass().getResourceAsStream("/test_multi_config.json")) { + Map config = + new Moshi.Builder() + .build() + .>adapter( + Types.newParameterizedType(Map.class, String.class, Object.class)) + .fromJson(Okio.buffer(Okio.source(stream))); + builder.addOrUpdateConfig("test", config); + } + when(context.run(anyMap(), any(), any())).thenReturn(Waf.ResultWithData.OK_NULL); + module = new WAFModule(); + module.setWafBuilder(builder); + module.config( + new AppSecModuleConfigurer() { + @Override + public void addSubConfigListener( + String key, AppSecModuleConfigurer.SubconfigListener value) { + listener = value; + } + + @Override + public void addTraceSegmentPostProcessor(TraceSegmentPostProcessor interceptor) {} + }); + callback = module.getDataSubscriptions().iterator().next(); + initialHandle = currentHandle(); + } + + private Object currentSnapshotField(String name) throws Exception { + Field referenceField = WAFModule.class.getDeclaredField("ctxAndAddresses"); + referenceField.setAccessible(true); + Object snapshot = ((AtomicReference) referenceField.get(module)).get(); + Field field = snapshot.getClass().getDeclaredField(name); + field.setAccessible(true); + return field.get(snapshot); + } + + private WafHandle currentHandle() throws Exception { + return (WafHandle) currentSnapshotField("ctx"); + } + + @AfterEach + void cleanup() throws Exception { + if (module != null) { + currentHandle().close(); + } + if (builder != null) { + builder.close(); + } + } + + private void reload() throws Exception { + listener.onNewSubconfig(null, AppSecModuleConfigurer.Reconfiguration.NOOP); + } + + private void evaluate() { + evaluate(request, new GatewayContext(false)); + } + + private void evaluate(AppSecRequestContext requestContext, GatewayContext gatewayContext) { + callback.onDataAvailable( + new ChangeableFlow(), requestContext, MapDataBundle.ofDelegate(emptyMap()), gatewayContext); + } + + @Test + void usesReplacementWhenReloadRetiresTheCapturedHandle() throws Exception { + when(request.isWafContextClosed()) + .thenAnswer( + invocation -> { + reload(); + return false; + }); + when(request.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())).thenReturn(context); + + evaluate(); + + verify(request) + .getOrCreateWafContext( + argThat(handle -> handle != initialHandle && handle.isOnline()), + anyBoolean(), + anyBoolean()); + verify(request, never()).getOrCreateWafContext(eq(initialHandle), anyBoolean(), anyBoolean()); + verify(context).run(anyMap(), any(), any()); + } + + @Test + void waitsForContextCreationBeforeClosingTheHandle() throws Exception { + CountDownLatch creating = new CountDownLatch(1); + CountDownLatch finishCreation = new CountDownLatch(1); + when(request.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())) + .thenAnswer( + invocation -> { + creating.countDown(); + assertTrue(finishCreation.await(10, TimeUnit.SECONDS)); + return context; + }); + FutureTask evaluation = + new FutureTask<>( + () -> { + evaluate(); + return null; + }); + FutureTask update = + new FutureTask<>( + () -> { + reload(); + return null; + }); + Thread evaluationThread = new Thread(evaluation, "waf-context-creation-test"); + Thread updateThread = new Thread(update, "waf-handle-reload-test"); + evaluationThread.start(); + try { + assertTrue(creating.await(10, TimeUnit.SECONDS)); + updateThread.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (!update.isDone() + && !blockedInHandleReplacement(updateThread) + && System.nanoTime() < deadline) { + Thread.yield(); + } + assertFalse(update.isDone(), "reload must wait for context creation"); + assertTrue(blockedInHandleReplacement(updateThread)); + assertTrue(initialHandle.isOnline(), "handle must remain online during context creation"); + } finally { + finishCreation.countDown(); + evaluationThread.join(10000); + updateThread.join(10000); + } + evaluation.get(10, TimeUnit.SECONDS); + update.get(10, TimeUnit.SECONDS); + assertFalse(initialHandle.isOnline()); + verify(context).run(anyMap(), any(), any()); + } + + private static boolean blockedInHandleReplacement(Thread thread) { + if (thread.getState() != Thread.State.WAITING) { + return false; + } + for (StackTraceElement frame : thread.getStackTrace()) { + if (frame.getClassName().equals(WAFModule.class.getName()) + && frame.getMethodName().equals("initOrUpdateWafHandle")) { + return true; + } + } + return false; + } + + @Test + void retriesWithoutWaitingForRetiredHandleDestruction() throws Exception { + CountDownLatch captured = new CountDownLatch(1); + CountDownLatch replaced = new CountDownLatch(1); + when(request.isWafContextClosed()) + .thenAnswer( + invocation -> { + captured.countDown(); + assertTrue(replaced.await(10, TimeUnit.SECONDS)); + return false; + }); + when(request.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())).thenReturn(context); + StampedLock oldLock = (StampedLock) currentSnapshotField("handleLock"); + FutureTask evaluation = + new FutureTask<>( + () -> { + evaluate(); + return null; + }); + FutureTask update = + new FutureTask<>( + () -> { + reload(); + return null; + }); + Thread evaluationThread = new Thread(evaluation, "waf-stale-snapshot-test"); + Thread updateThread = new Thread(update, "waf-retired-handle-test"); + long stamp = oldLock.writeLock(); + try { + evaluationThread.start(); + assertTrue(captured.await(10, TimeUnit.SECONDS)); + updateThread.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (currentHandle() == initialHandle && !update.isDone() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertTrue(currentHandle() != initialHandle, "replacement must be published"); + replaced.countDown(); + evaluation.get(10, TimeUnit.SECONDS); + assertFalse(update.isDone(), "old-handle destruction is still blocked"); + verify(request) + .getOrCreateWafContext( + argThat(handle -> handle != initialHandle && handle.isOnline()), + anyBoolean(), + anyBoolean()); + } finally { + replaced.countDown(); + oldLock.unlockWrite(stamp); + evaluationThread.join(10000); + updateThread.join(10000); + } + update.get(10, TimeUnit.SECONDS); + assertFalse(initialHandle.isOnline()); + } + + @Test + void usesAReadyContextWithoutTakingTheHandleLock() throws Exception { + when(request.getWafContextIfReady(anyBoolean(), anyBoolean())).thenReturn(context); + StampedLock handleLock = (StampedLock) currentSnapshotField("handleLock"); + FutureTask evaluation = + new FutureTask<>( + () -> { + evaluate(); + return null; + }); + long stamp = handleLock.writeLock(); + try { + Thread thread = new Thread(evaluation, "waf-ready-context-test"); + thread.start(); + evaluation.get(10, TimeUnit.SECONDS); + } finally { + handleLock.unlockWrite(stamp); + } + + verify(request, never()).getOrCreateWafContext(any(), anyBoolean(), anyBoolean()); + verify(context).run(anyMap(), any(), any()); + } + + @Test + void doesNotCreateAContextForARequestClosedAfterTheClosedCheck() { + AppSecRequestContext requestContext = spy(new AppSecRequestContext()); + evaluate(requestContext, new GatewayContext(false)); + assertNotNull(requestContext.getWafContextIfReady(true, false)); + // close between onDataAvailable's isWafContextClosed() check and the ready-context read + doAnswer( + invocation -> { + Object open = invocation.callRealMethod(); + requestContext.closeWafContext(); + return open; + }) + .when(requestContext) + .isWafContextClosed(); + + evaluate(requestContext, new GatewayContext(false)); + + verify(requestContext, times(2)).getOrCreateWafContext(any(), anyBoolean(), anyBoolean()); + assertTrue(requestContext.isWafContextClosed()); + assertNull(requestContext.getWafContextIfReady(false, false)); + verify(requestContext, never()).setWafErrors(); + } + + @Test + void skipsAReadyContextClosedBeforeItRuns() throws Exception { + WafContext closed = new WafContext(currentHandle()); + closed.close(); + when(request.getWafContextIfReady(anyBoolean(), anyBoolean())).thenReturn(closed); + // open at onDataAvailable's fast check, closed by the time the run fails + when(request.isWafContextClosed()).thenReturn(false, true); + + evaluate(); + + verify(request, never()).getOrCreateWafContext(any(), anyBoolean(), anyBoolean()); + verify(request, never()).setWafErrors(); + } + + @Test + void firstRaspCallCreatesRaspMetricsThenUsesTheReadyContext() { + AppSecRequestContext requestContext = spy(new AppSecRequestContext()); + GatewayContext rasp = new GatewayContext(false, RuleType.SQL_INJECTION); + try { + evaluate(requestContext, new GatewayContext(false)); + WafContext created = requestContext.getWafContextIfReady(true, false); + assertNotNull(created); + assertNull(requestContext.getRaspMetrics()); + + evaluate(requestContext, rasp); + evaluate(requestContext, rasp); + + verify(requestContext, times(1)).getOrCreateWafContext(any(), eq(true), eq(true)); + assertNotNull(requestContext.getRaspMetrics()); + assertEquals(2, requestContext.getRaspMetricsCounter().get()); + assertSame(created, requestContext.getWafContextIfReady(true, true)); + } finally { + requestContext.closeWafContext(); + } + } + + @Test + void allowsReloadWhileAnExistingContextEvaluates() throws Exception { + when(request.getOrCreateWafContext(any(), anyBoolean(), anyBoolean())).thenReturn(context); + when(context.run(anyMap(), any(), any())) + .thenAnswer( + invocation -> { + FutureTask update = + new FutureTask<>( + () -> { + reload(); + return null; + }); + Thread thread = new Thread(update, "waf-reload-during-evaluation-test"); + thread.start(); + update.get(10, TimeUnit.SECONDS); + assertFalse(initialHandle.isOnline()); + return Waf.ResultWithData.OK_NULL; + }); + + evaluate(); + + verify(context).run(anyMap(), any(), any()); + } +} 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 f9c3e78ddf2..d26e574ae98 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,52 @@ void closeBeforeFirstUsePreventsLaterCreation() { assertTrue(ctx.isWafContextClosed()); } + @Test + void readyContextIsAbsentBeforeFirstUse() { + AppSecRequestContext ctx = new AppSecRequestContext(); + + assertNull(ctx.getWafContextIfReady(false, false)); + } + + @Test + void readyContextIsTheCreatedContext() { + AppSecRequestContext ctx = new AppSecRequestContext(); + try { + WafContext created = ctx.getOrCreateWafContext(wafHandle, true, false); + + assertSame(created, ctx.getWafContextIfReady(true, false)); + assertSame(created, ctx.getWafContextIfReady(false, true)); + } finally { + ctx.closeWafContext(); + } + } + + @Test + void readyContextRequiresMetricsTheCallNeeds() { + AppSecRequestContext ctx = new AppSecRequestContext(); + try { + WafContext created = ctx.getOrCreateWafContext(wafHandle, true, false); + + // The first RASP call must go through getOrCreateWafContext so RASP metrics get created. + assertNull(ctx.getWafContextIfReady(true, true)); + assertSame(created, ctx.getOrCreateWafContext(wafHandle, true, true)); + assertNotNull(ctx.getRaspMetrics()); + assertSame(created, ctx.getWafContextIfReady(true, true)); + } finally { + ctx.closeWafContext(); + } + } + + @Test + void readyContextIsAbsentAfterClose() { + AppSecRequestContext ctx = new AppSecRequestContext(); + assertNotNull(ctx.getOrCreateWafContext(wafHandle, false, false)); + + ctx.closeWafContext(); + + assertNull(ctx.getWafContextIfReady(false, false)); + } + @Test void fallbackCloseDoesNotClosePendingApiSecurityWafContext() { AppSecRequestContext ctx = new AppSecRequestContext();