From daa207fde54919960d2c23ecaca738cc0a433757 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Mon, 12 May 2025 13:07:49 +0200 Subject: [PATCH 1/4] Fix --- .../datadog/appsec/gateway/GatewayBridge.java | 6 +- .../security/ApiSecuritySamplerTest.groovy | 19 ++++++ .../gateway/GatewayBridgeSpecification.groovy | 66 +++++++++++++++++++ 3 files changed, 90 insertions(+), 1 deletion(-) 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 7d29e6ee8b0..9ba2bdb07a3 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,6 +20,7 @@ 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; @@ -668,7 +669,10 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { Map tags = spanInfo.getTags(); if (maybeSampleForApiSecurity(ctx, spanInfo, tags)) { - ctx.setKeepOpenForApiSecurityPostProcessing(true); + if (!Config.get().isApmTracingEnabled()) { + traceSeg.setTagTop(Tags.ASM_KEEP, true); + traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); + } } else { ctx.closeWafContext(); } diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy index a4ef9984786..4f9980af136 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy @@ -25,6 +25,25 @@ class ApiSecuritySamplerTest extends DDSpecification { sampled } + void 'happy path with single request and tracing disabled'() { + given: + final ctx = createContext('route1', 'GET', 200) + final sampler = new ApiSecuritySamplerImpl() + + when: + final preSampled = sampler.preSampleRequest(ctx) + + then: + preSampled + + when: + ctx.setKeepOpenForApiSecurityPostProcessing(true) + final sampled = sampler.sampleRequest(ctx) + + then: + sampled + } + void 'second request is not sampled for the same endpoint'() { given: AppSecRequestContext ctx1 = createContext('route1', 'GET', 200) 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 df6bd6d24b3..cc6edb8d562 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 @@ -9,6 +9,9 @@ import com.datadog.appsec.event.data.DataBundle import com.datadog.appsec.event.data.KnownAddresses import com.datadog.appsec.report.AppSecEvent import com.datadog.appsec.report.AppSecEventWrapper +import datadog.trace.api.ProductTraceSource +import datadog.trace.api.config.GeneralConfig +import static datadog.trace.api.config.IastConfig.IAST_DEDUPLICATION_ENABLED import datadog.trace.api.function.TriConsumer import datadog.trace.api.function.TriFunction import datadog.trace.api.gateway.BlockResponseFunction @@ -22,6 +25,7 @@ import datadog.trace.api.internal.TraceSegment import datadog.trace.api.telemetry.LoginEvent import datadog.trace.api.telemetry.WafMetricCollector import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import datadog.trace.bootstrap.instrumentation.api.Tags import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase import datadog.trace.test.util.DDSpecification @@ -1162,4 +1166,66 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * eventDispatcher.getDataSubscribers(KnownAddresses.SESSION_ID) >> nonEmptyDsInfo 1 * eventDispatcher.publishDataEvent(_, _, _, _) } + + void 'test api security sampling'() { + given: + AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) + RequestContext mockCtx = Stub(RequestContext) { + getData(RequestContextSlot.APPSEC) >> mockAppSecCtx + getTraceSegment() >> traceSegment + } + IGSpanInfo spanInfo = Mock(AgentSpan) + + when: + def flow = requestEndedCB.apply(mockCtx, spanInfo) + + then: + 1 * mockAppSecCtx.transferCollectedEvents() >> [] + 1 * spanInfo.getTags() >> ['http.route': 'route'] + 1 * requestSampler.preSampleRequest(_) >> true + 0 * traceSegment.setTagTop(Tags.ASM_KEEP, true) + 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) + } + + void 'test api security sampling - trace excluded'() { + given: + AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) + RequestContext mockCtx = Stub(RequestContext) { + getData(RequestContextSlot.APPSEC) >> mockAppSecCtx + getTraceSegment() >> traceSegment + } + IGSpanInfo spanInfo = Mock(AgentSpan) + + when: + def flow = requestEndedCB.apply(mockCtx, spanInfo) + + then: + 1 * mockAppSecCtx.transferCollectedEvents() >> [] + 1 * spanInfo.getTags() >> ['http.route': 'route'] + 1 * requestSampler.preSampleRequest(_) >> false + 0 * traceSegment.setTagTop(Tags.ASM_KEEP, true) + 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) + } + + void 'test api security sampling with tracing disabled'() { + given: + injectSysConfig(GeneralConfig.APM_TRACING_ENABLED, "false") + AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) + RequestContext mockCtx = Stub(RequestContext) { + getData(RequestContextSlot.APPSEC) >> mockAppSecCtx + getTraceSegment() >> traceSegment + } + IGSpanInfo spanInfo = Mock(AgentSpan) + + when: + def flow = requestEndedCB.apply(mockCtx, spanInfo) + + then: + 1 * mockAppSecCtx.transferCollectedEvents() >> [] + 1 * spanInfo.getTags() >> ['http.route': 'route'] + 1 * requestSampler.preSampleRequest(_) >> true + 1 * traceSegment.setTagTop(Tags.ASM_KEEP, true) + 1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) + } + } From 5f1ea30f16f21f95233c9f7c57b92bf30a9e3be1 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Mon, 12 May 2025 13:12:59 +0200 Subject: [PATCH 2/4] undo --- .../security/ApiSecuritySamplerTest.groovy | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy index 4f9980af136..a4ef9984786 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy @@ -25,25 +25,6 @@ class ApiSecuritySamplerTest extends DDSpecification { sampled } - void 'happy path with single request and tracing disabled'() { - given: - final ctx = createContext('route1', 'GET', 200) - final sampler = new ApiSecuritySamplerImpl() - - when: - final preSampled = sampler.preSampleRequest(ctx) - - then: - preSampled - - when: - ctx.setKeepOpenForApiSecurityPostProcessing(true) - final sampled = sampler.sampleRequest(ctx) - - then: - sampled - } - void 'second request is not sampled for the same endpoint'() { given: AppSecRequestContext ctx1 = createContext('route1', 'GET', 200) From 0fbd784b04b429e613187feb16350cc32ec91bb9 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Mon, 12 May 2025 15:45:42 +0200 Subject: [PATCH 3/4] FIX --- .../datadog/appsec/gateway/GatewayBridge.java | 7 ++---- .../gateway/GatewayBridgeSpecification.groovy | 24 ++----------------- 2 files changed, 4 insertions(+), 27 deletions(-) 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 9ba2bdb07a3..467bcaa2c52 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; @@ -669,10 +668,8 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { Map tags = spanInfo.getTags(); if (maybeSampleForApiSecurity(ctx, spanInfo, tags)) { - if (!Config.get().isApmTracingEnabled()) { - traceSeg.setTagTop(Tags.ASM_KEEP, true); - traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); - } + traceSeg.setTagTop(Tags.ASM_KEEP, true); + traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); } else { ctx.closeWafContext(); } 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 cc6edb8d562..38448d79e92 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 @@ -1183,8 +1183,8 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * mockAppSecCtx.transferCollectedEvents() >> [] 1 * spanInfo.getTags() >> ['http.route': 'route'] 1 * requestSampler.preSampleRequest(_) >> true - 0 * traceSegment.setTagTop(Tags.ASM_KEEP, true) - 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) + 1 * traceSegment.setTagTop(Tags.ASM_KEEP, true) + 1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) } void 'test api security sampling - trace excluded'() { @@ -1207,25 +1207,5 @@ class GatewayBridgeSpecification extends DDSpecification { 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) } - void 'test api security sampling with tracing disabled'() { - given: - injectSysConfig(GeneralConfig.APM_TRACING_ENABLED, "false") - AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) - RequestContext mockCtx = Stub(RequestContext) { - getData(RequestContextSlot.APPSEC) >> mockAppSecCtx - getTraceSegment() >> traceSegment - } - IGSpanInfo spanInfo = Mock(AgentSpan) - - when: - def flow = requestEndedCB.apply(mockCtx, spanInfo) - - then: - 1 * mockAppSecCtx.transferCollectedEvents() >> [] - 1 * spanInfo.getTags() >> ['http.route': 'route'] - 1 * requestSampler.preSampleRequest(_) >> true - 1 * traceSegment.setTagTop(Tags.ASM_KEEP, true) - 1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) - } } From c1bd3d1e6a9684c96830631614f959773afea4d1 Mon Sep 17 00:00:00 2001 From: "alejandro.gonzalez" Date: Wed, 14 May 2025 12:09:27 +0200 Subject: [PATCH 4/4] wip --- .../datadog/appsec/gateway/GatewayBridge.java | 7 +++-- .../gateway/GatewayBridgeSpecification.groovy | 26 ++++++++++++++----- 2 files changed, 25 insertions(+), 8 deletions(-) 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 467bcaa2c52..9ba2bdb07a3 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,6 +20,7 @@ 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; @@ -668,8 +669,10 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { Map tags = spanInfo.getTags(); if (maybeSampleForApiSecurity(ctx, spanInfo, tags)) { - traceSeg.setTagTop(Tags.ASM_KEEP, true); - traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); + if (!Config.get().isApmTracingEnabled()) { + traceSeg.setTagTop(Tags.ASM_KEEP, true); + traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); + } } else { ctx.closeWafContext(); } 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 38448d79e92..4eb14da17dc 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 @@ -1175,16 +1175,14 @@ class GatewayBridgeSpecification extends DDSpecification { getTraceSegment() >> traceSegment } IGSpanInfo spanInfo = Mock(AgentSpan) - when: def flow = requestEndedCB.apply(mockCtx, spanInfo) - then: 1 * mockAppSecCtx.transferCollectedEvents() >> [] 1 * spanInfo.getTags() >> ['http.route': 'route'] 1 * requestSampler.preSampleRequest(_) >> true - 1 * traceSegment.setTagTop(Tags.ASM_KEEP, true) - 1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) + 0 * traceSegment.setTagTop(Tags.ASM_KEEP, true) + 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) } void 'test api security sampling - trace excluded'() { @@ -1195,10 +1193,8 @@ class GatewayBridgeSpecification extends DDSpecification { getTraceSegment() >> traceSegment } IGSpanInfo spanInfo = Mock(AgentSpan) - when: def flow = requestEndedCB.apply(mockCtx, spanInfo) - then: 1 * mockAppSecCtx.transferCollectedEvents() >> [] 1 * spanInfo.getTags() >> ['http.route': 'route'] @@ -1207,5 +1203,23 @@ class GatewayBridgeSpecification extends DDSpecification { 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) } + void 'test api security sampling with tracing disabled'() { + given: + injectSysConfig(GeneralConfig.APM_TRACING_ENABLED, "false") + AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) + RequestContext mockCtx = Stub(RequestContext) { + getData(RequestContextSlot.APPSEC) >> mockAppSecCtx + getTraceSegment() >> traceSegment + } + IGSpanInfo spanInfo = Mock(AgentSpan) + when: + def flow = requestEndedCB.apply(mockCtx, spanInfo) + then: + 1 * mockAppSecCtx.transferCollectedEvents() >> [] + 1 * spanInfo.getTags() >> ['http.route': 'route'] + 1 * requestSampler.preSampleRequest(_) >> true + 1 * traceSegment.setTagTop(Tags.ASM_KEEP, true) + 1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) + } }