From fef2c64601b5f29b3125529cc71c5208d4d10c2e Mon Sep 17 00:00:00 2001 From: Valentin Zakharov Date: Mon, 10 Feb 2025 00:25:00 +0100 Subject: [PATCH 01/15] Added AppSec post-processing --- .../datadog/appsec/gateway/GatewayBridge.java | 7 ++ .../gateway/GatewayBridgeSpecification.groovy | 3 + .../trace/common/writer/DDAgentWriter.java | 5 +- .../trace/common/writer/DDIntakeWriter.java | 6 +- .../AppSecSpanPostProcessor.java | 40 ++++++ .../AppSecSpanPostProcessorTest.groovy | 119 ++++++++++++++++++ .../datadog/trace/api/gateway/Events.java | 13 ++ .../api/gateway/InstrumentationGateway.java | 14 +++ .../gateway/InstrumentationGatewayTest.java | 18 +++ 9 files changed, 223 insertions(+), 2 deletions(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java create mode 100644 dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index 54665d12ca9..a3e58fa32e8 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -5,6 +5,7 @@ import static com.datadog.appsec.gateway.AppSecRequestContext.DEFAULT_REQUEST_HEADERS_ALLOW_LIST; import static com.datadog.appsec.gateway.AppSecRequestContext.REQUEST_HEADERS_ALLOW_LIST; import static com.datadog.appsec.gateway.AppSecRequestContext.RESPONSE_HEADERS_ALLOW_LIST; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import com.datadog.appsec.AppSecSystem; import com.datadog.appsec.api.security.ApiSecurityRequestSampler; @@ -151,6 +152,7 @@ public void init() { subscriptionService.registerCallback(EVENTS.shellCmd(), this::onShellCmd); subscriptionService.registerCallback(EVENTS.user(), this::onUser); subscriptionService.registerCallback(EVENTS.loginEvent(), this::onLoginEvent); + subscriptionService.registerCallback(EVENTS.postProcessing(), this::onPostProcessing); if (additionalIGEvents.contains(EVENTS.requestPathParams())) { subscriptionService.registerCallback(EVENTS.requestPathParams(), this::onRequestPathParams); @@ -787,6 +789,10 @@ private void onRequestHeader(RequestContext ctx_, String name, String value) { } } + private void onPostProcessing(RequestContext ctx_) { + // Do AppSec post-processing + } + public void stop() { subscriptionService.reset(); } @@ -912,6 +918,7 @@ private Flow maybePublishRequestData(AppSecRequestContext ctx) { try { GatewayContext gwCtx = new GatewayContext(false); + activeSpan().getLocalRootSpan().setRequiresPostProcessing(true); return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx); } catch (ExpiredSubscriberInfoException e) { this.initialReqDataSubInfo = null; diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index c13156d461b..3ff70b5afd2 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -26,6 +26,7 @@ import datadog.trace.test.util.DDSpecification import java.util.function.BiConsumer import java.util.function.BiFunction +import java.util.function.Consumer import java.util.function.Function import java.util.function.Supplier @@ -103,6 +104,7 @@ class GatewayBridgeSpecification extends DDSpecification { BiFunction> shellCmdCB BiFunction> userCB TriFunction> loginEventCB + Consumer postProcessingCB void setup() { callInitAndCaptureCBs() @@ -443,6 +445,7 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * ig.registerCallback(EVENTS.shellCmd(), _) >> { shellCmdCB = it[1]; null } 1 * ig.registerCallback(EVENTS.user(), _) >> { userCB = it[1]; null } 1 * ig.registerCallback(EVENTS.loginEvent(), _) >> { loginEventCB = it[1]; null } + 1 * ig.registerCallback(EVENTS.postProcessing(), _) >> { postProcessingCB = it[1]; null } 0 * ig.registerCallback(_, _) bridge.init() diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java index 153bcca6565..f576a3fe953 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java @@ -14,6 +14,8 @@ import datadog.trace.common.writer.ddagent.DDAgentMapperDiscovery; import datadog.trace.common.writer.ddagent.Prioritization; import datadog.trace.core.monitor.HealthMetrics; +import datadog.trace.core.postprocessor.AppSecSpanPostProcessor; +import datadog.trace.core.postprocessor.SpanPostProcessor; import java.util.concurrent.TimeUnit; import okhttp3.HttpUrl; import okhttp3.OkHttpClient; @@ -153,6 +155,7 @@ public DDAgentWriter build() { final DDAgentMapperDiscovery mapperDiscovery = new DDAgentMapperDiscovery(featureDiscovery); final PayloadDispatcher dispatcher = new PayloadDispatcherImpl(mapperDiscovery, agentApi, healthMetrics, monitoring); + final SpanPostProcessor spanPostProcessor = new AppSecSpanPostProcessor(); final TraceProcessingWorker traceProcessingWorker = new TraceProcessingWorker( traceBufferSize, @@ -163,7 +166,7 @@ public DDAgentWriter build() { flushIntervalMilliseconds, TimeUnit.MILLISECONDS, singleSpanSampler, - null); + spanPostProcessor); return new DDAgentWriter( traceProcessingWorker, diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java index 614865a6b5e..fd26f4db7eb 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java @@ -9,6 +9,8 @@ import datadog.trace.common.writer.ddagent.Prioritization; import datadog.trace.common.writer.ddintake.DDIntakeMapperDiscovery; import datadog.trace.core.monitor.HealthMetrics; +import datadog.trace.core.postprocessor.AppSecSpanPostProcessor; +import datadog.trace.core.postprocessor.SpanPostProcessor; import java.util.EnumMap; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -113,6 +115,8 @@ public DDIntakeWriter build() { dispatcher = new CompositePayloadDispatcher(dispatchers); } + SpanPostProcessor spanPostProcessor = new AppSecSpanPostProcessor(); + final TraceProcessingWorker traceProcessingWorker = new TraceProcessingWorker( traceBufferSize, @@ -123,7 +127,7 @@ public DDIntakeWriter build() { flushIntervalMilliseconds, TimeUnit.MILLISECONDS, singleSpanSampler, - null); + spanPostProcessor); return new DDIntakeWriter( traceProcessingWorker, diff --git a/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java b/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java new file mode 100644 index 00000000000..ab4a4ea40ff --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java @@ -0,0 +1,40 @@ +package datadog.trace.core.postprocessor; + +import static datadog.trace.api.gateway.Events.EVENTS; + +import datadog.trace.api.gateway.CallbackProvider; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.core.DDSpan; +import java.util.function.BooleanSupplier; +import java.util.function.Consumer; + +public class AppSecSpanPostProcessor implements SpanPostProcessor { + + // For testing purpose + protected AgentTracer.TracerAPI tracer() { + return AgentTracer.get(); + } + + @Override + public boolean process(DDSpan span, BooleanSupplier timeoutCheck) { + CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC); + if (cbp == null) { + return false; + } + + RequestContext ctx = span.getRequestContext(); + if (ctx == null) { + return false; + } + + Consumer postProcessingCallback = cbp.getCallback(EVENTS.postProcessing()); + if (postProcessingCallback == null) { + return false; + } + + postProcessingCallback.accept(ctx); + return true; + } +} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy new file mode 100644 index 00000000000..17fc937d0e6 --- /dev/null +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy @@ -0,0 +1,119 @@ +package datadog.trace.core.postprocessor + +import datadog.trace.api.gateway.CallbackProvider +import datadog.trace.api.gateway.RequestContext +import datadog.trace.api.gateway.RequestContextSlot +import datadog.trace.bootstrap.instrumentation.api.AgentTracer +import datadog.trace.core.DDSpan +import datadog.trace.core.DDSpanContext +import datadog.trace.core.PendingTrace +import datadog.trace.test.util.DDSpecification + +import java.util.function.Consumer +import java.util.function.BooleanSupplier +import static datadog.trace.api.gateway.Events.EVENTS + +class AppSecSpanPostProcessorTest extends DDSpecification { + def "process returns false if span context is null"() { + given: + def processor = new AppSecSpanPostProcessor() + def span = Mock(DDSpan) + def timeoutCheck = Mock(BooleanSupplier) + (span.context()) >> null + + expect: + !processor.process(span, timeoutCheck) + } + + def "process returns false if callback provider is null"() { + given: + AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) + tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> null + def processor = new AppSecSpanPostProcessor() { + @Override + protected AgentTracer.TracerAPI tracer() { + return tracer + } + } + def span = Mock(DDSpan) { + context() >> Mock(DDSpanContext) + } + def timeoutCheck = Mock(BooleanSupplier) + + expect: + !processor.process(span, timeoutCheck) + } + + def "process returns false if request context is null"() { + given: + AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) + def cbp = Mock(CallbackProvider) + tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> cbp + def processor = new AppSecSpanPostProcessor() { + @Override + protected AgentTracer.TracerAPI tracer() { + return tracer + } + } + def span = Mock(DDSpan) { + context() >> Mock(DDSpanContext) + getRequestContext() >> null + } + def timeoutCheck = Mock(BooleanSupplier) + + expect: + !processor.process(span, timeoutCheck) + } + + def "process returns false if post-processing callback is null"() { + given: + AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) + def cbp = Mock(CallbackProvider) + tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> cbp + cbp.getCallback(EVENTS.postProcessing()) >> null + def processor = new AppSecSpanPostProcessor() { + @Override + protected AgentTracer.TracerAPI tracer() { + return tracer + } + } + def span = Mock(DDSpan) { + context() >> Mock(DDSpanContext) + getRequestContext() >> Mock(RequestContext) + } + def timeoutCheck = Mock(BooleanSupplier) + + expect: + !processor.process(span, timeoutCheck) + } + + def "process returns true when all components are properly configured"() { + given: + def callback = Mock(Consumer) + AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) + def cbp = Mock(CallbackProvider) + tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> cbp + cbp.getCallback(EVENTS.postProcessing()) >> callback + def processor = new AppSecSpanPostProcessor() { + @Override + protected AgentTracer.TracerAPI tracer() { + return tracer + } + } + def span = DDSpan.create("test", 0, Mock(DDSpanContext) { + isRequiresPostProcessing() >> true + getTraceCollector() >> Mock(PendingTrace) { + getCurrentTimeNano() >> 0 + } + getRequestContext() >> Mock(RequestContext) + }, []) + def timeoutCheck = Mock(BooleanSupplier) + + when: + boolean result = processor.process(span, timeoutCheck) + + then: + result + 1 * callback.accept(_) + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java index d840cad01c3..6f63ddf7049 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java @@ -9,6 +9,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import java.util.function.BiFunction; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; @@ -312,6 +313,18 @@ public EventType>> shellCmd() { return (EventType>>) SHELL_CMD; } + static final int POST_PROCESSING_ID = 26; + + @SuppressWarnings("rawtypes") + private static final EventType POST_PROCESSING = + new ET<>("trace.post.processing", POST_PROCESSING_ID); + + /** The span post-processing */ + @SuppressWarnings("unchecked") + public EventType> postProcessing() { + return (EventType>) POST_PROCESSING; + } + static final int MAX_EVENTS = nextId.get(); private static final class ET extends EventType { diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java b/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java index cd5219f1dfc..e770fcd1952 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java @@ -10,6 +10,7 @@ import static datadog.trace.api.gateway.Events.LOGIN_EVENT_ID; import static datadog.trace.api.gateway.Events.MAX_EVENTS; import static datadog.trace.api.gateway.Events.NETWORK_CONNECTION_ID; +import static datadog.trace.api.gateway.Events.POST_PROCESSING_ID; import static datadog.trace.api.gateway.Events.REQUEST_BODY_CONVERTED_ID; import static datadog.trace.api.gateway.Events.REQUEST_BODY_DONE_ID; import static datadog.trace.api.gateway.Events.REQUEST_BODY_START_ID; @@ -37,6 +38,7 @@ import java.util.concurrent.atomic.AtomicReferenceArray; import java.util.function.BiConsumer; import java.util.function.BiFunction; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; import org.slf4j.Logger; @@ -459,6 +461,18 @@ public Flow apply(RequestContext ctx, String[] arg) { } } }; + case POST_PROCESSING_ID: + return (C) + new Consumer() { + @Override + public void accept(RequestContext ctx) { + try { + ((Consumer) callback).accept(ctx); + } catch (Throwable t) { + log.warn("Callback for {} threw.", eventType, t); + } + } + }; default: log.warn("Unwrapped callback for {}", eventType); return callback; diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java index 8f2840d4308..b5ab581c479 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java @@ -13,6 +13,7 @@ import java.util.Collections; import java.util.function.BiConsumer; import java.util.function.BiFunction; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; import org.assertj.core.api.ThrowableAssert; @@ -223,6 +224,8 @@ public void testNormalCalls() { cbp.getCallback(events.execCmd()).apply(null, null); ss.registerCallback(events.shellCmd(), callback); cbp.getCallback(events.shellCmd()).apply(null, null); + ss.registerCallback(events.postProcessing(), callback); + cbp.getCallback(events.postProcessing()).accept(null); assertThat(callback.count).isEqualTo(Events.MAX_EVENTS); } @@ -293,6 +296,8 @@ public void testThrowableBlocking() { cbp.getCallback(events.execCmd()).apply(null, null); ss.registerCallback(events.shellCmd(), throwback); cbp.getCallback(events.shellCmd()).apply(null, null); + ss.registerCallback(events.postProcessing(), throwback); + cbp.getCallback(events.postProcessing()).accept(null); assertThat(throwback.count).isEqualTo(Events.MAX_EVENTS); } @@ -438,6 +443,7 @@ public void mergeFlowReturnsLatestNonNullResult() { private static class Callback implements Supplier>, + Consumer, BiConsumer, TriConsumer, BiFunction>, @@ -470,6 +476,11 @@ public Flow get() { return new Flow.ResultFlow<>((D) ctxt); } + @Override + public void accept(RequestContext requestContext) { + count++; + } + @Override public void accept(RequestContext requestContext, T s, T s2) { count++; @@ -510,6 +521,7 @@ public Flow apply(RequestContext requestContext, T t, T t2) { private static class Throwback implements Supplier>, + Consumer, BiConsumer, TriConsumer, BiFunction>, @@ -535,6 +547,12 @@ public Flow get() { throw new IllegalArgumentException(); } + @Override + public void accept(RequestContext requestContext) { + count++; + throw new IllegalArgumentException(); + } + @Override public void accept(RequestContext requestContext, T s, T s2) { count++; From 21a297588de8c69cfe4bc9d0ac8839b42a5ecb50 Mon Sep 17 00:00:00 2001 From: Valentin Zakharov Date: Mon, 10 Feb 2025 22:24:02 +0100 Subject: [PATCH 02/15] Api Security schema computation moved to post-processing stage in serialization thread --- .../appsec/gateway/AppSecRequestContext.java | 23 ++++--------------- .../datadog/appsec/gateway/GatewayBridge.java | 17 ++++++++------ .../AppSecRequestContextSpecification.groovy | 15 +++++------- .../gateway/GatewayBridgeSpecification.groovy | 2 +- .../java/datadog/trace/core/CoreTracer.java | 4 +++- 5 files changed, 24 insertions(+), 37 deletions(-) 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 34536773602..7afe609991e 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 @@ -10,8 +10,6 @@ import datadog.trace.api.Config; import datadog.trace.api.http.StoredBodySupplier; import datadog.trace.api.internal.TraceSegment; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.util.stacktrace.StackTraceEvent; import io.sqreen.powerwaf.Additive; import io.sqreen.powerwaf.PowerwafContext; @@ -569,15 +567,6 @@ public String getSessionId() { @Override public void close() { - final AgentSpan span = AgentTracer.activeSpan(); - close(span != null && span.isRequiresPostProcessing()); - } - - /* end interface for GatewayBridge */ - - /* Should be accessible from the modules */ - - public void close(boolean requiresPostProcessing) { if (!requestEndCalled) { log.debug(SEND_TELEMETRY, "Request end event was not called before close"); } @@ -586,18 +575,14 @@ public void close(boolean requiresPostProcessing) { SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)"); closeAdditive(); } - derivatives = null; - - // check if we might need to further post process data related to the span in order to not free - // related data - if (requiresPostProcessing) { - return; - } - collectedCookies = null; requestHeaders.clear(); responseHeaders.clear(); persistentData.clear(); + if (derivatives != null) { + derivatives.clear(); + derivatives = null; + } } /** @return the portion of the body read so far, if any */ diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index a3e58fa32e8..d8592343a62 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -663,11 +663,6 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { } ctx.setRequestEndCalled(); - maybeExtractSchemas(ctx); - - // WAF call - ctx.closeAdditive(); - TraceSegment traceSeg = ctx_.getTraceSegment(); // AppSec report metric and events for web span only @@ -730,7 +725,9 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { } } - ctx.close(spanInfo.isRequiresPostProcessing()); + if (!spanInfo.isRequiresPostProcessing()) { + ctx.close(); + } return NoopFlow.INSTANCE; } @@ -790,7 +787,13 @@ private void onRequestHeader(RequestContext ctx_, String name, String value) { } private void onPostProcessing(RequestContext ctx_) { - // Do AppSec post-processing + AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); + if (ctx == null) { + return; + } + + maybeExtractSchemas(ctx); + ctx.close(); } public void stop() { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy index 92d22cb6113..a29d68d5715 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy @@ -251,7 +251,6 @@ class AppSecRequestContextSpecification extends DDSpecification { void 'test that internal data is cleared on close'() { setup: final ctx = new AppSecRequestContext() - final fullCleanup = !postProcessing when: ctx.requestHeaders.put('Accept', ['*']) @@ -260,19 +259,17 @@ class AppSecRequestContextSpecification extends DDSpecification { ctx.persistentData.put(KnownAddresses.REQUEST_METHOD, 'GET') ctx.derivatives = ['a': 'b'] ctx.additive = createAdditive() - ctx.close(postProcessing) + ctx.close() then: ctx.additive == null ctx.derivatives == null + ctx.additive == null - ctx.requestHeaders.isEmpty() == fullCleanup - ctx.responseHeaders.isEmpty() == fullCleanup - ctx.cookies.isEmpty() == fullCleanup - ctx.persistentData.isEmpty() == fullCleanup - - where: - postProcessing << [true, false] + ctx.requestHeaders.isEmpty() + ctx.responseHeaders.isEmpty() + ctx.cookies.isEmpty() + ctx.persistentData.isEmpty() } def "test increase and get WafTimeouts"() { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 3ff70b5afd2..50269370666 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -161,7 +161,7 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * spanInfo.getTags() >> ['http.client_ip': '1.1.1.1'] 1 * mockAppSecCtx.transferCollectedEvents() >> [event] 1 * mockAppSecCtx.peerAddress >> '2001::1' - 1 * mockAppSecCtx.close(false) + 1 * mockAppSecCtx.close() 1 * traceSegment.setTagTop("_dd.appsec.enabled", 1) 1 * traceSegment.setTagTop("_dd.runtime_family", "jvm") 1 * traceSegment.setTagTop('appsec.event', true) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 5d033f13c78..cfe1d2c79bb 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -277,7 +277,9 @@ void onRootSpanPublished(final AgentSpan root) { RequestContext requestContext = root.getRequestContext(); if (requestContext != null) { try { - requestContext.close(); + if (!root.isRequiresPostProcessing()) { + requestContext.close(); + } } catch (IOException e) { log.warn("Error closing request context data", e); } From 0c2cec690e3eca27d67198e7808df4b751242ad3 Mon Sep 17 00:00:00 2001 From: Valentin Zakharov Date: Tue, 11 Feb 2025 18:17:14 +0100 Subject: [PATCH 03/15] Implemented time-based sampling per endpoint --- .../appsec/api/security/ApiAccessTracker.java | 102 ++++++++++++++++++ .../security/ApiSecurityRequestSampler.java | 58 ++++------ .../appsec/gateway/AppSecRequestContext.java | 17 ++- .../datadog/appsec/gateway/GatewayBridge.java | 9 +- .../api/security/ApiAccessTrackerTest.groovy | 55 ++++++++++ .../ApiSecurityRequestSamplerTest.groovy | 58 ++++------ .../gateway/GatewayBridgeSpecification.groovy | 1 - 7 files changed, 217 insertions(+), 83 deletions(-) create mode 100644 dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java create mode 100644 dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java new file mode 100644 index 00000000000..32aaa761717 --- /dev/null +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java @@ -0,0 +1,102 @@ +package com.datadog.appsec.api.security; + +import java.util.Deque; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedDeque; + +/** + * The ApiAccessTracker class provides a mechanism to track API access events, managing them within + * a specified capacity limit. Each event is associated with a unique combination of route, method, + * and status code, which is used to generate a unique key for tracking access timestamps. + * + *

Usage: - When an API access event occurs, the `updateApiAccessIfExpired` method is called with + * the route, method, and status code of the API request. - If the access event for the given + * parameters is new or has expired (based on the expirationTimeInMs threshold), the event's + * timestamp is updated, effectively moving the event to the end of the tracking list. - If the + * tracker's capacity is reached, the oldest event is automatically removed to make room for new + * events. - This mechanism ensures that the tracker always contains the most recent access events + * within the specified capacity limit, with older, less relevant events being discarded. + */ +public class ApiAccessTracker { + private static final int INTERVAL_SECONDS = 30; + private static final int MAX_SIZE = 4096; + private final Map apiAccessMap; // Map + private final Deque apiAccessQueue; // hashes ordered by access time + private final long expirationTimeInMs; + private final int capacity; + + public ApiAccessTracker() { + this(MAX_SIZE, INTERVAL_SECONDS * 1000); + } + + public ApiAccessTracker(int capacity, long expirationTimeInMs) { + this.capacity = capacity; + this.expirationTimeInMs = expirationTimeInMs; + this.apiAccessMap = new ConcurrentHashMap<>(); + this.apiAccessQueue = new ConcurrentLinkedDeque<>(); + } + + /** + * Updates the API access log with the given route, method, and status code. If the record exists + * and is outdated, it is updated by moving to the end of the list. If the record does not exist, + * a new record is added. If the capacity limit is reached, the oldest record is removed. Returns + * true if the record was updated or added, false otherwise. + * + * @param route The route of the API endpoint request + * @param method The method of the API request + * @param statusCode The HTTP response status code of the API request + * @return return true if the record was updated or added, false otherwise + */ + public boolean updateApiAccessIfExpired(String route, String method, int statusCode) { + long currentTime = System.currentTimeMillis(); + long hash = computeApiHash(route, method, statusCode); + + // New or updated record + boolean isNewOrUpdated = false; + if (!apiAccessMap.containsKey(hash) + || currentTime - apiAccessMap.get(hash) > expirationTimeInMs) { + + cleanupExpiredEntries(currentTime); + + apiAccessMap.put(hash, currentTime); // Update timestamp + // move hash to the end of the queue + apiAccessQueue.remove(hash); + apiAccessQueue.addLast(hash); + isNewOrUpdated = true; + + // Remove the oldest hash if capacity is reached + while (apiAccessQueue.size() > this.capacity) { + Long oldestHash = apiAccessQueue.pollFirst(); + if (oldestHash != null) { + apiAccessMap.remove(oldestHash); + } + } + } + + return isNewOrUpdated; + } + + private void cleanupExpiredEntries(long currentTime) { + while (!apiAccessQueue.isEmpty()) { + Long oldestHash = apiAccessQueue.peekFirst(); + if (oldestHash == null) break; + + Long lastAccessTime = apiAccessMap.get(oldestHash); + if (lastAccessTime == null || currentTime - lastAccessTime > expirationTimeInMs) { + apiAccessQueue.pollFirst(); // remove from head + apiAccessMap.remove(oldestHash); + } else { + break; // is up-to-date + } + } + } + + private long computeApiHash(String route, String method, int statusCode) { + long result = 17; + result = 31 * result + route.hashCode(); + result = 31 * result + method.hashCode(); + result = 31 * result + statusCode; + return result; + } +} diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java index 38b5c631a4e..ddd8a10255b 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java @@ -1,56 +1,38 @@ package com.datadog.appsec.api.security; +import com.datadog.appsec.gateway.AppSecRequestContext; import datadog.trace.api.Config; -import java.util.concurrent.atomic.AtomicLong; public class ApiSecurityRequestSampler { - private volatile int sampling; - private final AtomicLong cumulativeCounter = new AtomicLong(); + private final ApiAccessTracker apiAccessTracker; + private final Config config; public ApiSecurityRequestSampler(final Config config) { - sampling = computeSamplingParameter(config.getApiSecurityRequestSampleRate()); + this.apiAccessTracker = new ApiAccessTracker(); + this.config = config; } - /** - * Sets the new sampling parameter - * - * @return {@code true} if the value changed - */ - public boolean setSampling(final float newSamplingFloat) { - int newSampling = computeSamplingParameter(newSamplingFloat); - if (newSampling != sampling) { - sampling = newSampling; - cumulativeCounter.set(0); // Reset current sampling counter - return true; + public boolean sampleRequest(AppSecRequestContext ctx) { + if (!config.isApiSecurityEnabled() || ctx == null) { + return false; } - return false; - } - - public int getSampling() { - return sampling; - } - public boolean sampleRequest() { - long prevValue = cumulativeCounter.getAndAdd(sampling); - long newValue = prevValue + sampling; - if (newValue / 100 == prevValue / 100 + 1) { - // Sample request - return true; + String route = ctx.getRoute(); + if (route == null) { + return false; } - // Skipped by sampling - return false; - } - static int computeSamplingParameter(final float pct) { - if (pct >= 1) { - return 100; + String method = ctx.getMethod(); + if (method == null) { + return false; } - if (pct < 0) { - // Api security can only be disabled by setting the sampling to zero, so we set it to 100%. - // TODO: We probably want a warning here. - return 100; + + int statusCode = ctx.getResponseStatus(); + if (statusCode == 0) { + return false; } - return (int) (pct * 100); + + return apiAccessTracker.updateApiAccessIfExpired(route, method, statusCode); } } 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 7afe609991e..71e7badbd82 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 @@ -90,6 +90,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private String scheme; private String method; private String savedRawURI; + private String route; private final Map> requestHeaders = new LinkedHashMap<>(); private final Map> responseHeaders = new LinkedHashMap<>(); private volatile Map> collectedCookies; @@ -360,7 +361,7 @@ void setScheme(String scheme) { this.scheme = scheme; } - String getMethod() { + public String getMethod() { return method; } @@ -368,7 +369,7 @@ void setMethod(String method) { this.method = method; } - String getSavedRawURI() { + public String getSavedRawURI() { return savedRawURI; } @@ -380,6 +381,18 @@ void setRawURI(String savedRawURI) { this.savedRawURI = savedRawURI; } + public String getRoute() { + return route; + } + + void setRoute(String route) { + if (this.route != null && this.route.compareToIgnoreCase(route) != 0) { + throw new IllegalStateException( + "Forbidden attempt to set different route for given request context"); + } + this.route = route; + } + void addRequestHeader(String name, String value) { if (finishedRequestHeaders) { throw new IllegalStateException("Request headers were said to be finished before"); diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index d8592343a62..78cf752f00e 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -725,8 +725,10 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { } } - if (!spanInfo.isRequiresPostProcessing()) { - ctx.close(); + // Route used in post-processing + Object route = spanInfo.getTags().get(Tags.HTTP_ROUTE); + if (route instanceof String) { + ctx.setRoute((String) route); } return NoopFlow.INSTANCE; } @@ -786,6 +788,7 @@ private void onRequestHeader(RequestContext ctx_, String name, String value) { } } + // This handler is executed in a separate thread due the computation possible overhead private void onPostProcessing(RequestContext ctx_) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null) { @@ -965,7 +968,7 @@ private Flow maybePublishResponseData(AppSecRequestContext ctx) { private void maybeExtractSchemas(AppSecRequestContext ctx) { boolean extractSchema = false; if (Config.get().isApiSecurityEnabled() && requestSampler != null) { - extractSchema = requestSampler.sampleRequest(); + extractSchema = requestSampler.sampleRequest(ctx); } if (!extractSchema) { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy new file mode 100644 index 00000000000..783b938fd4b --- /dev/null +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy @@ -0,0 +1,55 @@ +package com.datadog.appsec.api.security + +import datadog.trace.test.util.DDSpecification + +class ApiAccessTrackerTest extends DDSpecification { + def "should add new api access and update if expired"() { + given: "An ApiAccessTracker with capacity 2 and expiration time 1 second" + def tracker = new ApiAccessTracker(2, 1000) + + when: "Adding new api access" + tracker.updateApiAccessIfExpired("route1", "GET", 200) + def firstAccessTime = tracker.apiAccessMap.values().iterator().next() + + then: "The access is added" + tracker.apiAccessMap.size() == 1 + + when: "Waiting more than expiration time and adding another access with the same key" + Thread.sleep(1100) // Waiting more than 1 second to ensure expiration + tracker.updateApiAccessIfExpired("route1", "GET", 200) + def secondAccessTime = tracker.apiAccessMap.values().iterator().next() + + then: "The access is updated and moved to the end" + tracker.apiAccessMap.size() == 1 + secondAccessTime > firstAccessTime + } + + def "should remove the oldest access when capacity is exceeded"() { + given: "An ApiAccessTracker with capacity 1" + def tracker = new ApiAccessTracker(1, 1000) + + when: "Adding two api accesses" + tracker.updateApiAccessIfExpired("route1", "GET", 200) + Thread.sleep(100) // Delay to ensure different timestamps + tracker.updateApiAccessIfExpired("route2", "POST", 404) + + then: "The oldest access is removed" + tracker.apiAccessMap.size() == 1 + !tracker.apiAccessMap.containsKey(tracker.computeApiHash("route1", "GET", 200)) + tracker.apiAccessMap.containsKey(tracker.computeApiHash("route2", "POST", 404)) + } + + def "should not update access if not expired"() { + given: "An ApiAccessTracker with a short expiration time" + def tracker = new ApiAccessTracker(2, 2000) // 2 seconds expiration + + when: "Adding an api access and updating it before it expires" + tracker.updateApiAccessIfExpired("route1", "GET", 200) + def updateTime = System.currentTimeMillis() + boolean updatedBeforeExpiration = tracker.updateApiAccessIfExpired("route1", "GET", 200) + + then: "The access is not updated" + !updatedBeforeExpiration + tracker.apiAccessMap.get(tracker.computeApiHash("route1", "GET", 200)) == updateTime + } +} \ No newline at end of file diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy index ff256bc21fb..26c1dd6b42d 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy @@ -1,55 +1,35 @@ package com.datadog.appsec.api.security +import com.datadog.appsec.gateway.AppSecRequestContext import datadog.trace.api.Config import datadog.trace.test.util.DDSpecification -import spock.lang.Shared class ApiSecurityRequestSamplerTest extends DDSpecification { - @Shared - static final float DEFAULT_SAMPLE_RATE = Config.get().getApiSecurityRequestSampleRate() + def config = Mock(Config) { + isApiSecurityEnabled() >> true + } - void 'Api Security Request Sample Rate'() { - given: - def config = Spy(Config.get()) - config.getApiSecurityRequestSampleRate() >> sampleRate - def sampler = new ApiSecurityRequestSampler(config) + def sampler = new ApiSecurityRequestSampler(config) + void 'Api Security Sample Request'() { when: - def numOfRequest = expectedSampledRequests.size() - def results = new int[numOfRequest] - for (int i = 0; i < numOfRequest; i++) { - results[i] = sampler.sampleRequest() ? 1 : 0 + def span = Mock(AppSecRequestContext) { + getSavedRawURI() >> route + getMethod() >> method + getResponseStatus() >> statusCode } + def sample = sampler.sampleRequest(span) then: - results == expectedSampledRequests as int[] + sample == sampleResult where: - sampleRate | expectedSampledRequests - DEFAULT_SAMPLE_RATE | [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0] // Default sample rate - 10% - 0.0 | [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] - 0.1 | [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0] - 0.25 | [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1] - 0.33 | [0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1] - 0.5 | [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - 0.75 | [0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1] - 0.9 | [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0] - 0.99 | [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] - 1.0 | [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] - 1.25 | [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] // Wrong sample rate - use 100% - -0.5 | [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] // Wrong sample rate - use 100% - } - - void 'update sample rate'() { - given: - def config = Spy(Config.get()) - def sampler = new ApiSecurityRequestSampler(config) - - when: - sampler.setSampling(0.2) - - then: - sampler.sampling == 20 + method | route | statusCode | sampleResult + 'GET' | 'route1' | 200 | true + 'GET' | 'route2' | null | false + 'GET' | null | 404 | false + 'TOP' | 999 | 404 | true + null | '999' | 404 | false } -} +} \ No newline at end of file diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 50269370666..f10acdb6031 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -169,7 +169,6 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * traceSegment.setTagTop('http.request.headers.accept', 'header_value') 1 * traceSegment.setTagTop('http.response.headers.content-type', 'text/html; charset=UTF-8') 1 * traceSegment.setTagTop('network.client.ip', '2001::1') - 1 * mockAppSecCtx.closeAdditive() flow.result == null flow.action == Flow.Action.Noop.INSTANCE } From 30a1592263582de7d6954570c8e9190b473c46e1 Mon Sep 17 00:00:00 2001 From: Valentin Zakharov Date: Wed, 12 Feb 2025 10:43:33 +0100 Subject: [PATCH 04/15] Added preSampler for short-circuit Api Security --- .../appsec/api/security/ApiAccessTracker.java | 7 +++++ .../security/ApiSecurityRequestSampler.java | 27 ++++++++++--------- .../datadog/appsec/gateway/GatewayBridge.java | 12 ++++++--- .../ApiSecurityRequestSamplerTest.groovy | 2 +- .../gateway/GatewayBridgeSpecification.groovy | 11 +++++--- 5 files changed, 39 insertions(+), 20 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java index 32aaa761717..4c53d63a526 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java @@ -77,6 +77,13 @@ public boolean updateApiAccessIfExpired(String route, String method, int statusC return isNewOrUpdated; } + public boolean isApiAccessExpired(String route, String method, int statusCode) { + long currentTime = System.currentTimeMillis(); + long hash = computeApiHash(route, method, statusCode); + return !apiAccessMap.containsKey(hash) + || currentTime - apiAccessMap.get(hash) > expirationTimeInMs; + } + private void cleanupExpiredEntries(long currentTime) { while (!apiAccessQueue.isEmpty()) { Long oldestHash = apiAccessQueue.peekFirst(); diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java index ddd8a10255b..739fd9d89a6 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java @@ -14,25 +14,28 @@ public ApiSecurityRequestSampler(final Config config) { } public boolean sampleRequest(AppSecRequestContext ctx) { - if (!config.isApiSecurityEnabled() || ctx == null) { + if (!isValid(ctx)) { return false; } - String route = ctx.getRoute(); - if (route == null) { - return false; - } + return apiAccessTracker.updateApiAccessIfExpired( + ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus()); + } - String method = ctx.getMethod(); - if (method == null) { + public boolean preSampleRequest(AppSecRequestContext ctx) { + if (!isValid(ctx)) { return false; } - int statusCode = ctx.getResponseStatus(); - if (statusCode == 0) { - return false; - } + return apiAccessTracker.isApiAccessExpired( + ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus()); + } - return apiAccessTracker.updateApiAccessIfExpired(route, method, statusCode); + private boolean isValid(AppSecRequestContext ctx) { + return config.isApiSecurityEnabled() + && ctx != null + && ctx.getRoute() != null + && ctx.getMethod() != null + && ctx.getResponseStatus() != 0; } } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index 78cf752f00e..6c8e2c7c448 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -5,7 +5,6 @@ import static com.datadog.appsec.gateway.AppSecRequestContext.DEFAULT_REQUEST_HEADERS_ALLOW_LIST; import static com.datadog.appsec.gateway.AppSecRequestContext.REQUEST_HEADERS_ALLOW_LIST; import static com.datadog.appsec.gateway.AppSecRequestContext.RESPONSE_HEADERS_ALLOW_LIST; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import com.datadog.appsec.AppSecSystem; import com.datadog.appsec.api.security.ApiSecurityRequestSampler; @@ -664,6 +663,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { ctx.setRequestEndCalled(); TraceSegment traceSeg = ctx_.getTraceSegment(); + Map tags = spanInfo.getTags(); // AppSec report metric and events for web span only if (traceSeg != null) { @@ -685,7 +685,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { traceSeg.setTagTop("network.client.ip", ctx.getPeerAddress()); // Reflect client_ip as actor.ip for backward compatibility - Object clientIp = spanInfo.getTags().get(Tags.HTTP_CLIENT_IP); + Object clientIp = tags.get(Tags.HTTP_CLIENT_IP); if (clientIp != null) { traceSeg.setTagTop("actor.ip", clientIp); } @@ -726,10 +726,15 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { } // Route used in post-processing - Object route = spanInfo.getTags().get(Tags.HTTP_ROUTE); + Object route = tags.get(Tags.HTTP_ROUTE); if (route instanceof String) { ctx.setRoute((String) route); } + if (requestSampler.preSampleRequest(ctx)) { + // The request is pre-sampled - we need to post-process it + spanInfo.setRequiresPostProcessing(true); + } + return NoopFlow.INSTANCE; } @@ -924,7 +929,6 @@ private Flow maybePublishRequestData(AppSecRequestContext ctx) { try { GatewayContext gwCtx = new GatewayContext(false); - activeSpan().getLocalRootSpan().setRequiresPostProcessing(true); return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx); } catch (ExpiredSubscriberInfoException e) { this.initialReqDataSubInfo = null; diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy index 26c1dd6b42d..de9fb0a49da 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy @@ -15,7 +15,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { void 'Api Security Sample Request'() { when: def span = Mock(AppSecRequestContext) { - getSavedRawURI() >> route + getRoute() >> route getMethod() >> method getResponseStatus() >> statusCode } diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index f10acdb6031..70d3703fc55 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -1,6 +1,7 @@ package com.datadog.appsec.gateway import com.datadog.appsec.AppSecSystem +import com.datadog.appsec.api.security.ApiSecurityRequestSampler import com.datadog.appsec.config.TraceSegmentPostProcessor import com.datadog.appsec.event.EventDispatcher import com.datadog.appsec.event.EventProducerService @@ -76,7 +77,10 @@ class GatewayBridgeSpecification extends DDSpecification { } TraceSegmentPostProcessor pp = Mock() - GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, [pp]) + ApiSecurityRequestSampler requestSampler = Mock(ApiSecurityRequestSampler) { + preSampleRequest(_ as AppSecRequestContext) >> false + } + GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, requestSampler, [pp]) Supplier> requestStartedCB BiFunction> requestEndedCB @@ -161,7 +165,6 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * spanInfo.getTags() >> ['http.client_ip': '1.1.1.1'] 1 * mockAppSecCtx.transferCollectedEvents() >> [event] 1 * mockAppSecCtx.peerAddress >> '2001::1' - 1 * mockAppSecCtx.close() 1 * traceSegment.setTagTop("_dd.appsec.enabled", 1) 1 * traceSegment.setTagTop("_dd.runtime_family", "jvm") 1 * traceSegment.setTagTop('appsec.event', true) @@ -971,7 +974,9 @@ class GatewayBridgeSpecification extends DDSpecification { getData(RequestContextSlot.APPSEC) >> mockAppSecCtx getTraceSegment() >> traceSegment } - final spanInfo = Mock(AgentSpan) + final spanInfo = Mock(AgentSpan) { + getTags() >> ['http.route':'/'] + } when: requestEndedCB.apply(mockCtx, spanInfo) From 6a719869e8b4477fa1f03fbc9b4fb30feef17dc5 Mon Sep 17 00:00:00 2001 From: Valentin Zakharov Date: Wed, 12 Feb 2025 13:15:17 +0100 Subject: [PATCH 05/15] Added post-processing limit --- .../com/datadog/iast/overhead/OverheadContext.java | 2 +- .../com/datadog/iast/overhead/OverheadController.java | 2 +- .../datadog/appsec/api/security/ApiAccessTracker.java | 11 ++++++----- .../com/datadog/appsec/gateway/GatewayBridge.java | 9 ++++++++- .../datadog/trace}/util/NonBlockingSemaphore.java | 2 +- .../trace}/util/NonBlockingSemaphoreTest.groovy | 2 +- 6 files changed, 18 insertions(+), 10 deletions(-) rename {dd-java-agent/agent-iast/src/main/java/com/datadog/iast => internal-api/src/main/java/datadog/trace}/util/NonBlockingSemaphore.java (98%) rename {dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast => internal-api/src/test/groovy/datadog/trace}/util/NonBlockingSemaphoreTest.groovy (99%) diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java index b34b533f43a..c6893ccdd8b 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java @@ -2,7 +2,7 @@ import static datadog.trace.api.iast.IastDetectionMode.UNLIMITED; -import com.datadog.iast.util.NonBlockingSemaphore; +import datadog.trace.util.NonBlockingSemaphore; public class OverheadContext { diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java index 0cd4575056a..d30be2432d6 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java @@ -4,7 +4,6 @@ import com.datadog.iast.IastRequestContext; import com.datadog.iast.IastSystem; -import com.datadog.iast.util.NonBlockingSemaphore; import datadog.trace.api.Config; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; @@ -13,6 +12,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.util.AgentTaskScheduler; +import datadog.trace.util.NonBlockingSemaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java index 4c53d63a526..edf22bae2bb 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java @@ -38,10 +38,11 @@ public ApiAccessTracker(int capacity, long expirationTimeInMs) { } /** - * Updates the API access log with the given route, method, and status code. If the record exists - * and is outdated, it is updated by moving to the end of the list. If the record does not exist, - * a new record is added. If the capacity limit is reached, the oldest record is removed. Returns - * true if the record was updated or added, false otherwise. + * Updates the API access log with the given route, method, and status code. If the record already + * exists and is outdated, it is updated by moving to the end of the list. If the record does not + * exist, a new record is added. If the capacity limit is reached, the oldest record is removed. + * This method should not be called concurrently by multiple threads, due absence of additional + * synchronization for updating data structures is not required. * * @param route The route of the API endpoint request * @param method The method of the API request @@ -66,7 +67,7 @@ public boolean updateApiAccessIfExpired(String route, String method, int statusC isNewOrUpdated = true; // Remove the oldest hash if capacity is reached - while (apiAccessQueue.size() > this.capacity) { + while (apiAccessMap.size() > this.capacity) { Long oldestHash = apiAccessQueue.pollFirst(); if (oldestHash != null) { apiAccessMap.remove(oldestHash); diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index 6c8e2c7c448..f6bf2646e8f 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -35,6 +35,7 @@ import datadog.trace.api.telemetry.WafMetricCollector; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; +import datadog.trace.util.NonBlockingSemaphore; import datadog.trace.util.stacktrace.StackTraceEvent; import datadog.trace.util.stacktrace.StackUtils; import java.net.URI; @@ -82,6 +83,10 @@ public class GatewayBridge { private static final String METASTRUCT_EXPLOIT = "exploit"; + private final int MAX_POST_PROCESSING_TASKS = 16; + private final NonBlockingSemaphore postProcessingCounter = + NonBlockingSemaphore.withPermitCount(MAX_POST_PROCESSING_TASKS); + private final SubscriptionService subscriptionService; private final EventProducerService producerService; private final ApiSecurityRequestSampler requestSampler; @@ -730,7 +735,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { if (route instanceof String) { ctx.setRoute((String) route); } - if (requestSampler.preSampleRequest(ctx)) { + if (requestSampler.preSampleRequest(ctx) && postProcessingCounter.acquire()) { // The request is pre-sampled - we need to post-process it spanInfo.setRequiresPostProcessing(true); } @@ -802,6 +807,8 @@ private void onPostProcessing(RequestContext ctx_) { maybeExtractSchemas(ctx); ctx.close(); + // Decrease the counter to allow the next request to be post-processed + postProcessingCounter.release(); } public void stop() { diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java b/internal-api/src/main/java/datadog/trace/util/NonBlockingSemaphore.java similarity index 98% rename from dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java rename to internal-api/src/main/java/datadog/trace/util/NonBlockingSemaphore.java index ef6e529980a..e4f2249c5e5 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java +++ b/internal-api/src/main/java/datadog/trace/util/NonBlockingSemaphore.java @@ -1,4 +1,4 @@ -package com.datadog.iast.util; +package datadog.trace.util; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy b/internal-api/src/test/groovy/datadog/trace/util/NonBlockingSemaphoreTest.groovy similarity index 99% rename from dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy rename to internal-api/src/test/groovy/datadog/trace/util/NonBlockingSemaphoreTest.groovy index a8247cac497..1c74ddcc826 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/util/NonBlockingSemaphoreTest.groovy @@ -1,4 +1,4 @@ -package com.datadog.iast.util +package datadog.trace.util import groovy.transform.CompileDynamic From 64539edea675fcfece7211efd9fefdcab291fdbb Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Mon, 24 Feb 2025 13:50:10 +0100 Subject: [PATCH 06/15] wip --- .../iast/overhead/OverheadContext.java | 2 +- .../iast/overhead/OverheadController.java | 2 +- .../iast}/util/NonBlockingSemaphore.java | 2 +- .../util/NonBlockingSemaphoreTest.groovy | 2 +- dd-java-agent/appsec/build.gradle | 1 + .../java/com/datadog/appsec/AppSecSystem.java | 11 +- .../appsec/api/security/ApiAccessTracker.java | 110 -------- .../security/ApiSecurityRequestSampler.java | 187 +++++++++++-- .../api/security/AppSecSpanPostProcessor.java | 90 +++++++ .../config/AppSecConfigServiceImpl.java | 3 - .../appsec/gateway/AppSecRequestContext.java | 66 +++-- .../datadog/appsec/gateway/GatewayBridge.java | 82 ++---- .../api/security/ApiAccessTrackerTest.groovy | 55 ---- .../ApiSecurityRequestSamplerTest.groovy | 222 ++++++++++++++-- .../AppSecSpanPostProcessorTest.groovy | 245 ++++++++++++++++++ ...ppSecConfigServiceImplSpecification.groovy | 4 - .../gateway/GatewayBridgeSpecification.groovy | 6 +- .../springboot/controller/WebController.java | 5 + .../appsec/SpringBootSmokeTest.groovy | 35 +++ .../AbstractAppSecServerSmokeTest.groovy | 4 + .../datadog/trace/api/ConfigDefaults.java | 2 +- .../trace/api/config/AppSecConfig.java | 2 +- .../trace/common/writer/DDAgentWriter.java | 6 +- .../trace/common/writer/DDIntakeWriter.java | 7 +- .../common/writer/TraceProcessingWorker.java | 63 ++--- .../java/datadog/trace/core/CoreTracer.java | 4 +- .../main/java/datadog/trace/core/DDSpan.java | 10 - .../datadog/trace/core/DDSpanContext.java | 9 - .../AppSecSpanPostProcessor.java | 40 --- .../core/postprocessor/SpanPostProcessor.java | 26 -- .../writer/TraceProcessingWorkerTest.groovy | 28 +- .../AppSecSpanPostProcessorTest.groovy | 119 --------- internal-api/build.gradle | 5 +- .../main/java/datadog/trace/api/Config.java | 13 +- .../datadog/trace/api/gateway/Events.java | 13 - .../datadog/trace/api/gateway/IGSpanInfo.java | 4 - .../api/gateway/InstrumentationGateway.java | 14 - .../instrumentation/api/ExtractedSpan.java | 5 - .../instrumentation/api/ImmutableSpan.java | 3 - .../instrumentation/api/NoopSpan.java | 5 - .../api/SpanPostProcessor.java | 37 +++ .../gateway/InstrumentationGatewayTest.java | 18 -- 42 files changed, 917 insertions(+), 650 deletions(-) rename {internal-api/src/main/java/datadog/trace => dd-java-agent/agent-iast/src/main/java/com/datadog/iast}/util/NonBlockingSemaphore.java (98%) rename {internal-api/src/test/groovy/datadog/trace => dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast}/util/NonBlockingSemaphoreTest.groovy (99%) delete mode 100644 dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java create mode 100644 dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java delete mode 100644 dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy create mode 100644 dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy delete mode 100644 dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java delete mode 100644 dd-trace-core/src/main/java/datadog/trace/core/postprocessor/SpanPostProcessor.java delete mode 100644 dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy create mode 100644 internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanPostProcessor.java diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java index c6893ccdd8b..b34b533f43a 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java @@ -2,7 +2,7 @@ import static datadog.trace.api.iast.IastDetectionMode.UNLIMITED; -import datadog.trace.util.NonBlockingSemaphore; +import com.datadog.iast.util.NonBlockingSemaphore; public class OverheadContext { diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java index d30be2432d6..0cd4575056a 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java @@ -4,6 +4,7 @@ import com.datadog.iast.IastRequestContext; import com.datadog.iast.IastSystem; +import com.datadog.iast.util.NonBlockingSemaphore; import datadog.trace.api.Config; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; @@ -12,7 +13,6 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.util.AgentTaskScheduler; -import datadog.trace.util.NonBlockingSemaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; diff --git a/internal-api/src/main/java/datadog/trace/util/NonBlockingSemaphore.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java similarity index 98% rename from internal-api/src/main/java/datadog/trace/util/NonBlockingSemaphore.java rename to dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java index e4f2249c5e5..ef6e529980a 100644 --- a/internal-api/src/main/java/datadog/trace/util/NonBlockingSemaphore.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java @@ -1,4 +1,4 @@ -package datadog.trace.util; +package com.datadog.iast.util; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; diff --git a/internal-api/src/test/groovy/datadog/trace/util/NonBlockingSemaphoreTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy similarity index 99% rename from internal-api/src/test/groovy/datadog/trace/util/NonBlockingSemaphoreTest.groovy rename to dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy index 1c74ddcc826..a8247cac497 100644 --- a/internal-api/src/test/groovy/datadog/trace/util/NonBlockingSemaphoreTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy @@ -1,4 +1,4 @@ -package datadog.trace.util +package com.datadog.iast.util import groovy.transform.CompileDynamic diff --git a/dd-java-agent/appsec/build.gradle b/dd-java-agent/appsec/build.gradle index 38f999b62aa..de195400063 100644 --- a/dd-java-agent/appsec/build.gradle +++ b/dd-java-agent/appsec/build.gradle @@ -82,6 +82,7 @@ ext { 'com.datadog.appsec.config.AppSecFeatures.ApiSecurity', 'com.datadog.appsec.config.AppSecFeatures.AutoUserInstrum', 'com.datadog.appsec.event.ReplaceableEventProducerService', + 'com.datadog.appsec.api.security.ApiSecurityRequestSampler.NoOp', ] excludedClassesBranchCoverage = [ 'com.datadog.appsec.gateway.GatewayBridge', diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java index 9d86b681cc1..9900268ba73 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java @@ -1,6 +1,7 @@ package com.datadog.appsec; import com.datadog.appsec.api.security.ApiSecurityRequestSampler; +import com.datadog.appsec.api.security.AppSecSpanPostProcessor; import com.datadog.appsec.blocking.BlockingServiceImpl; import com.datadog.appsec.config.AppSecConfigService; import com.datadog.appsec.config.AppSecConfigServiceImpl; @@ -21,6 +22,7 @@ import datadog.trace.api.telemetry.ProductChange; import datadog.trace.api.telemetry.ProductChangeCollector; import datadog.trace.bootstrap.ActiveSubsystems; +import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -66,7 +68,14 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s EventDispatcher eventDispatcher = new EventDispatcher(); REPLACEABLE_EVENT_PRODUCER.replaceEventProducerService(eventDispatcher); - ApiSecurityRequestSampler requestSampler = new ApiSecurityRequestSampler(config); + ApiSecurityRequestSampler requestSampler; + if (Config.get().isApiSecurityEnabled()) { + requestSampler = new ApiSecurityRequestSampler(); + SpanPostProcessor.Holder.INSTANCE = + new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER); + } else { + requestSampler = new ApiSecurityRequestSampler.NoOp(); + } ConfigurationPoller configurationPoller = sco.configurationPoller(config); // may throw and abort startup diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java deleted file mode 100644 index edf22bae2bb..00000000000 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java +++ /dev/null @@ -1,110 +0,0 @@ -package com.datadog.appsec.api.security; - -import java.util.Deque; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentLinkedDeque; - -/** - * The ApiAccessTracker class provides a mechanism to track API access events, managing them within - * a specified capacity limit. Each event is associated with a unique combination of route, method, - * and status code, which is used to generate a unique key for tracking access timestamps. - * - *

Usage: - When an API access event occurs, the `updateApiAccessIfExpired` method is called with - * the route, method, and status code of the API request. - If the access event for the given - * parameters is new or has expired (based on the expirationTimeInMs threshold), the event's - * timestamp is updated, effectively moving the event to the end of the tracking list. - If the - * tracker's capacity is reached, the oldest event is automatically removed to make room for new - * events. - This mechanism ensures that the tracker always contains the most recent access events - * within the specified capacity limit, with older, less relevant events being discarded. - */ -public class ApiAccessTracker { - private static final int INTERVAL_SECONDS = 30; - private static final int MAX_SIZE = 4096; - private final Map apiAccessMap; // Map - private final Deque apiAccessQueue; // hashes ordered by access time - private final long expirationTimeInMs; - private final int capacity; - - public ApiAccessTracker() { - this(MAX_SIZE, INTERVAL_SECONDS * 1000); - } - - public ApiAccessTracker(int capacity, long expirationTimeInMs) { - this.capacity = capacity; - this.expirationTimeInMs = expirationTimeInMs; - this.apiAccessMap = new ConcurrentHashMap<>(); - this.apiAccessQueue = new ConcurrentLinkedDeque<>(); - } - - /** - * Updates the API access log with the given route, method, and status code. If the record already - * exists and is outdated, it is updated by moving to the end of the list. If the record does not - * exist, a new record is added. If the capacity limit is reached, the oldest record is removed. - * This method should not be called concurrently by multiple threads, due absence of additional - * synchronization for updating data structures is not required. - * - * @param route The route of the API endpoint request - * @param method The method of the API request - * @param statusCode The HTTP response status code of the API request - * @return return true if the record was updated or added, false otherwise - */ - public boolean updateApiAccessIfExpired(String route, String method, int statusCode) { - long currentTime = System.currentTimeMillis(); - long hash = computeApiHash(route, method, statusCode); - - // New or updated record - boolean isNewOrUpdated = false; - if (!apiAccessMap.containsKey(hash) - || currentTime - apiAccessMap.get(hash) > expirationTimeInMs) { - - cleanupExpiredEntries(currentTime); - - apiAccessMap.put(hash, currentTime); // Update timestamp - // move hash to the end of the queue - apiAccessQueue.remove(hash); - apiAccessQueue.addLast(hash); - isNewOrUpdated = true; - - // Remove the oldest hash if capacity is reached - while (apiAccessMap.size() > this.capacity) { - Long oldestHash = apiAccessQueue.pollFirst(); - if (oldestHash != null) { - apiAccessMap.remove(oldestHash); - } - } - } - - return isNewOrUpdated; - } - - public boolean isApiAccessExpired(String route, String method, int statusCode) { - long currentTime = System.currentTimeMillis(); - long hash = computeApiHash(route, method, statusCode); - return !apiAccessMap.containsKey(hash) - || currentTime - apiAccessMap.get(hash) > expirationTimeInMs; - } - - private void cleanupExpiredEntries(long currentTime) { - while (!apiAccessQueue.isEmpty()) { - Long oldestHash = apiAccessQueue.peekFirst(); - if (oldestHash == null) break; - - Long lastAccessTime = apiAccessMap.get(oldestHash); - if (lastAccessTime == null || currentTime - lastAccessTime > expirationTimeInMs) { - apiAccessQueue.pollFirst(); // remove from head - apiAccessMap.remove(oldestHash); - } else { - break; // is up-to-date - } - } - } - - private long computeApiHash(String route, String method, int statusCode) { - long result = 17; - result = 31 * result + route.hashCode(); - result = 31 * result + method.hashCode(); - result = 31 * result + statusCode; - return result; - } -} diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java index 739fd9d89a6..56c566c3937 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java @@ -2,40 +2,191 @@ import com.datadog.appsec.gateway.AppSecRequestContext; import datadog.trace.api.Config; +import datadog.trace.api.time.SystemTimeSource; +import datadog.trace.api.time.TimeSource; +import java.util.Deque; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedDeque; +import java.util.concurrent.Semaphore; +import javax.annotation.Nonnull; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class ApiSecurityRequestSampler { - private final ApiAccessTracker apiAccessTracker; - private final Config config; + private static final Logger log = LoggerFactory.getLogger(ApiSecurityRequestSampler.class); - public ApiSecurityRequestSampler(final Config config) { - this.apiAccessTracker = new ApiAccessTracker(); - this.config = config; + /** + * A maximum number of request contexts we'll keep open past the end of request at any given time. + * This will avoid excessive memory usage in case of a high number of concurrent requests, and + * should also prevent memory leaks. + */ + private static final int MAX_POST_PROCESSING_TASKS = 4; + /** Maximum number of entries in the access map. */ + private static final int MAX_SIZE = 4096; + /** Mapping from endpoint hash to last access timestamp in millis. */ + private final ConcurrentHashMap accessMap; + /** Deque of endpoint hashes ordered by access time. Oldest is always first. */ + private final Deque accessDeque; + + private final long expirationTimeInMs; + private final int capacity; + private final TimeSource timeSource; + private final Semaphore counter = new Semaphore(MAX_POST_PROCESSING_TASKS); + + public ApiSecurityRequestSampler() { + this( + MAX_SIZE, + (long) (Config.get().getApiSecuritySampleDelay() * 1_000), + SystemTimeSource.INSTANCE); + } + + public ApiSecurityRequestSampler( + int capacity, long expirationTimeInMs, @Nonnull TimeSource timeSource) { + this.capacity = capacity; + this.expirationTimeInMs = expirationTimeInMs; + this.accessMap = new ConcurrentHashMap<>(); + this.accessDeque = new ConcurrentLinkedDeque<>(); + this.timeSource = timeSource; + } + + /** + * Prepare a request context for later sampling decision. This method should be called at request + * end, and is thread-safe. If a request can potentially be sampled, this method will return true. + * If this method returns true, the caller MUST call {@link #releaseOne()} once the context is not + * needed anymore. + */ + public boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx) { + final String route = ctx.getRoute(); + if (route == null) { + log.debug("Route is null, skipping API security sampling"); + return false; + } + final String method = ctx.getMethod(); + if (method == null) { + log.debug("Method is null, skipping API security sampling"); + return false; + } + final int statusCode = ctx.getResponseStatus(); + if (statusCode == 0) { + log.debug("Status code is 0, skipping API security sampling"); + return false; + } + long hash = computeApiHash(route, method, statusCode); + ctx.setApiSecurityEndpointHash(hash); + if (!isApiAccessExpired(hash)) { + log.debug("API security sampling is not required for this request"); + return false; + } + if (counter.tryAcquire()) { + log.debug("API security sampling is required for this request (presampled)"); + ctx.setKeepOpenForApiSecurityPostProcessing(true); + return true; + } + return false; } + /** Get the final sampling decision. This method is NOT thread-safe. */ public boolean sampleRequest(AppSecRequestContext ctx) { - if (!isValid(ctx)) { + if (ctx == null) { return false; } + final Long hash = ctx.getApiSecurityEndpointHash(); + if (hash == null) { + // This should never happen, it should have been short-circuited before. + return false; + } + return updateApiAccessIfExpired(hash); + } - return apiAccessTracker.updateApiAccessIfExpired( - ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus()); + /** Release one permit for the sampler. This must be called after processing a span. */ + public void releaseOne() { + counter.release(); } - public boolean preSampleRequest(AppSecRequestContext ctx) { - if (!isValid(ctx)) { + private boolean updateApiAccessIfExpired(final long hash) { + final long currentTime = timeSource.getCurrentTimeMillis(); + + Long lastAccess = accessMap.get(hash); + if (lastAccess != null && currentTime - lastAccess < expirationTimeInMs) { return false; } - return apiAccessTracker.isApiAccessExpired( - ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus()); + if (accessMap.put(hash, currentTime) == null) { + accessDeque.addLast(hash); + // If we added a new entry, we perform purging. + cleanupExpiredEntries(currentTime); + } else { + // This is now the most recently accessed entry. + accessDeque.remove(hash); + accessDeque.addLast(hash); + } + + return true; + } + + private boolean isApiAccessExpired(final long hash) { + final long currentTime = timeSource.getCurrentTimeMillis(); + final Long lastAccess = accessMap.get(hash); + return lastAccess == null || currentTime - lastAccess >= expirationTimeInMs; + } + + private void cleanupExpiredEntries(final long currentTime) { + // Purge all expired entries. + while (!accessDeque.isEmpty()) { + final Long oldestHash = accessDeque.peekFirst(); + if (oldestHash == null) { + // Should never happen + continue; + } + + final Long lastAccessTime = accessMap.get(oldestHash); + if (lastAccessTime == null) { + // Should never happen + continue; + } + + if (currentTime - lastAccessTime < expirationTimeInMs) { + // The oldest hash is up-to-date, so stop here. + break; + } + + accessDeque.pollFirst(); + accessMap.remove(oldestHash); + } + + // If we went over capacity, remove the oldest entries until we are within the limit. + // This should never be more than 1. + final int toRemove = accessMap.size() - this.capacity; + for (int i = 0; i < toRemove; i++) { + Long oldestHash = accessDeque.pollFirst(); + if (oldestHash != null) { + accessMap.remove(oldestHash); + } + } + } + + private long computeApiHash(final String route, final String method, final int statusCode) { + long result = 17; + result = 31 * result + route.hashCode(); + result = 31 * result + method.hashCode(); + result = 31 * result + statusCode; + return result; } - private boolean isValid(AppSecRequestContext ctx) { - return config.isApiSecurityEnabled() - && ctx != null - && ctx.getRoute() != null - && ctx.getMethod() != null - && ctx.getResponseStatus() != 0; + public static final class NoOp extends ApiSecurityRequestSampler { + public NoOp() { + super(0, 0, SystemTimeSource.INSTANCE); + } + + @Override + public boolean preSampleRequest(@Nonnull AppSecRequestContext ctx) { + return false; + } + + @Override + public boolean sampleRequest(AppSecRequestContext ctx) { + return false; + } } } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java new file mode 100644 index 00000000000..5fd0729b200 --- /dev/null +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java @@ -0,0 +1,90 @@ +package com.datadog.appsec.api.security; + +import com.datadog.appsec.event.EventProducerService; +import com.datadog.appsec.event.ExpiredSubscriberInfoException; +import com.datadog.appsec.event.data.DataBundle; +import com.datadog.appsec.event.data.KnownAddresses; +import com.datadog.appsec.event.data.SingletonDataBundle; +import com.datadog.appsec.gateway.AppSecRequestContext; +import com.datadog.appsec.gateway.GatewayContext; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.internal.TraceSegment; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor; +import java.util.Collections; +import java.util.function.BooleanSupplier; +import javax.annotation.Nonnull; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class AppSecSpanPostProcessor implements SpanPostProcessor { + + private static final Logger log = LoggerFactory.getLogger(AppSecSpanPostProcessor.class); + private final ApiSecurityRequestSampler sampler; + private final EventProducerService producerService; + + public AppSecSpanPostProcessor( + ApiSecurityRequestSampler sampler, EventProducerService producerService) { + this.sampler = sampler; + this.producerService = producerService; + } + + @Override + public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutCheck) { + final RequestContext ctx_ = span.getRequestContext(); + if (ctx_ == null) { + return; + } + final AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); + if (ctx == null) { + return; + } + + if (!ctx.isKeepOpenForApiSecurityPostProcessing()) { + log.debug("Request context not marked for API security post-processing"); + return; + } + + try { + if (timeoutCheck.getAsBoolean()) { + log.debug("Timeout detected, skipping API security post-processing"); + return; + } + if (!sampler.sampleRequest(ctx)) { + log.debug("Request not sampled, skipping API security post-processing"); + return; + } + log.debug("Request sampled, processing API security post-processing"); + extractSchemas(ctx, ctx_.getTraceSegment()); + } finally { + ctx.setKeepOpenForApiSecurityPostProcessing(false); + try { + ctx.close(); + } catch (Exception e) { + log.debug("Error closing AppSecRequestContext", e); + } + sampler.releaseOne(); + } + } + + private void extractSchemas(final AppSecRequestContext ctx, final TraceSegment traceSegment) { + final EventProducerService.DataSubscriberInfo sub = + producerService.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR); + if (sub == null || sub.isEmpty()) { + log.debug("No subscribers for schema extraction"); + return; + } + + final DataBundle bundle = + new SingletonDataBundle<>( + KnownAddresses.WAF_CONTEXT_PROCESSOR, Collections.singletonMap("extract-schema", true)); + try { + GatewayContext gwCtx = new GatewayContext(false); + producerService.publishDataEvent(sub, ctx, bundle, gwCtx); + ctx.commitDerivatives(traceSegment); + } catch (ExpiredSubscriberInfoException e) { + log.debug("Subscriber info expired", e); + } + } +} diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java index 069733a03fd..0ebf33bb137 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java @@ -2,7 +2,6 @@ import static com.datadog.appsec.util.StandardizedLogging.RulesInvalidReason.INVALID_JSON_FILE; import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_ACTIVATION; -import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE; import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE; import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE; import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_CUSTOM_RULES; @@ -200,7 +199,6 @@ private void subscribeAsmFeatures() { log.debug("Will not subscribe report CAPABILITY_ASM_ACTIVATION (AppSec explicitly enabled)"); } this.configurationPoller.addCapabilities(CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE); - this.configurationPoller.addCapabilities(CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE); } private void distributeSubConfigurations( @@ -362,7 +360,6 @@ public void close() { | CAPABILITY_ASM_CUSTOM_RULES | CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE | CAPABILITY_ASM_TRUSTED_IPS - | CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE | CAPABILITY_ASM_RASP_SQLI | CAPABILITY_ASM_RASP_SSRF | CAPABILITY_ASM_RASP_LFI 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 71e7badbd82..1e748e3048a 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 @@ -142,6 +142,9 @@ public class AppSecRequestContext implements DataBundle, Closeable { // Used to detect missing request-end event at close. private volatile boolean requestEndCalled; + private volatile boolean keepOpenForApiSecurityPostProcessing; + private volatile Long apiSecurityEndpointHash; + private static final AtomicIntegerFieldUpdater WAF_TIMEOUTS_UPDATER = AtomicIntegerFieldUpdater.newUpdater(AppSecRequestContext.class, "wafTimeouts"); private static final AtomicIntegerFieldUpdater RASP_TIMEOUTS_UPDATER = @@ -369,7 +372,7 @@ void setMethod(String method) { this.method = method; } - public String getSavedRawURI() { + String getSavedRawURI() { return savedRawURI; } @@ -385,14 +388,26 @@ public String getRoute() { return route; } - void setRoute(String route) { - if (this.route != null && this.route.compareToIgnoreCase(route) != 0) { - throw new IllegalStateException( - "Forbidden attempt to set different route for given request context"); - } + public void setRoute(String route) { this.route = route; } + public void setKeepOpenForApiSecurityPostProcessing(final boolean flag) { + this.keepOpenForApiSecurityPostProcessing = flag; + } + + public boolean isKeepOpenForApiSecurityPostProcessing() { + return this.keepOpenForApiSecurityPostProcessing; + } + + public void setApiSecurityEndpointHash(long hash) { + this.apiSecurityEndpointHash = hash; + } + + public Long getApiSecurityEndpointHash() { + return this.apiSecurityEndpointHash; + } + void addRequestHeader(String name, String value) { if (finishedRequestHeaders) { throw new IllegalStateException("Request headers were said to be finished before"); @@ -578,23 +593,33 @@ public String getSessionId() { return sessionId; } + /** + * Close the context and release all resources. This method is idempotent and can be called + * multiple times. For each root span, this method is always called from + * CoreTracer#onRootSpaPublished. + */ @Override public void close() { if (!requestEndCalled) { log.debug(SEND_TELEMETRY, "Request end event was not called before close"); } - if (additive != null) { - log.debug( - SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)"); - closeAdditive(); - } - collectedCookies = null; - requestHeaders.clear(); - responseHeaders.clear(); - persistentData.clear(); - if (derivatives != null) { - derivatives.clear(); - derivatives = null; + // For API Security, we sometimes keep contexts open for late processing. In that case, this + // flag needs to be + // later reset by the API Security post-processor and close must be called again. + if (!keepOpenForApiSecurityPostProcessing) { + if (additive != null) { + log.debug( + SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)"); + closeAdditive(); + } + collectedCookies = null; + requestHeaders.clear(); + responseHeaders.clear(); + persistentData.clear(); + if (derivatives != null) { + derivatives.clear(); + derivatives = null; + } } } @@ -661,6 +686,7 @@ List getStackTraces() { } public void reportDerivatives(Map data) { + log.debug("Reporting derivatives: {}", data); if (data == null || data.isEmpty()) return; if (derivatives == null) { @@ -670,11 +696,13 @@ public void reportDerivatives(Map data) { } } - boolean commitDerivatives(TraceSegment traceSegment) { + public boolean commitDerivatives(TraceSegment traceSegment) { + log.debug("Committing derivatives: {} for {}", derivatives, traceSegment); if (traceSegment == null || derivatives == null) { return false; } derivatives.forEach(traceSegment::setTagTop); + log.debug("Committed derivatives: {} for {}", derivatives, traceSegment); derivatives = null; return true; } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index f6bf2646e8f..a63806d0d92 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -20,7 +20,6 @@ import com.datadog.appsec.event.data.SingletonDataBundle; import com.datadog.appsec.report.AppSecEvent; import com.datadog.appsec.report.AppSecEventWrapper; -import datadog.trace.api.Config; import datadog.trace.api.ProductTraceSource; import datadog.trace.api.gateway.Events; import datadog.trace.api.gateway.Flow; @@ -35,7 +34,6 @@ import datadog.trace.api.telemetry.WafMetricCollector; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; -import datadog.trace.util.NonBlockingSemaphore; import datadog.trace.util.stacktrace.StackTraceEvent; import datadog.trace.util.stacktrace.StackUtils; import java.net.URI; @@ -83,10 +81,6 @@ public class GatewayBridge { private static final String METASTRUCT_EXPLOIT = "exploit"; - private final int MAX_POST_PROCESSING_TASKS = 16; - private final NonBlockingSemaphore postProcessingCounter = - NonBlockingSemaphore.withPermitCount(MAX_POST_PROCESSING_TASKS); - private final SubscriptionService subscriptionService; private final EventProducerService producerService; private final ApiSecurityRequestSampler requestSampler; @@ -156,7 +150,6 @@ public void init() { subscriptionService.registerCallback(EVENTS.shellCmd(), this::onShellCmd); subscriptionService.registerCallback(EVENTS.user(), this::onUser); subscriptionService.registerCallback(EVENTS.loginEvent(), this::onLoginEvent); - subscriptionService.registerCallback(EVENTS.postProcessing(), this::onPostProcessing); if (additionalIGEvents.contains(EVENTS.requestPathParams())) { subscriptionService.registerCallback(EVENTS.requestPathParams(), this::onRequestPathParams); @@ -670,6 +663,12 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { TraceSegment traceSeg = ctx_.getTraceSegment(); Map tags = spanInfo.getTags(); + if (maybeSampleForApiSecurity(ctx, spanInfo, tags)) { + ctx.setKeepOpenForApiSecurityPostProcessing(true); + } else { + ctx.closeAdditive(); + } + // AppSec report metric and events for web span only if (traceSeg != null) { traceSeg.setTagTop("_dd.appsec.enabled", 1); @@ -730,19 +729,21 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { } } - // Route used in post-processing - Object route = tags.get(Tags.HTTP_ROUTE); - if (route instanceof String) { - ctx.setRoute((String) route); - } - if (requestSampler.preSampleRequest(ctx) && postProcessingCounter.acquire()) { - // The request is pre-sampled - we need to post-process it - spanInfo.setRequiresPostProcessing(true); - } - + ctx.close(); return NoopFlow.INSTANCE; } + private boolean maybeSampleForApiSecurity( + AppSecRequestContext ctx, IGSpanInfo spanInfo, Map tags) { + log.debug("Checking API Security for end of request handler on span: {}", spanInfo.getSpanId()); + // API Security sampling requires http.route tag. + final Object route = tags.get(Tags.HTTP_ROUTE); + if (route != null) { + ctx.setRoute(route.toString()); + } + return requestSampler.preSampleRequest(ctx); + } + private Flow onRequestHeadersDone(RequestContext ctx_) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null || ctx.isReqDataPublished()) { @@ -798,19 +799,6 @@ private void onRequestHeader(RequestContext ctx_, String name, String value) { } } - // This handler is executed in a separate thread due the computation possible overhead - private void onPostProcessing(RequestContext ctx_) { - AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); - if (ctx == null) { - return; - } - - maybeExtractSchemas(ctx); - ctx.close(); - // Decrease the counter to allow the next request to be post-processed - postProcessingCounter.release(); - } - public void stop() { subscriptionService.reset(); } @@ -976,40 +964,6 @@ private Flow maybePublishResponseData(AppSecRequestContext ctx) { } } - private void maybeExtractSchemas(AppSecRequestContext ctx) { - boolean extractSchema = false; - if (Config.get().isApiSecurityEnabled() && requestSampler != null) { - extractSchema = requestSampler.sampleRequest(ctx); - } - - if (!extractSchema) { - return; - } - - while (true) { - DataSubscriberInfo subInfo = requestEndSubInfo; - if (subInfo == null) { - subInfo = producerService.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR); - requestEndSubInfo = subInfo; - } - if (subInfo == null || subInfo.isEmpty()) { - return; - } - - DataBundle bundle = - new SingletonDataBundle<>( - KnownAddresses.WAF_CONTEXT_PROCESSOR, - Collections.singletonMap("extract-schema", true)); - try { - GatewayContext gwCtx = new GatewayContext(false); - producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx); - return; - } catch (ExpiredSubscriberInfoException e) { - requestEndSubInfo = null; - } - } - } - private static Map> parseQueryStringParams( String queryString, Charset uriEncoding) { if (queryString == null) { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy deleted file mode 100644 index 783b938fd4b..00000000000 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiAccessTrackerTest.groovy +++ /dev/null @@ -1,55 +0,0 @@ -package com.datadog.appsec.api.security - -import datadog.trace.test.util.DDSpecification - -class ApiAccessTrackerTest extends DDSpecification { - def "should add new api access and update if expired"() { - given: "An ApiAccessTracker with capacity 2 and expiration time 1 second" - def tracker = new ApiAccessTracker(2, 1000) - - when: "Adding new api access" - tracker.updateApiAccessIfExpired("route1", "GET", 200) - def firstAccessTime = tracker.apiAccessMap.values().iterator().next() - - then: "The access is added" - tracker.apiAccessMap.size() == 1 - - when: "Waiting more than expiration time and adding another access with the same key" - Thread.sleep(1100) // Waiting more than 1 second to ensure expiration - tracker.updateApiAccessIfExpired("route1", "GET", 200) - def secondAccessTime = tracker.apiAccessMap.values().iterator().next() - - then: "The access is updated and moved to the end" - tracker.apiAccessMap.size() == 1 - secondAccessTime > firstAccessTime - } - - def "should remove the oldest access when capacity is exceeded"() { - given: "An ApiAccessTracker with capacity 1" - def tracker = new ApiAccessTracker(1, 1000) - - when: "Adding two api accesses" - tracker.updateApiAccessIfExpired("route1", "GET", 200) - Thread.sleep(100) // Delay to ensure different timestamps - tracker.updateApiAccessIfExpired("route2", "POST", 404) - - then: "The oldest access is removed" - tracker.apiAccessMap.size() == 1 - !tracker.apiAccessMap.containsKey(tracker.computeApiHash("route1", "GET", 200)) - tracker.apiAccessMap.containsKey(tracker.computeApiHash("route2", "POST", 404)) - } - - def "should not update access if not expired"() { - given: "An ApiAccessTracker with a short expiration time" - def tracker = new ApiAccessTracker(2, 2000) // 2 seconds expiration - - when: "Adding an api access and updating it before it expires" - tracker.updateApiAccessIfExpired("route1", "GET", 200) - def updateTime = System.currentTimeMillis() - boolean updatedBeforeExpiration = tracker.updateApiAccessIfExpired("route1", "GET", 200) - - then: "The access is not updated" - !updatedBeforeExpiration - tracker.apiAccessMap.get(tracker.computeApiHash("route1", "GET", 200)) == updateTime - } -} \ No newline at end of file diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy index de9fb0a49da..b66402c99ce 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy @@ -1,35 +1,219 @@ package com.datadog.appsec.api.security import com.datadog.appsec.gateway.AppSecRequestContext -import datadog.trace.api.Config +import datadog.trace.api.time.ControllableTimeSource import datadog.trace.test.util.DDSpecification class ApiSecurityRequestSamplerTest extends DDSpecification { - def config = Mock(Config) { - isApiSecurityEnabled() >> true + void 'happy path with single request'() { + given: + def ctx = createContext('route1', 'GET', 200) + def sampler = new ApiSecurityRequestSampler() + + when: + final preSampled = sampler.preSampleRequest(ctx) + + then: + preSampled + + when: + ctx.setKeepOpenForApiSecurityPostProcessing(true) + final sampled = sampler.sampleRequest(ctx) + + then: + sampled } - def sampler = new ApiSecurityRequestSampler(config) + void 'second request is not sampled for the same endpoint'() { + given: + AppSecRequestContext ctx1 = createContext('route1', 'GET', 200) + AppSecRequestContext ctx2 = createContext('route1', 'GET', 200) + def sampler = new ApiSecurityRequestSampler() - void 'Api Security Sample Request'() { when: - def span = Mock(AppSecRequestContext) { - getRoute() >> route - getMethod() >> method - getResponseStatus() >> statusCode + final preSampled1 = sampler.preSampleRequest(ctx1) + ctx1.setKeepOpenForApiSecurityPostProcessing(true) + final sampled1 = sampler.sampleRequest(ctx1) + sampler.releaseOne() + + then: + preSampled1 + sampled1 + + when: + final preSampled2 = sampler.preSampleRequest(ctx2) + + then: + !preSampled2 + } + + void 'preSampleRequest with maximum concurrent contexts'() { + given: + final ctx1 = Spy(createContext('route2', 'GET', 200)) + final ctx2 = Spy(createContext('route3', 'GET', 200)) + final sampler = new ApiSecurityRequestSampler() + assert sampler.MAX_POST_PROCESSING_TASKS > 0 + + when: 'exhaust the maximum number of concurrent contexts' + final List preSampled1 = (1..sampler.MAX_POST_PROCESSING_TASKS).collect { + sampler.preSampleRequest(createContext('route1', 'GET', 200 + it)) } - def sample = sampler.sampleRequest(span) then: - sample == sampleResult + preSampled1.every { it } + + and: 'try to sample one more' + final preSampled2 = sampler.preSampleRequest(ctx1) + + then: + !preSampled2 + + when: 'release one context' + sampler.releaseOne() + + and: 'next can be sampled' + final preSampled3 = sampler.preSampleRequest(ctx2) + + then: + preSampled3 + } + + void 'preSampleRequest with null route'() { + given: + def ctx = createContext(null, 'GET', 200) + def sampler = new ApiSecurityRequestSampler() + + when: + def preSampled = sampler.preSampleRequest(ctx) + + then: + !preSampled + } + + void 'preSampleRequest with null method'() { + given: + def ctx = createContext('route1', null, 200) + def sampler = new ApiSecurityRequestSampler() + + when: + def preSampled = sampler.preSampleRequest(ctx) + + then: + !preSampled + } + + void 'preSampleRequest with 0 status code'() { + given: + def ctx = createContext('route1', 'GET', 0) + def sampler = new ApiSecurityRequestSampler() + + when: + def preSampled = sampler.preSampleRequest(ctx) + + then: + !preSampled + } + + void 'sampleRequest with null context throws'() { + given: + def sampler = new ApiSecurityRequestSampler() + + when: + sampler.preSampleRequest(null) + + then: + thrown(NullPointerException) + } + + void 'sampleRequest without prior preSampleRequest never works'() { + given: + def sampler = new ApiSecurityRequestSampler() + def ctx = createContext('route1', 'GET', 200) + + when: + def sampled = sampler.sampleRequest(ctx) + + then: + !sampled + } + + void 'sampleRequest honors expiration'() { + given: + def ctx = createContext('route1', 'GET', 200) + ctx.setApiSecurityEndpointHash(42L) + ctx.setKeepOpenForApiSecurityPostProcessing(true) + final timeSource = new ControllableTimeSource() + timeSource.set(0) + final long expirationTimeInMs = 10L + final long expirationTimeInNs = expirationTimeInMs * 1_000_000 + def sampler = new ApiSecurityRequestSampler(10, expirationTimeInMs, timeSource) + + when: + def sampled = sampler.sampleRequest(ctx) + + then: + sampled + + when: + sampled = sampler.sampleRequest(ctx) + + then: 'second request is not sampled' + !sampled + + when: 'expiration time has passed' + timeSource.advance(expirationTimeInNs) + sampled = sampler.sampleRequest(ctx) + + then: 'request is sampled again' + sampled + } + + void 'internal accessMap never goes beyond capacity'() { + given: + final timeSource = new ControllableTimeSource() + timeSource.set(0) + final long expirationTimeInMs = 10_000 + final int maxCapacity = 10 + ApiSecurityRequestSampler sampler = new ApiSecurityRequestSampler(10, expirationTimeInMs, timeSource) + + expect: + for (int i = 0; i < maxCapacity * 10; i++) { + timeSource.advance(1_000_000) + final ctx = createContext('route1', 'GET', 200 + 1) + ctx.setApiSecurityEndpointHash(i as long) + ctx.setKeepOpenForApiSecurityPostProcessing(true) + assert sampler.sampleRequest(ctx) + assert sampler.accessMap.size() <= maxCapacity + } + } + + void 'expired entries are purged from internal accessMap'() { + given: + final timeSource = new ControllableTimeSource() + timeSource.set(0) + final long expirationTimeInMs = 10_000 + final int maxCapacity = 10 + ApiSecurityRequestSampler sampler = new ApiSecurityRequestSampler(10, expirationTimeInMs, timeSource) + + expect: + for (int i = 0; i < maxCapacity * 10; i++) { + final ctx = createContext('route1', 'GET', 200 + 1) + ctx.setApiSecurityEndpointHash(i as long) + ctx.setKeepOpenForApiSecurityPostProcessing(true) + assert sampler.sampleRequest(ctx) + assert sampler.accessMap.size() <= 2 + if (i % 2) { + timeSource.advance(expirationTimeInMs * 1_000_000) + } + } + } - where: - method | route | statusCode | sampleResult - 'GET' | 'route1' | 200 | true - 'GET' | 'route2' | null | false - 'GET' | null | 404 | false - 'TOP' | 999 | 404 | true - null | '999' | 404 | false + private static AppSecRequestContext createContext(final String route, final String method, int statusCode) { + final AppSecRequestContext ctx = new AppSecRequestContext() + ctx.setRoute(route) + ctx.setMethod(method) + ctx.setResponseStatus(statusCode) + ctx } -} \ No newline at end of file +} diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy new file mode 100644 index 00000000000..b77c98c353c --- /dev/null +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy @@ -0,0 +1,245 @@ +package com.datadog.appsec.api.security + +import com.datadog.appsec.event.EventProducerService +import com.datadog.appsec.event.ExpiredSubscriberInfoException +import com.datadog.appsec.event.data.KnownAddresses +import com.datadog.appsec.gateway.AppSecRequestContext +import datadog.trace.api.gateway.RequestContext +import datadog.trace.api.internal.TraceSegment +import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import datadog.trace.test.util.DDSpecification + +class AppSecSpanPostProcessorTest extends DDSpecification { + + void 'schema extracted on happy path'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def subInfo = Mock(EventProducerService.DataSubscriberInfo) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def traceSegment = Mock(TraceSegment) + def ctx = Mock(AppSecRequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> ctx + 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true + 1 * sampler.sampleRequest(_) >> true + 1 * reqCtx.getTraceSegment() >> traceSegment + 1 * producer.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR) >> subInfo + 1 * subInfo.isEmpty() >> false + 1 * producer.publishDataEvent(_, ctx, _, _) + 1 * ctx.commitDerivatives(traceSegment) + 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.close() + 1 * sampler.releaseOne() + 0 * _ + } + + void 'no schema extracted if sampling is false'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def ctx = Mock(AppSecRequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> ctx + 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true + 1 * sampler.sampleRequest(_) >> false + 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.close() + 1 * sampler.releaseOne() + 0 * _ + } + + void 'permit is released even if request context close throws'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def traceSegment = Mock(TraceSegment) + def ctx = Mock(AppSecRequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> ctx + 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true + 1 * sampler.sampleRequest(_) >> true + 1 * reqCtx.getTraceSegment() >> traceSegment + 1 * producer.getDataSubscribers(_) >> null + 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.close() >> { throw new RuntimeException() } + 1 * sampler.releaseOne() + 0 * _ + } + + void 'context is cleaned up on timeout'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def ctx = Mock(AppSecRequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { true }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> ctx + 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true + 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.close() + 1 * sampler.releaseOne() + 0 * _ + } + + void 'process null request context does nothing'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def span = Mock(AgentSpan) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> null + 0 * _ + } + + void 'process null appsec request context does nothing'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> null + 0 * _ + } + + void 'process already closed context does nothing'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def ctx = Mock(AppSecRequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> ctx + 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> false + 0 * _ + } + + void 'process throws on null span'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(null, { false }) + + then: + thrown(NullPointerException) + 0 * _ + } + + void 'empty event subscription does not break the process'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def subInfo = Mock(EventProducerService.DataSubscriberInfo) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def traceSegment = Mock(TraceSegment) + def ctx = Mock(AppSecRequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> ctx + 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true + 1 * sampler.sampleRequest(_) >> true + 1 * reqCtx.getTraceSegment() >> traceSegment + 1 * producer.getDataSubscribers(_) >> subInfo + 1 * subInfo.isEmpty() >> true + 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.close() + 1 * sampler.releaseOne() + 0 * _ + } + + void 'expired event subscription does not break the process'() { + given: + def sampler = Mock(ApiSecurityRequestSampler) + def producer = Mock(EventProducerService) + def subInfo = Mock(EventProducerService.DataSubscriberInfo) + def span = Mock(AgentSpan) + def reqCtx = Mock(RequestContext) + def traceSegment = Mock(TraceSegment) + def ctx = Mock(AppSecRequestContext) + def processor = new AppSecSpanPostProcessor(sampler, producer) + + when: + processor.process(span, { false }) + + then: + noExceptionThrown() + 1 * span.getRequestContext() >> reqCtx + 1 * reqCtx.getData(_) >> ctx + 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true + 1 * sampler.sampleRequest(_) >> true + 1 * reqCtx.getTraceSegment() >> traceSegment + 1 * producer.getDataSubscribers(_) >> subInfo + 1 * subInfo.isEmpty() >> false + 1 * producer.publishDataEvent(_, ctx, _, _) >> { throw new ExpiredSubscriberInfoException() } + 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.close() + 1 * sampler.releaseOne() + 0 * _ + } +} diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy index 33f07f611b5..fd622945d6e 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy @@ -14,7 +14,6 @@ import datadog.trace.test.util.DDSpecification import java.nio.file.Files import java.nio.file.Path -import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_ACTIVATION import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE @@ -258,7 +257,6 @@ class AppSecConfigServiceImplSpecification extends DDSpecification { } 1 * poller.addConfigurationEndListener(_) >> { listeners.savedConfEndListener = it[0] } 1 * poller.addCapabilities(CAPABILITY_ASM_ACTIVATION) - 1 * poller.addCapabilities(CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE) 1 * poller.addCapabilities(CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE) 1 * poller.addCapabilities(CAPABILITY_ASM_DD_RULES | CAPABILITY_ASM_IP_BLOCKING @@ -411,7 +409,6 @@ class AppSecConfigServiceImplSpecification extends DDSpecification { } 1 * poller.addConfigurationEndListener(_) >> { listeners.savedConfEndListener = it[0] } 1 * poller.addCapabilities(CAPABILITY_ASM_ACTIVATION) - 1 * poller.addCapabilities(CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE) 1 * poller.addCapabilities(CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE) 1 * poller.addCapabilities(CAPABILITY_ASM_DD_RULES | CAPABILITY_ASM_IP_BLOCKING @@ -495,7 +492,6 @@ class AppSecConfigServiceImplSpecification extends DDSpecification { | CAPABILITY_ASM_CUSTOM_RULES | CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE | CAPABILITY_ASM_TRUSTED_IPS - | CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE | CAPABILITY_ASM_RASP_SQLI | CAPABILITY_ASM_RASP_SSRF | CAPABILITY_ASM_RASP_CMDI diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 70d3703fc55..b8cb74e0992 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -77,9 +77,7 @@ class GatewayBridgeSpecification extends DDSpecification { } TraceSegmentPostProcessor pp = Mock() - ApiSecurityRequestSampler requestSampler = Mock(ApiSecurityRequestSampler) { - preSampleRequest(_ as AppSecRequestContext) >> false - } + ApiSecurityRequestSampler requestSampler = Mock(ApiSecurityRequestSampler) GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, requestSampler, [pp]) Supplier> requestStartedCB @@ -165,6 +163,7 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * spanInfo.getTags() >> ['http.client_ip': '1.1.1.1'] 1 * mockAppSecCtx.transferCollectedEvents() >> [event] 1 * mockAppSecCtx.peerAddress >> '2001::1' + 1 * mockAppSecCtx.close() 1 * traceSegment.setTagTop("_dd.appsec.enabled", 1) 1 * traceSegment.setTagTop("_dd.runtime_family", "jvm") 1 * traceSegment.setTagTop('appsec.event', true) @@ -447,7 +446,6 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * ig.registerCallback(EVENTS.shellCmd(), _) >> { shellCmdCB = it[1]; null } 1 * ig.registerCallback(EVENTS.user(), _) >> { userCB = it[1]; null } 1 * ig.registerCallback(EVENTS.loginEvent(), _) >> { loginEventCB = it[1]; null } - 1 * ig.registerCallback(EVENTS.postProcessing(), _) >> { postProcessingCB = it[1]; null } 0 * ig.registerCallback(_, _) bridge.init() diff --git a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java index 8529a048fbd..a71fa52b689 100644 --- a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java +++ b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java @@ -205,6 +205,11 @@ public String shiCmdParamsAndFile( return "EXECUTED"; } + @GetMapping("/api_security/sampling/{status_code}") + public ResponseEntity apiSecuritySampling(@PathVariable("status_code") int statusCode) { + return ResponseEntity.status(statusCode).body("EXECUTED"); + } + private void withProcess(final Operation op) { Process process = null; try { diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy index e7aaea58024..42768f998b3 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy +++ b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy @@ -6,12 +6,18 @@ import okhttp3.FormBody import okhttp3.MediaType import okhttp3.Request import okhttp3.RequestBody +import okhttp3.Response import spock.lang.Shared import java.nio.charset.StandardCharsets class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest { + @Override + def logLevel() { + 'DEBUG' + } + @Shared String buildDir = new File(System.getProperty("datadog.smoketest.builddir")).absolutePath @Shared @@ -650,7 +656,36 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest { 'cmd' | ['$(cat /etc/passwd 1>&2 ; echo .)'] | null 'cmdWithParams' | ['$(cat /etc/passwd 1>&2 ; echo .)'] | ['param'] 'cmdParamsAndFile' | ['$(cat /etc/passwd 1>&2 ; echo .)'] | ['param'] + } + void 'API Security samples only one request per endpoint'() { + given: + def url = "http://localhost:${httpPort}/api_security/sampling/200?test=value" + def client = OkHttpUtils.clientBuilder().build() + def request = new Request.Builder() + .url(url) + .addHeader('X-My-Header', "value") + .get() + .build() + + when: + List responses = (1..3).collect { + client.newCall(request).execute() + } + + then: + responses.each { + assert it.code() == 200 + } + waitForTraceCount(3) + def spans = rootSpans.toList().toSorted { it.span.duration } + spans.size() == 3 + def sampledSpans = spans.findAll { it.meta.keySet().any { it.startsWith('_dd.appsec.s.req.') } } + sampledSpans.size() == 1 + def span = sampledSpans[0] + span.meta.containsKey('_dd.appsec.s.req.query') + span.meta.containsKey('_dd.appsec.s.req.params') + span.meta.containsKey('_dd.appsec.s.req.headers') } } diff --git a/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy b/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy index ef240f03752..0d8eb04f393 100644 --- a/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy +++ b/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy @@ -44,6 +44,10 @@ abstract class AbstractAppSecServerSmokeTest extends AbstractServerSmokeTest { protected String[] defaultAppSecProperties = [ "-Ddd.appsec.enabled=${System.getProperty('smoke_test.appsec.enabled') ?: 'true'}", "-Ddd.profiling.enabled=false", + // TODO: Remove once this is the default value + "-Ddd.api-security.enabled=true", + "-Ddd.appsec.waf.timeout=300000", + "-DPOWERWAF_EXIT_ON_LEAK=true", // disable AppSec rate limit "-Ddd.appsec.trace.rate.limit=-1" ] + (System.getProperty('smoke_test.appsec.enabled') == 'inactive' ? diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java index 5e620d344e1..172b8f2051b 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java @@ -114,7 +114,7 @@ public final class ConfigDefaults { static final boolean DEFAULT_APPSEC_WAF_METRICS = true; static final int DEFAULT_APPSEC_WAF_TIMEOUT = 100000; // 0.1 s static final boolean DEFAULT_API_SECURITY_ENABLED = false; - static final float DEFAULT_API_SECURITY_REQUEST_SAMPLE_RATE = 0.1f; // 10 % + static final float DEFAULT_API_SECURITY_SAMPLE_DELAY = 30.0f; static final boolean DEFAULT_APPSEC_RASP_ENABLED = true; static final boolean DEFAULT_APPSEC_STACK_TRACE_ENABLED = true; static final int DEFAULT_APPSEC_MAX_STACK_TRACES = 2; diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/AppSecConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/AppSecConfig.java index 454ef0a5425..a7f71b6f664 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/AppSecConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/AppSecConfig.java @@ -26,7 +26,7 @@ public final class AppSecConfig { public static final String API_SECURITY_ENABLED = "api-security.enabled"; public static final String API_SECURITY_ENABLED_EXPERIMENTAL = "experimental.api-security.enabled"; - public static final String API_SECURITY_REQUEST_SAMPLE_RATE = "api-security.request.sample.rate"; + public static final String API_SECURITY_SAMPLE_DELAY = "api-security.sample.delay"; public static final String APPSEC_SCA_ENABLED = "appsec.sca.enabled"; public static final String APPSEC_RASP_ENABLED = "appsec.rasp.enabled"; diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java index f576a3fe953..50b9aba2cdb 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java @@ -14,8 +14,6 @@ import datadog.trace.common.writer.ddagent.DDAgentMapperDiscovery; import datadog.trace.common.writer.ddagent.Prioritization; import datadog.trace.core.monitor.HealthMetrics; -import datadog.trace.core.postprocessor.AppSecSpanPostProcessor; -import datadog.trace.core.postprocessor.SpanPostProcessor; import java.util.concurrent.TimeUnit; import okhttp3.HttpUrl; import okhttp3.OkHttpClient; @@ -155,7 +153,6 @@ public DDAgentWriter build() { final DDAgentMapperDiscovery mapperDiscovery = new DDAgentMapperDiscovery(featureDiscovery); final PayloadDispatcher dispatcher = new PayloadDispatcherImpl(mapperDiscovery, agentApi, healthMetrics, monitoring); - final SpanPostProcessor spanPostProcessor = new AppSecSpanPostProcessor(); final TraceProcessingWorker traceProcessingWorker = new TraceProcessingWorker( traceBufferSize, @@ -165,8 +162,7 @@ public DDAgentWriter build() { null == prioritization ? FAST_LANE : prioritization, flushIntervalMilliseconds, TimeUnit.MILLISECONDS, - singleSpanSampler, - spanPostProcessor); + singleSpanSampler); return new DDAgentWriter( traceProcessingWorker, diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java index fd26f4db7eb..f298a5e8f6a 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/DDIntakeWriter.java @@ -9,8 +9,6 @@ import datadog.trace.common.writer.ddagent.Prioritization; import datadog.trace.common.writer.ddintake.DDIntakeMapperDiscovery; import datadog.trace.core.monitor.HealthMetrics; -import datadog.trace.core.postprocessor.AppSecSpanPostProcessor; -import datadog.trace.core.postprocessor.SpanPostProcessor; import java.util.EnumMap; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -115,8 +113,6 @@ public DDIntakeWriter build() { dispatcher = new CompositePayloadDispatcher(dispatchers); } - SpanPostProcessor spanPostProcessor = new AppSecSpanPostProcessor(); - final TraceProcessingWorker traceProcessingWorker = new TraceProcessingWorker( traceBufferSize, @@ -126,8 +122,7 @@ public DDIntakeWriter build() { prioritization, flushIntervalMilliseconds, TimeUnit.MILLISECONDS, - singleSpanSampler, - spanPostProcessor); + singleSpanSampler); return new DDIntakeWriter( traceProcessingWorker, diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java index d2ed59f3f8a..6cd0ecdaed2 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java @@ -7,6 +7,7 @@ import datadog.communication.ddagent.DroppingPolicy; import datadog.trace.api.Config; +import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor; import datadog.trace.common.sampling.SingleSpanSampler; import datadog.trace.common.writer.ddagent.FlushEvent; import datadog.trace.common.writer.ddagent.Prioritization; @@ -14,8 +15,6 @@ import datadog.trace.core.CoreSpan; import datadog.trace.core.DDSpan; import datadog.trace.core.monitor.HealthMetrics; -import datadog.trace.core.postprocessor.SpanPostProcessor; -import java.util.ArrayList; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -53,8 +52,7 @@ public TraceProcessingWorker( final Prioritization prioritization, final long flushInterval, final TimeUnit timeUnit, - final SingleSpanSampler singleSpanSampler, - final SpanPostProcessor spanPostProcessor) { + final SingleSpanSampler singleSpanSampler) { this.capacity = capacity; this.primaryQueue = createQueue(capacity); this.secondaryQueue = createQueue(capacity); @@ -75,13 +73,7 @@ public TraceProcessingWorker( this.serializingHandler = new TraceSerializingHandler( - primaryQueue, - secondaryQueue, - healthMetrics, - dispatcher, - flushInterval, - timeUnit, - spanPostProcessor); + primaryQueue, secondaryQueue, healthMetrics, dispatcher, flushInterval, timeUnit); this.serializerThread = newAgentThread(TRACE_PROCESSOR, serializingHandler); } @@ -142,7 +134,6 @@ public static class TraceSerializingHandler implements Runnable { private final boolean doTimeFlush; private final PayloadDispatcher payloadDispatcher; private long lastTicks; - private final SpanPostProcessor spanPostProcessor; public TraceSerializingHandler( final MpscBlockingConsumerArrayQueue primaryQueue, @@ -150,8 +141,7 @@ public TraceSerializingHandler( final HealthMetrics healthMetrics, final PayloadDispatcher payloadDispatcher, final long flushInterval, - final TimeUnit timeUnit, - final SpanPostProcessor spanPostProcessor) { + final TimeUnit timeUnit) { this.primaryQueue = primaryQueue; this.secondaryQueue = secondaryQueue; this.healthMetrics = healthMetrics; @@ -163,7 +153,6 @@ public TraceSerializingHandler( } else { this.ticksRequiredToFlush = Long.MAX_VALUE; } - this.spanPostProcessor = spanPostProcessor; } @Override @@ -262,36 +251,26 @@ private void maybeTracePostProcessing(List trace) { return; } - // Filter spans that need post-processing - List spansToPostProcess = null; - for (DDSpan span : trace) { - if (span.isRequiresPostProcessing()) { - if (spansToPostProcess == null) { - spansToPostProcess = new ArrayList<>(); - } - spansToPostProcess.add(span); - } - } - - if (spansToPostProcess == null) { - return; - } - + final SpanPostProcessor postProcessor = SpanPostProcessor.Holder.INSTANCE; try { - long timeout = Config.get().getTracePostProcessingTimeout(); - long deadline = System.currentTimeMillis() + timeout; - BooleanSupplier timeoutCheck = () -> System.currentTimeMillis() > deadline; - - for (DDSpan span : spansToPostProcess) { - if (!spanPostProcessor.process(span, timeoutCheck)) { - log.debug("Span post-processing interrupted due to timeout."); - break; - } + final long timeout = Config.get().getTracePostProcessingTimeout(); + final long deadline = System.currentTimeMillis() + timeout; + final boolean[] timedOut = {false}; + final BooleanSupplier timeoutCheck = + () -> { + if (timedOut[0]) { + return true; + } + if (System.currentTimeMillis() > deadline) { + timedOut[0] = true; + } + return timedOut[0]; + }; + for (DDSpan span : trace) { + postProcessor.process(span, timeoutCheck); } } catch (Throwable e) { - if (log.isDebugEnabled()) { - log.debug("Error while trace post-processing", e); - } + log.debug("Error while trace post-processing", e); } } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index cfe1d2c79bb..5d033f13c78 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -277,9 +277,7 @@ void onRootSpanPublished(final AgentSpan root) { RequestContext requestContext = root.getRequestContext(); if (requestContext != null) { try { - if (!root.isRequiresPostProcessing()) { - requestContext.close(); - } + requestContext.close(); } catch (IOException e) { log.warn("Error closing request context data", e); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java index 771d22e3758..19f3852a0a6 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java @@ -825,16 +825,6 @@ public void addLink(AgentSpanLink link) { } } - @Override - public boolean isRequiresPostProcessing() { - return context.isRequiresPostProcessing(); - } - - @Override - public void setRequiresPostProcessing(boolean requiresPostProcessing) { - context.setRequiresPostProcessing(requiresPostProcessing); - } - // to be accessible in Spock spies, which the field wouldn't otherwise be public long getStartTimeNano() { return startTimeNano; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java index e6980e2b41d..70b96f2ff55 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java @@ -142,7 +142,6 @@ public class DDSpanContext private final boolean injectBaggageAsTags; private volatile int encodedOperationName; private volatile int encodedResourceName; - private volatile boolean requiresPostProcessing; private volatile CharSequence lastParentId; private final boolean isRemote; @@ -1044,14 +1043,6 @@ public void setMetaStructCurrent(String field, Object value) { setMetaStruct(field, value); } - public void setRequiresPostProcessing(boolean postProcessing) { - this.requiresPostProcessing = postProcessing; - } - - public boolean isRequiresPostProcessing() { - return requiresPostProcessing; - } - public CharSequence getLastParentId() { return lastParentId; } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java b/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java deleted file mode 100644 index ab4a4ea40ff..00000000000 --- a/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/AppSecSpanPostProcessor.java +++ /dev/null @@ -1,40 +0,0 @@ -package datadog.trace.core.postprocessor; - -import static datadog.trace.api.gateway.Events.EVENTS; - -import datadog.trace.api.gateway.CallbackProvider; -import datadog.trace.api.gateway.RequestContext; -import datadog.trace.api.gateway.RequestContextSlot; -import datadog.trace.bootstrap.instrumentation.api.AgentTracer; -import datadog.trace.core.DDSpan; -import java.util.function.BooleanSupplier; -import java.util.function.Consumer; - -public class AppSecSpanPostProcessor implements SpanPostProcessor { - - // For testing purpose - protected AgentTracer.TracerAPI tracer() { - return AgentTracer.get(); - } - - @Override - public boolean process(DDSpan span, BooleanSupplier timeoutCheck) { - CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC); - if (cbp == null) { - return false; - } - - RequestContext ctx = span.getRequestContext(); - if (ctx == null) { - return false; - } - - Consumer postProcessingCallback = cbp.getCallback(EVENTS.postProcessing()); - if (postProcessingCallback == null) { - return false; - } - - postProcessingCallback.accept(ctx); - return true; - } -} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/SpanPostProcessor.java b/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/SpanPostProcessor.java deleted file mode 100644 index 8751959456f..00000000000 --- a/dd-trace-core/src/main/java/datadog/trace/core/postprocessor/SpanPostProcessor.java +++ /dev/null @@ -1,26 +0,0 @@ -package datadog.trace.core.postprocessor; - -import datadog.trace.core.DDSpan; -import java.util.function.BooleanSupplier; - -/** - * Span Post-processing with a timeout check capability. - * - *

Implementations of this interface should carry out post-processing of spans while supporting - * interruption when a specified time limit is exceeded. The method {@code process} must check the - * state of {@code timeoutCheck} while processing span. If {@code timeoutCheck.getAsBoolean()} - * returns {@code true}, processing should be immediately halted, and the method should return - * {@code false}. If post-processing completes successfully before the timeout, the method should - * return {@code true}. - */ -public interface SpanPostProcessor { - /** - * Post-processes a span. - * - * @param span The span to be post-processed. - * @param timeoutCheck A timeout check returning {@code true} if the allotted time has elapsed. - * @return {@code true} if the span was successfully processed; {@code false} in case of a - * timeout. - */ - boolean process(DDSpan span, BooleanSupplier timeoutCheck); -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/writer/TraceProcessingWorkerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/writer/TraceProcessingWorkerTest.groovy index eeef4121580..2a37a69ec64 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/common/writer/TraceProcessingWorkerTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/common/writer/TraceProcessingWorkerTest.groovy @@ -7,7 +7,7 @@ import datadog.trace.core.DDSpan import datadog.trace.core.DDSpanContext import datadog.trace.core.PendingTrace import datadog.trace.core.monitor.HealthMetrics -import datadog.trace.core.postprocessor.SpanPostProcessor +import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor import datadog.trace.test.util.DDSpecification import spock.util.concurrent.PollingConditions @@ -44,7 +44,6 @@ class TraceProcessingWorkerTest extends DDSpecification { FAST_LANE, 1, TimeUnit.NANOSECONDS, - null, null ) // stop heartbeats from being throttled @@ -71,7 +70,6 @@ class TraceProcessingWorkerTest extends DDSpecification { FAST_LANE, 1, TimeUnit.NANOSECONDS, - null, null ) // stop heartbeats from being throttled def timeConditions = new PollingConditions(timeout: 1, initialDelay: 1, factor: 1.25) @@ -97,7 +95,7 @@ class TraceProcessingWorkerTest extends DDSpecification { false }, FAST_LANE, - 100, TimeUnit.SECONDS, null, null) // prevent heartbeats from helping the flush happen + 100, TimeUnit.SECONDS, null) // prevent heartbeats from helping the flush happen when: "there is pending work it is completed before a flush" // processing this span will throw an exception, but it should be caught @@ -136,7 +134,7 @@ class TraceProcessingWorkerTest extends DDSpecification { throwingDispatcher, { false }, FAST_LANE, - 100, TimeUnit.SECONDS, null, null) // prevent heartbeats from helping the flush happen + 100, TimeUnit.SECONDS, null) // prevent heartbeats from helping the flush happen worker.start() when: "a trace is processed but can't be passed on" @@ -163,9 +161,7 @@ class TraceProcessingWorkerTest extends DDSpecification { } HealthMetrics healthMetrics = Mock(HealthMetrics) - // Span 1 - should be post-processed def span1 = DDSpan.create("test", 0, Mock(DDSpanContext) { - isRequiresPostProcessing() >> true getTraceCollector() >> Mock(PendingTrace) { getCurrentTimeNano() >> 0 } @@ -174,14 +170,13 @@ class TraceProcessingWorkerTest extends DDSpecification { // Span 2 - should NOT be post-processed def span2 = DDSpan.create("test", 0, Mock(DDSpanContext) { - isRequiresPostProcessing() >> false getTraceCollector() >> Mock(PendingTrace) { getCurrentTimeNano() >> 0 } }, []) def processedSpan2 = false - SpanPostProcessor spanPostProcessor = Mock(SpanPostProcessor) { + SpanPostProcessor.Holder.INSTANCE = Mock(SpanPostProcessor) { process(span1, _) >> { processedSpan1 = true } process(span2, _) >> { processedSpan2 = true } } @@ -189,7 +184,7 @@ class TraceProcessingWorkerTest extends DDSpecification { TraceProcessingWorker worker = new TraceProcessingWorker(10, healthMetrics, countingDispatcher, { false - }, FAST_LANE, 100, TimeUnit.SECONDS, null, spanPostProcessor) + }, FAST_LANE, 100, TimeUnit.SECONDS, null) worker.start() when: "traces are submitted" @@ -198,11 +193,12 @@ class TraceProcessingWorkerTest extends DDSpecification { then: "traces are passed through unless rejected on submission" conditions.eventually { - assert processedSpan1 == true - assert processedSpan2 == false + assert processedSpan1 + assert processedSpan2 } cleanup: + SpanPostProcessor.Holder.INSTANCE = SpanPostProcessor.Holder.NOOP worker.close() } @@ -217,7 +213,7 @@ class TraceProcessingWorkerTest extends DDSpecification { TraceProcessingWorker worker = new TraceProcessingWorker(10, healthMetrics, countingDispatcher, { false - }, FAST_LANE, 100, TimeUnit.SECONDS, null, null) + }, FAST_LANE, 100, TimeUnit.SECONDS, null) // prevent heartbeats from helping the flush happen worker.start() @@ -268,7 +264,7 @@ class TraceProcessingWorkerTest extends DDSpecification { TraceProcessingWorker worker = new TraceProcessingWorker(10, healthMetrics, countingDispatcher, { false - }, FAST_LANE, 100, TimeUnit.SECONDS, null, null) + }, FAST_LANE, 100, TimeUnit.SECONDS, null) worker.start() worker.close() int queueSize = 0 @@ -305,7 +301,7 @@ class TraceProcessingWorkerTest extends DDSpecification { return false } } - TraceProcessingWorker worker = new TraceProcessingWorker(10, healthMetrics, countingDispatcher, { true }, FAST_LANE, 100, TimeUnit.SECONDS, singleSpanSampler, null) + TraceProcessingWorker worker = new TraceProcessingWorker(10, healthMetrics, countingDispatcher, { true }, FAST_LANE, 100, TimeUnit.SECONDS, singleSpanSampler) worker.start() when: "traces are submitted" @@ -381,7 +377,7 @@ class TraceProcessingWorkerTest extends DDSpecification { return false } } - TraceProcessingWorker worker = new TraceProcessingWorker(10, healthMetrics, countingDispatcher, { false }, FAST_LANE, 100, TimeUnit.SECONDS, singleSpanSampler, null) + TraceProcessingWorker worker = new TraceProcessingWorker(10, healthMetrics, countingDispatcher, { false }, FAST_LANE, 100, TimeUnit.SECONDS, singleSpanSampler) worker.start() when: "traces are submitted" diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy deleted file mode 100644 index 17fc937d0e6..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/postprocessor/AppSecSpanPostProcessorTest.groovy +++ /dev/null @@ -1,119 +0,0 @@ -package datadog.trace.core.postprocessor - -import datadog.trace.api.gateway.CallbackProvider -import datadog.trace.api.gateway.RequestContext -import datadog.trace.api.gateway.RequestContextSlot -import datadog.trace.bootstrap.instrumentation.api.AgentTracer -import datadog.trace.core.DDSpan -import datadog.trace.core.DDSpanContext -import datadog.trace.core.PendingTrace -import datadog.trace.test.util.DDSpecification - -import java.util.function.Consumer -import java.util.function.BooleanSupplier -import static datadog.trace.api.gateway.Events.EVENTS - -class AppSecSpanPostProcessorTest extends DDSpecification { - def "process returns false if span context is null"() { - given: - def processor = new AppSecSpanPostProcessor() - def span = Mock(DDSpan) - def timeoutCheck = Mock(BooleanSupplier) - (span.context()) >> null - - expect: - !processor.process(span, timeoutCheck) - } - - def "process returns false if callback provider is null"() { - given: - AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) - tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> null - def processor = new AppSecSpanPostProcessor() { - @Override - protected AgentTracer.TracerAPI tracer() { - return tracer - } - } - def span = Mock(DDSpan) { - context() >> Mock(DDSpanContext) - } - def timeoutCheck = Mock(BooleanSupplier) - - expect: - !processor.process(span, timeoutCheck) - } - - def "process returns false if request context is null"() { - given: - AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) - def cbp = Mock(CallbackProvider) - tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> cbp - def processor = new AppSecSpanPostProcessor() { - @Override - protected AgentTracer.TracerAPI tracer() { - return tracer - } - } - def span = Mock(DDSpan) { - context() >> Mock(DDSpanContext) - getRequestContext() >> null - } - def timeoutCheck = Mock(BooleanSupplier) - - expect: - !processor.process(span, timeoutCheck) - } - - def "process returns false if post-processing callback is null"() { - given: - AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) - def cbp = Mock(CallbackProvider) - tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> cbp - cbp.getCallback(EVENTS.postProcessing()) >> null - def processor = new AppSecSpanPostProcessor() { - @Override - protected AgentTracer.TracerAPI tracer() { - return tracer - } - } - def span = Mock(DDSpan) { - context() >> Mock(DDSpanContext) - getRequestContext() >> Mock(RequestContext) - } - def timeoutCheck = Mock(BooleanSupplier) - - expect: - !processor.process(span, timeoutCheck) - } - - def "process returns true when all components are properly configured"() { - given: - def callback = Mock(Consumer) - AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) - def cbp = Mock(CallbackProvider) - tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> cbp - cbp.getCallback(EVENTS.postProcessing()) >> callback - def processor = new AppSecSpanPostProcessor() { - @Override - protected AgentTracer.TracerAPI tracer() { - return tracer - } - } - def span = DDSpan.create("test", 0, Mock(DDSpanContext) { - isRequiresPostProcessing() >> true - getTraceCollector() >> Mock(PendingTrace) { - getCurrentTimeNano() >> 0 - } - getRequestContext() >> Mock(RequestContext) - }, []) - def timeoutCheck = Mock(BooleanSupplier) - - when: - boolean result = processor.process(span, timeoutCheck) - - then: - result - 1 * callback.accept(_) - } -} diff --git a/internal-api/build.gradle b/internal-api/build.gradle index 9d0107ae991..0cd1ad54a37 100644 --- a/internal-api/build.gradle +++ b/internal-api/build.gradle @@ -213,7 +213,10 @@ excludedClassesCoverage += [ 'datadog.trace.util.stacktrace.StackTraceFrame', 'datadog.trace.api.iast.VulnerabilityMarks', 'datadog.trace.api.iast.securitycontrol.SecurityControlHelper', - 'datadog.trace.api.iast.securitycontrol.SecurityControl' + 'datadog.trace.api.iast.securitycontrol.SecurityControl', + // Trivial holder and no-op + 'datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor.Holder', + 'datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor.NoOpSpanPostProcessor', ] excludedClassesBranchCoverage = [ 'datadog.trace.api.ProductActivationConfig', diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index 6a0745c6745..eb4091e1156 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -292,7 +292,7 @@ public static String getHostName() { private final int appSecMaxStackTraces; private final int appSecMaxStackTraceDepth; private final boolean apiSecurityEnabled; - private final float apiSecurityRequestSampleRate; + private final float apiSecuritySampleDelay; private final IastDetectionMode iastDetectionMode; private final int iastMaxConcurrentRequests; @@ -1385,9 +1385,8 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment()) apiSecurityEnabled = configProvider.getBoolean( API_SECURITY_ENABLED, DEFAULT_API_SECURITY_ENABLED, API_SECURITY_ENABLED_EXPERIMENTAL); - apiSecurityRequestSampleRate = - configProvider.getFloat( - API_SECURITY_REQUEST_SAMPLE_RATE, DEFAULT_API_SECURITY_REQUEST_SAMPLE_RATE); + apiSecuritySampleDelay = + configProvider.getFloat(API_SECURITY_SAMPLE_DELAY, DEFAULT_API_SECURITY_SAMPLE_DELAY); iastDebugEnabled = configProvider.getBoolean(IAST_DEBUG_ENABLED, DEFAULT_IAST_DEBUG_ENABLED); @@ -2763,8 +2762,8 @@ public boolean isApiSecurityEnabled() { return apiSecurityEnabled; } - public float getApiSecurityRequestSampleRate() { - return apiSecurityRequestSampleRate; + public float getApiSecuritySampleDelay() { + return apiSecuritySampleDelay; } public ProductActivation getIastActivation() { @@ -4799,8 +4798,6 @@ public String toString() { + appSecHttpBlockedTemplateJson + ", apiSecurityEnabled=" + apiSecurityEnabled - + ", apiSecurityRequestSampleRate=" - + apiSecurityRequestSampleRate + ", cwsEnabled=" + cwsEnabled + ", cwsTlsRefresh=" diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java index 6f63ddf7049..d840cad01c3 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java @@ -9,7 +9,6 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import java.util.function.BiFunction; -import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; @@ -313,18 +312,6 @@ public EventType>> shellCmd() { return (EventType>>) SHELL_CMD; } - static final int POST_PROCESSING_ID = 26; - - @SuppressWarnings("rawtypes") - private static final EventType POST_PROCESSING = - new ET<>("trace.post.processing", POST_PROCESSING_ID); - - /** The span post-processing */ - @SuppressWarnings("unchecked") - public EventType> postProcessing() { - return (EventType>) POST_PROCESSING; - } - static final int MAX_EVENTS = nextId.get(); private static final class ET extends EventType { diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java b/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java index 28ddbe27ad3..60bca47f184 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java @@ -16,8 +16,4 @@ public interface IGSpanInfo { void setRequestBlockingAction(Flow.Action.RequestBlockingAction rba); Flow.Action.RequestBlockingAction getRequestBlockingAction(); - - boolean isRequiresPostProcessing(); - - void setRequiresPostProcessing(boolean requiresPostProcessing); } diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java b/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java index e770fcd1952..cd5219f1dfc 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java @@ -10,7 +10,6 @@ import static datadog.trace.api.gateway.Events.LOGIN_EVENT_ID; import static datadog.trace.api.gateway.Events.MAX_EVENTS; import static datadog.trace.api.gateway.Events.NETWORK_CONNECTION_ID; -import static datadog.trace.api.gateway.Events.POST_PROCESSING_ID; import static datadog.trace.api.gateway.Events.REQUEST_BODY_CONVERTED_ID; import static datadog.trace.api.gateway.Events.REQUEST_BODY_DONE_ID; import static datadog.trace.api.gateway.Events.REQUEST_BODY_START_ID; @@ -38,7 +37,6 @@ import java.util.concurrent.atomic.AtomicReferenceArray; import java.util.function.BiConsumer; import java.util.function.BiFunction; -import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; import org.slf4j.Logger; @@ -461,18 +459,6 @@ public Flow apply(RequestContext ctx, String[] arg) { } } }; - case POST_PROCESSING_ID: - return (C) - new Consumer() { - @Override - public void accept(RequestContext ctx) { - try { - ((Consumer) callback).accept(ctx); - } catch (Throwable t) { - log.warn("Callback for {} threw.", eventType, t); - } - } - }; default: log.warn("Unwrapped callback for {}", eventType); return callback; diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ExtractedSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ExtractedSpan.java index 84b9d55ecd1..59bf633b834 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ExtractedSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ExtractedSpan.java @@ -159,11 +159,6 @@ public RequestBlockingAction getRequestBlockingAction() { return null; } - @Override - public boolean isRequiresPostProcessing() { - return false; - } - @Override public String toString() { return "ExtractedSpan{spanContext=" + this.spanContext + '}'; diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ImmutableSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ImmutableSpan.java index 620d51b57a1..c7a63cdfccf 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ImmutableSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ImmutableSpan.java @@ -187,7 +187,4 @@ public AgentSpan setMetaStruct(String field, Object value) { @Override public void setRequestBlockingAction(RequestBlockingAction rba) {} - - @Override - public void setRequiresPostProcessing(boolean requiresPostProcessing) {} } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NoopSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NoopSpan.java index 5d476c39b20..a4ea00f23de 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NoopSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NoopSpan.java @@ -139,9 +139,4 @@ public TraceConfig traceConfig() { public boolean isOutbound() { return false; } - - @Override - public boolean isRequiresPostProcessing() { - return false; - } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanPostProcessor.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanPostProcessor.java new file mode 100644 index 00000000000..137ddedba1d --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanPostProcessor.java @@ -0,0 +1,37 @@ +package datadog.trace.bootstrap.instrumentation.api; + +import java.util.function.BooleanSupplier; +import javax.annotation.Nonnull; + +/** + * Applies post-processing of spans before serialization. + * + *

Post-processing runs in TraceProcessingWorker thread. This provides the following properties: + *

  • Runs in a single thread. Post-processing for each span runs sequentially. + *
  • Runs after the request end, and does not block the main thread. + *
  • Runs at a point where the sampler decision is already available. + */ +public interface SpanPostProcessor { + + /** + * Post-processes a span, if needed. + * + *

    Implementations should use {@code timeoutCheck}, and if true, they should halt processing as + * much as possible. This method is guaranteed to be called even if post-processing of previous + * spans have timed out. + */ + void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutCheck); + + class Holder { + public static final SpanPostProcessor NOOP = new NoOpSpanPostProcessor(); + + // XXX: At the moment, a single post-processor can be registered, and only AppSec defines one. + // If other products add their own, we'll need to refactor this to support multiple processors. + public static volatile SpanPostProcessor INSTANCE = NOOP; + } + + class NoOpSpanPostProcessor implements SpanPostProcessor { + @Override + public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutCheck) {} + } +} diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java index b5ab581c479..8f2840d4308 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java @@ -13,7 +13,6 @@ import java.util.Collections; import java.util.function.BiConsumer; import java.util.function.BiFunction; -import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; import org.assertj.core.api.ThrowableAssert; @@ -224,8 +223,6 @@ public void testNormalCalls() { cbp.getCallback(events.execCmd()).apply(null, null); ss.registerCallback(events.shellCmd(), callback); cbp.getCallback(events.shellCmd()).apply(null, null); - ss.registerCallback(events.postProcessing(), callback); - cbp.getCallback(events.postProcessing()).accept(null); assertThat(callback.count).isEqualTo(Events.MAX_EVENTS); } @@ -296,8 +293,6 @@ public void testThrowableBlocking() { cbp.getCallback(events.execCmd()).apply(null, null); ss.registerCallback(events.shellCmd(), throwback); cbp.getCallback(events.shellCmd()).apply(null, null); - ss.registerCallback(events.postProcessing(), throwback); - cbp.getCallback(events.postProcessing()).accept(null); assertThat(throwback.count).isEqualTo(Events.MAX_EVENTS); } @@ -443,7 +438,6 @@ public void mergeFlowReturnsLatestNonNullResult() { private static class Callback implements Supplier>, - Consumer, BiConsumer, TriConsumer, BiFunction>, @@ -476,11 +470,6 @@ public Flow get() { return new Flow.ResultFlow<>((D) ctxt); } - @Override - public void accept(RequestContext requestContext) { - count++; - } - @Override public void accept(RequestContext requestContext, T s, T s2) { count++; @@ -521,7 +510,6 @@ public Flow apply(RequestContext requestContext, T t, T t2) { private static class Throwback implements Supplier>, - Consumer, BiConsumer, TriConsumer, BiFunction>, @@ -547,12 +535,6 @@ public Flow get() { throw new IllegalArgumentException(); } - @Override - public void accept(RequestContext requestContext) { - count++; - throw new IllegalArgumentException(); - } - @Override public void accept(RequestContext requestContext, T s, T s2) { count++; From d4b08e310e598a67669c0bf172f2c676e36d27d5 Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Mon, 3 Mar 2025 15:26:52 +0100 Subject: [PATCH 07/15] ST commit --- .circleci/config.continue.yml.j2 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.continue.yml.j2 b/.circleci/config.continue.yml.j2 index b0d6de96190..0702fce232d 100644 --- a/.circleci/config.continue.yml.j2 +++ b/.circleci/config.continue.yml.j2 @@ -36,7 +36,7 @@ instrumentation_modules: &instrumentation_modules "dd-java-agent/instrumentation debugger_modules: &debugger_modules "dd-java-agent/agent-debugger|dd-java-agent/agent-bootstrap|dd-java-agent/agent-builder|internal-api|communication|dd-trace-core" profiling_modules: &profiling_modules "dd-java-agent/agent-profiling" -default_system_tests_commit: &default_system_tests_commit 69a5e874384dd256e2e3f42fc1c95807a67efbe6 +default_system_tests_commit: &default_system_tests_commit 1ef00a34ad1f83ae999887e510ef1ea1c27b151b parameters: nightly: From 5e32bf8e7682f32838ad16c5c1f8e8bc30aa8a26 Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Mon, 3 Mar 2025 17:01:34 +0100 Subject: [PATCH 08/15] wip --- dd-java-agent/appsec/build.gradle | 2 +- .../java/com/datadog/appsec/AppSecSystem.java | 10 +++--- .../api/security/ApiSecuritySampler.java | 35 +++++++++++++++++++ ...mpler.java => ApiSecuritySamplerImpl.java} | 34 ++++-------------- .../api/security/AppSecSpanPostProcessor.java | 5 ++- .../datadog/appsec/gateway/GatewayBridge.java | 6 ++-- ...t.groovy => ApiSecuritySamplerTest.groovy} | 26 +++++++------- .../AppSecSpanPostProcessorTest.groovy | 20 +++++------ .../gateway/GatewayBridgeSpecification.groovy | 5 ++- 9 files changed, 79 insertions(+), 64 deletions(-) create mode 100644 dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySampler.java rename dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/{ApiSecurityRequestSampler.java => ApiSecuritySamplerImpl.java} (85%) rename dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/{ApiSecurityRequestSamplerTest.groovy => ApiSecuritySamplerTest.groovy} (86%) diff --git a/dd-java-agent/appsec/build.gradle b/dd-java-agent/appsec/build.gradle index de195400063..6d9856e93d6 100644 --- a/dd-java-agent/appsec/build.gradle +++ b/dd-java-agent/appsec/build.gradle @@ -82,7 +82,7 @@ ext { 'com.datadog.appsec.config.AppSecFeatures.ApiSecurity', 'com.datadog.appsec.config.AppSecFeatures.AutoUserInstrum', 'com.datadog.appsec.event.ReplaceableEventProducerService', - 'com.datadog.appsec.api.security.ApiSecurityRequestSampler.NoOp', + 'com.datadog.appsec.api.security.ApiSecuritySampler.NoOp', ] excludedClassesBranchCoverage = [ 'com.datadog.appsec.gateway.GatewayBridge', diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java index 9900268ba73..e775989a83b 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java @@ -1,6 +1,7 @@ package com.datadog.appsec; -import com.datadog.appsec.api.security.ApiSecurityRequestSampler; +import com.datadog.appsec.api.security.ApiSecuritySampler; +import com.datadog.appsec.api.security.ApiSecuritySamplerImpl; import com.datadog.appsec.api.security.AppSecSpanPostProcessor; import com.datadog.appsec.blocking.BlockingServiceImpl; import com.datadog.appsec.config.AppSecConfigService; @@ -68,13 +69,14 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s EventDispatcher eventDispatcher = new EventDispatcher(); REPLACEABLE_EVENT_PRODUCER.replaceEventProducerService(eventDispatcher); - ApiSecurityRequestSampler requestSampler; + ApiSecuritySampler requestSampler; if (Config.get().isApiSecurityEnabled()) { - requestSampler = new ApiSecurityRequestSampler(); + // TODO: Address support for 1-click enablement + requestSampler = new ApiSecuritySamplerImpl(); SpanPostProcessor.Holder.INSTANCE = new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER); } else { - requestSampler = new ApiSecurityRequestSampler.NoOp(); + requestSampler = new ApiSecuritySampler.NoOp(); } ConfigurationPoller configurationPoller = sco.configurationPoller(config); diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySampler.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySampler.java new file mode 100644 index 00000000000..4412a5d6303 --- /dev/null +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySampler.java @@ -0,0 +1,35 @@ +package com.datadog.appsec.api.security; + +import com.datadog.appsec.gateway.AppSecRequestContext; +import javax.annotation.Nonnull; + +public interface ApiSecuritySampler { + /** + * Prepare a request context for later sampling decision. This method should be called at request + * end, and is thread-safe. If a request can potentially be sampled, this method will return true. + * If this method returns true, the caller MUST call {@link #releaseOne()} once the context is not + * needed anymore. + */ + boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx); + + /** Get the final sampling decision. This method is NOT required to be thread-safe. */ + boolean sampleRequest(AppSecRequestContext ctx); + + /** Release one permit for the sampler. This must be called after processing a span. */ + void releaseOne(); + + final class NoOp implements ApiSecuritySampler { + @Override + public boolean preSampleRequest(@Nonnull AppSecRequestContext ctx) { + return false; + } + + @Override + public boolean sampleRequest(AppSecRequestContext ctx) { + return false; + } + + @Override + public void releaseOne() {} + } +} diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java similarity index 85% rename from dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java rename to dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java index 56c566c3937..21c13ebf585 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java @@ -12,9 +12,9 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -public class ApiSecurityRequestSampler { +public class ApiSecuritySamplerImpl implements ApiSecuritySampler { - private static final Logger log = LoggerFactory.getLogger(ApiSecurityRequestSampler.class); + private static final Logger log = LoggerFactory.getLogger(ApiSecuritySamplerImpl.class); /** * A maximum number of request contexts we'll keep open past the end of request at any given time. @@ -34,14 +34,14 @@ public class ApiSecurityRequestSampler { private final TimeSource timeSource; private final Semaphore counter = new Semaphore(MAX_POST_PROCESSING_TASKS); - public ApiSecurityRequestSampler() { + public ApiSecuritySamplerImpl() { this( MAX_SIZE, (long) (Config.get().getApiSecuritySampleDelay() * 1_000), SystemTimeSource.INSTANCE); } - public ApiSecurityRequestSampler( + public ApiSecuritySamplerImpl( int capacity, long expirationTimeInMs, @Nonnull TimeSource timeSource) { this.capacity = capacity; this.expirationTimeInMs = expirationTimeInMs; @@ -50,12 +50,7 @@ public ApiSecurityRequestSampler( this.timeSource = timeSource; } - /** - * Prepare a request context for later sampling decision. This method should be called at request - * end, and is thread-safe. If a request can potentially be sampled, this method will return true. - * If this method returns true, the caller MUST call {@link #releaseOne()} once the context is not - * needed anymore. - */ + @Override public boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx) { final String route = ctx.getRoute(); if (route == null) { @@ -87,6 +82,7 @@ public boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx) { } /** Get the final sampling decision. This method is NOT thread-safe. */ + @Override public boolean sampleRequest(AppSecRequestContext ctx) { if (ctx == null) { return false; @@ -99,7 +95,7 @@ public boolean sampleRequest(AppSecRequestContext ctx) { return updateApiAccessIfExpired(hash); } - /** Release one permit for the sampler. This must be called after processing a span. */ + @Override public void releaseOne() { counter.release(); } @@ -173,20 +169,4 @@ private long computeApiHash(final String route, final String method, final int s result = 31 * result + statusCode; return result; } - - public static final class NoOp extends ApiSecurityRequestSampler { - public NoOp() { - super(0, 0, SystemTimeSource.INSTANCE); - } - - @Override - public boolean preSampleRequest(@Nonnull AppSecRequestContext ctx) { - return false; - } - - @Override - public boolean sampleRequest(AppSecRequestContext ctx) { - return false; - } - } } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java index 5fd0729b200..cf41d9f0fdb 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java @@ -21,11 +21,10 @@ public class AppSecSpanPostProcessor implements SpanPostProcessor { private static final Logger log = LoggerFactory.getLogger(AppSecSpanPostProcessor.class); - private final ApiSecurityRequestSampler sampler; + private final ApiSecuritySampler sampler; private final EventProducerService producerService; - public AppSecSpanPostProcessor( - ApiSecurityRequestSampler sampler, EventProducerService producerService) { + public AppSecSpanPostProcessor(ApiSecuritySampler sampler, EventProducerService producerService) { this.sampler = sampler; this.producerService = producerService; } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index a63806d0d92..95831f90c53 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -7,7 +7,7 @@ import static com.datadog.appsec.gateway.AppSecRequestContext.RESPONSE_HEADERS_ALLOW_LIST; import com.datadog.appsec.AppSecSystem; -import com.datadog.appsec.api.security.ApiSecurityRequestSampler; +import com.datadog.appsec.api.security.ApiSecuritySampler; import com.datadog.appsec.config.TraceSegmentPostProcessor; import com.datadog.appsec.event.EventProducerService; import com.datadog.appsec.event.EventProducerService.DataSubscriberInfo; @@ -83,7 +83,7 @@ public class GatewayBridge { private final SubscriptionService subscriptionService; private final EventProducerService producerService; - private final ApiSecurityRequestSampler requestSampler; + private final ApiSecuritySampler requestSampler; private final List traceSegmentPostProcessors; // subscriber cache @@ -109,7 +109,7 @@ public class GatewayBridge { public GatewayBridge( SubscriptionService subscriptionService, EventProducerService producerService, - ApiSecurityRequestSampler requestSampler, + ApiSecuritySampler requestSampler, List traceSegmentPostProcessors) { this.subscriptionService = subscriptionService; this.producerService = producerService; diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy similarity index 86% rename from dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy rename to dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy index b66402c99ce..a4ef9984786 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy @@ -4,12 +4,12 @@ import com.datadog.appsec.gateway.AppSecRequestContext import datadog.trace.api.time.ControllableTimeSource import datadog.trace.test.util.DDSpecification -class ApiSecurityRequestSamplerTest extends DDSpecification { +class ApiSecuritySamplerTest extends DDSpecification { void 'happy path with single request'() { given: - def ctx = createContext('route1', 'GET', 200) - def sampler = new ApiSecurityRequestSampler() + final ctx = createContext('route1', 'GET', 200) + final sampler = new ApiSecuritySamplerImpl() when: final preSampled = sampler.preSampleRequest(ctx) @@ -29,7 +29,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { given: AppSecRequestContext ctx1 = createContext('route1', 'GET', 200) AppSecRequestContext ctx2 = createContext('route1', 'GET', 200) - def sampler = new ApiSecurityRequestSampler() + final sampler = new ApiSecuritySamplerImpl() when: final preSampled1 = sampler.preSampleRequest(ctx1) @@ -52,7 +52,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { given: final ctx1 = Spy(createContext('route2', 'GET', 200)) final ctx2 = Spy(createContext('route3', 'GET', 200)) - final sampler = new ApiSecurityRequestSampler() + final sampler = new ApiSecuritySamplerImpl() assert sampler.MAX_POST_PROCESSING_TASKS > 0 when: 'exhaust the maximum number of concurrent contexts' @@ -82,7 +82,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { void 'preSampleRequest with null route'() { given: def ctx = createContext(null, 'GET', 200) - def sampler = new ApiSecurityRequestSampler() + def sampler = new ApiSecuritySamplerImpl() when: def preSampled = sampler.preSampleRequest(ctx) @@ -94,7 +94,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { void 'preSampleRequest with null method'() { given: def ctx = createContext('route1', null, 200) - def sampler = new ApiSecurityRequestSampler() + def sampler = new ApiSecuritySamplerImpl() when: def preSampled = sampler.preSampleRequest(ctx) @@ -106,7 +106,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { void 'preSampleRequest with 0 status code'() { given: def ctx = createContext('route1', 'GET', 0) - def sampler = new ApiSecurityRequestSampler() + def sampler = new ApiSecuritySamplerImpl() when: def preSampled = sampler.preSampleRequest(ctx) @@ -117,7 +117,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { void 'sampleRequest with null context throws'() { given: - def sampler = new ApiSecurityRequestSampler() + def sampler = new ApiSecuritySamplerImpl() when: sampler.preSampleRequest(null) @@ -128,7 +128,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { void 'sampleRequest without prior preSampleRequest never works'() { given: - def sampler = new ApiSecurityRequestSampler() + def sampler = new ApiSecuritySamplerImpl() def ctx = createContext('route1', 'GET', 200) when: @@ -147,7 +147,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { timeSource.set(0) final long expirationTimeInMs = 10L final long expirationTimeInNs = expirationTimeInMs * 1_000_000 - def sampler = new ApiSecurityRequestSampler(10, expirationTimeInMs, timeSource) + def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource) when: def sampled = sampler.sampleRequest(ctx) @@ -175,7 +175,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { timeSource.set(0) final long expirationTimeInMs = 10_000 final int maxCapacity = 10 - ApiSecurityRequestSampler sampler = new ApiSecurityRequestSampler(10, expirationTimeInMs, timeSource) + ApiSecuritySamplerImpl sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource) expect: for (int i = 0; i < maxCapacity * 10; i++) { @@ -194,7 +194,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification { timeSource.set(0) final long expirationTimeInMs = 10_000 final int maxCapacity = 10 - ApiSecurityRequestSampler sampler = new ApiSecurityRequestSampler(10, expirationTimeInMs, timeSource) + ApiSecuritySamplerImpl sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource) expect: for (int i = 0; i < maxCapacity * 10; i++) { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy index b77c98c353c..aea3fd68232 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy @@ -13,7 +13,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'schema extracted on happy path'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def subInfo = Mock(EventProducerService.DataSubscriberInfo) def span = Mock(AgentSpan) @@ -44,7 +44,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'no schema extracted if sampling is false'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def span = Mock(AgentSpan) def reqCtx = Mock(RequestContext) @@ -68,7 +68,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'permit is released even if request context close throws'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def span = Mock(AgentSpan) def reqCtx = Mock(RequestContext) @@ -95,7 +95,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'context is cleaned up on timeout'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def span = Mock(AgentSpan) def reqCtx = Mock(RequestContext) @@ -118,7 +118,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'process null request context does nothing'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def span = Mock(AgentSpan) def processor = new AppSecSpanPostProcessor(sampler, producer) @@ -134,7 +134,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'process null appsec request context does nothing'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def span = Mock(AgentSpan) def reqCtx = Mock(RequestContext) @@ -152,7 +152,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'process already closed context does nothing'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def span = Mock(AgentSpan) def reqCtx = Mock(RequestContext) @@ -172,7 +172,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'process throws on null span'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def processor = new AppSecSpanPostProcessor(sampler, producer) @@ -186,7 +186,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'empty event subscription does not break the process'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def subInfo = Mock(EventProducerService.DataSubscriberInfo) def span = Mock(AgentSpan) @@ -215,7 +215,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { void 'expired event subscription does not break the process'() { given: - def sampler = Mock(ApiSecurityRequestSampler) + def sampler = Mock(ApiSecuritySamplerImpl) def producer = Mock(EventProducerService) def subInfo = Mock(EventProducerService.DataSubscriberInfo) def span = Mock(AgentSpan) diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index b8cb74e0992..961838e01bd 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -1,7 +1,7 @@ package com.datadog.appsec.gateway import com.datadog.appsec.AppSecSystem -import com.datadog.appsec.api.security.ApiSecurityRequestSampler +import com.datadog.appsec.api.security.ApiSecuritySamplerImpl import com.datadog.appsec.config.TraceSegmentPostProcessor import com.datadog.appsec.event.EventDispatcher import com.datadog.appsec.event.EventProducerService @@ -27,7 +27,6 @@ import datadog.trace.test.util.DDSpecification import java.util.function.BiConsumer import java.util.function.BiFunction -import java.util.function.Consumer import java.util.function.Function import java.util.function.Supplier @@ -77,7 +76,7 @@ class GatewayBridgeSpecification extends DDSpecification { } TraceSegmentPostProcessor pp = Mock() - ApiSecurityRequestSampler requestSampler = Mock(ApiSecurityRequestSampler) + ApiSecuritySamplerImpl requestSampler = Mock(ApiSecuritySamplerImpl) GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, requestSampler, [pp]) Supplier> requestStartedCB From e9d1e7540f073893abc57e637248d107629f5a73 Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Wed, 5 Mar 2025 15:15:33 +0100 Subject: [PATCH 09/15] wip --- .../appsec/src/main/java/com/datadog/appsec/AppSecSystem.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java index e775989a83b..d30287c0cd5 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java @@ -71,8 +71,9 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s ApiSecuritySampler requestSampler; if (Config.get().isApiSecurityEnabled()) { - // TODO: Address support for 1-click enablement requestSampler = new ApiSecuritySamplerImpl(); + // When DD_API_SECURITY_ENABLED=true, ths post-processor is set even when AppSec is inactive. This should be + // low overhead since the post-processor exits early when there is no AppSec request context. SpanPostProcessor.Holder.INSTANCE = new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER); } else { From 35514640e113f7495ba288aee25a419d7387011d Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Thu, 6 Mar 2025 10:27:02 +0100 Subject: [PATCH 10/15] wip --- .../src/main/java/com/datadog/appsec/AppSecSystem.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java index d30287c0cd5..a8d18c2f515 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java @@ -72,8 +72,9 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s ApiSecuritySampler requestSampler; if (Config.get().isApiSecurityEnabled()) { requestSampler = new ApiSecuritySamplerImpl(); - // When DD_API_SECURITY_ENABLED=true, ths post-processor is set even when AppSec is inactive. This should be - // low overhead since the post-processor exits early when there is no AppSec request context. + // When DD_API_SECURITY_ENABLED=true, ths post-processor is set even when AppSec is inactive. + // This should be low overhead since the post-processor exits early if there's no AppSec + // context. SpanPostProcessor.Holder.INSTANCE = new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER); } else { From 21e1114b7cc22cd78e186efa2fb46556d11d58ae Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Thu, 6 Mar 2025 11:30:30 +0100 Subject: [PATCH 11/15] wip --- .../appsec/api/security/AppSecSpanPostProcessor.java | 4 ++++ .../appsec/api/security/AppSecSpanPostProcessorTest.groovy | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java index cf41d9f0fdb..77be93254f2 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java @@ -59,6 +59,10 @@ public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutChe } finally { ctx.setKeepOpenForApiSecurityPostProcessing(false); try { + // XXX: Close the additive first. This is not strictly needed, but it'll prevent getting it + // detected as a + // missed request-ended event. + ctx.closeAdditive(); ctx.close(); } catch (Exception e) { log.debug("Error closing AppSecRequestContext", e); diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy index aea3fd68232..e2287682bcd 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy @@ -37,6 +37,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { 1 * producer.publishDataEvent(_, ctx, _, _) 1 * ctx.commitDerivatives(traceSegment) 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.closeAdditive() 1 * ctx.close() 1 * sampler.releaseOne() 0 * _ @@ -61,6 +62,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true 1 * sampler.sampleRequest(_) >> false 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.closeAdditive() 1 * ctx.close() 1 * sampler.releaseOne() 0 * _ @@ -88,6 +90,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { 1 * reqCtx.getTraceSegment() >> traceSegment 1 * producer.getDataSubscribers(_) >> null 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.closeAdditive() 1 * ctx.close() >> { throw new RuntimeException() } 1 * sampler.releaseOne() 0 * _ @@ -111,6 +114,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { 1 * reqCtx.getData(_) >> ctx 1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.closeAdditive() 1 * ctx.close() 1 * sampler.releaseOne() 0 * _ @@ -208,6 +212,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { 1 * producer.getDataSubscribers(_) >> subInfo 1 * subInfo.isEmpty() >> true 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.closeAdditive() 1 * ctx.close() 1 * sampler.releaseOne() 0 * _ @@ -238,6 +243,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification { 1 * subInfo.isEmpty() >> false 1 * producer.publishDataEvent(_, ctx, _, _) >> { throw new ExpiredSubscriberInfoException() } 1 * ctx.setKeepOpenForApiSecurityPostProcessing(false) + 1 * ctx.closeAdditive() 1 * ctx.close() 1 * sampler.releaseOne() 0 * _ From 7ecd889d9545884bc0b1f966c1d4edbb42395fb5 Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Mon, 17 Mar 2025 08:23:35 +0100 Subject: [PATCH 12/15] Remove excessive logs and tweak condition --- .../datadog/appsec/api/security/ApiSecuritySamplerImpl.java | 6 +----- .../appsec/api/security/AppSecSpanPostProcessor.java | 1 - 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java index 21c13ebf585..c51bd46ef44 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java @@ -54,23 +54,19 @@ public ApiSecuritySamplerImpl( public boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx) { final String route = ctx.getRoute(); if (route == null) { - log.debug("Route is null, skipping API security sampling"); return false; } final String method = ctx.getMethod(); if (method == null) { - log.debug("Method is null, skipping API security sampling"); return false; } final int statusCode = ctx.getResponseStatus(); - if (statusCode == 0) { - log.debug("Status code is 0, skipping API security sampling"); + if (statusCode <= 0) { return false; } long hash = computeApiHash(route, method, statusCode); ctx.setApiSecurityEndpointHash(hash); if (!isApiAccessExpired(hash)) { - log.debug("API security sampling is not required for this request"); return false; } if (counter.tryAcquire()) { diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java index 77be93254f2..cb87b572137 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java @@ -41,7 +41,6 @@ public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutChe } if (!ctx.isKeepOpenForApiSecurityPostProcessing()) { - log.debug("Request context not marked for API security post-processing"); return; } From c429be8c564766686a7ed270156ff0333bb27d96 Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Mon, 17 Mar 2025 08:24:32 +0100 Subject: [PATCH 13/15] typo --- .../java/com/datadog/appsec/gateway/AppSecRequestContext.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 1e748e3048a..ad43e6bfec3 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 @@ -596,7 +596,7 @@ public String getSessionId() { /** * Close the context and release all resources. This method is idempotent and can be called * multiple times. For each root span, this method is always called from - * CoreTracer#onRootSpaPublished. + * CoreTracer#onRootSpanPublished. */ @Override public void close() { From 5488b19d566567ee617fcb18c881f3b47fc81d94 Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Mon, 17 Mar 2025 08:59:16 +0100 Subject: [PATCH 14/15] fix test --- .../com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy | 1 - 1 file changed, 1 deletion(-) diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 961838e01bd..c2453c52bf2 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -105,7 +105,6 @@ class GatewayBridgeSpecification extends DDSpecification { BiFunction> shellCmdCB BiFunction> userCB TriFunction> loginEventCB - Consumer postProcessingCB void setup() { callInitAndCaptureCBs() From 98d957d662cd80dd30a901421c5f49b1557acb22 Mon Sep 17 00:00:00 2001 From: Santiago Mola Date: Wed, 19 Mar 2025 10:48:10 +0100 Subject: [PATCH 15/15] fix merge conflict --- .../appsec/config/AppSecConfigServiceImplSpecification.groovy | 1 - 1 file changed, 1 deletion(-) diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy index fd622945d6e..56fd63f89f8 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy @@ -592,7 +592,6 @@ class AppSecConfigServiceImplSpecification extends DDSpecification { listeners.savedFeaturesListener = it[2] } 1 * poller.addConfigurationEndListener(_) >> { listeners.savedConfEndListener = it[0] } - 1 * poller.addCapabilities(CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE) 1 * poller.addCapabilities(CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE) 1 * poller.addCapabilities(CAPABILITY_ASM_DD_RULES | CAPABILITY_ASM_IP_BLOCKING