@@ -9,6 +9,9 @@ import com.datadog.appsec.event.data.DataBundle
9
9
import com.datadog.appsec.event.data.KnownAddresses
10
10
import com.datadog.appsec.report.AppSecEvent
11
11
import com.datadog.appsec.report.AppSecEventWrapper
12
+ import datadog.trace.api.ProductTraceSource
13
+ import datadog.trace.api.config.GeneralConfig
14
+ import static datadog.trace.api.config.IastConfig.IAST_DEDUPLICATION_ENABLED
12
15
import datadog.trace.api.function.TriConsumer
13
16
import datadog.trace.api.function.TriFunction
14
17
import datadog.trace.api.gateway.BlockResponseFunction
@@ -22,6 +25,7 @@ import datadog.trace.api.internal.TraceSegment
22
25
import datadog.trace.api.telemetry.LoginEvent
23
26
import datadog.trace.api.telemetry.WafMetricCollector
24
27
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
28
+ import datadog.trace.bootstrap.instrumentation.api.Tags
25
29
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
26
30
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase
27
31
import datadog.trace.test.util.DDSpecification
@@ -1162,4 +1166,60 @@ class GatewayBridgeSpecification extends DDSpecification {
1162
1166
1 * eventDispatcher.getDataSubscribers(KnownAddresses.SESSION_ID) >> nonEmptyDsInfo
1163
1167
1 * eventDispatcher.publishDataEvent(_, _, _, _)
1164
1168
}
1169
+
1170
+ void 'test api security sampling'() {
1171
+ given:
1172
+ AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext)
1173
+ RequestContext mockCtx = Stub(RequestContext) {
1174
+ getData(RequestContextSlot.APPSEC) >> mockAppSecCtx
1175
+ getTraceSegment() >> traceSegment
1176
+ }
1177
+ IGSpanInfo spanInfo = Mock(AgentSpan)
1178
+ when:
1179
+ def flow = requestEndedCB.apply(mockCtx, spanInfo)
1180
+ then:
1181
+ 1 * mockAppSecCtx.transferCollectedEvents() >> []
1182
+ 1 * spanInfo.getTags() >> ['http.route': 'route']
1183
+ 1 * requestSampler.preSampleRequest(_) >> true
1184
+ 0 * traceSegment.setTagTop(Tags.ASM_KEEP, true)
1185
+ 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM)
1186
+ }
1187
+
1188
+ void 'test api security sampling - trace excluded'() {
1189
+ given:
1190
+ AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext)
1191
+ RequestContext mockCtx = Stub(RequestContext) {
1192
+ getData(RequestContextSlot.APPSEC) >> mockAppSecCtx
1193
+ getTraceSegment() >> traceSegment
1194
+ }
1195
+ IGSpanInfo spanInfo = Mock(AgentSpan)
1196
+ when:
1197
+ def flow = requestEndedCB.apply(mockCtx, spanInfo)
1198
+ then:
1199
+ 1 * mockAppSecCtx.transferCollectedEvents() >> []
1200
+ 1 * spanInfo.getTags() >> ['http.route': 'route']
1201
+ 1 * requestSampler.preSampleRequest(_) >> false
1202
+ 0 * traceSegment.setTagTop(Tags.ASM_KEEP, true)
1203
+ 0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM)
1204
+ }
1205
+
1206
+ void 'test api security sampling with tracing disabled'() {
1207
+ given:
1208
+ injectSysConfig(GeneralConfig.APM_TRACING_ENABLED, " false " )
1209
+ AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext)
1210
+ RequestContext mockCtx = Stub(RequestContext) {
1211
+ getData(RequestContextSlot.APPSEC) >> mockAppSecCtx
1212
+ getTraceSegment() >> traceSegment
1213
+ }
1214
+ IGSpanInfo spanInfo = Mock(AgentSpan)
1215
+ when:
1216
+ def flow = requestEndedCB.apply(mockCtx, spanInfo)
1217
+ then:
1218
+ 1 * mockAppSecCtx.transferCollectedEvents() >> []
1219
+ 1 * spanInfo.getTags() >> ['http.route': 'route']
1220
+ 1 * requestSampler.preSampleRequest(_) >> true
1221
+ 1 * traceSegment.setTagTop(Tags.ASM_KEEP, true)
1222
+ 1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM)
1223
+ }
1224
+
1165
1225
}
0 commit comments