Skip to content

Commit 1358e8e

Browse files
claponcetdevflow.devflow-routing-intake
claponcet
and
devflow.devflow-routing-intake
authored
Support DD_APM_TRACING_ENABLED=false for AI Guard traces (#11885)
Support DD_APM_TRACING_ENABLED=false for AI Guard traces Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Fix remaining EVENT_TAG reference in Groovy test Merge branch 'master' into clara.poncet/ai-guard-standalone Set AI Guard bit on _dd.p.ts propagated trace source tag Merge branch 'master' into clara.poncet/ai-guard-standalone Fix flaky AIGuardSamplingTest by asserting on force-keep, not sampler rate-limit state Merge branch 'master' into clara.poncet/ai-guard-standalone Merge branch 'master' into clara.poncet/ai-guard-standalone Replace isProductMarked varargs overload with allocation-free two-arg overload Merge branch 'master' into clara.poncet/ai-guard-standalone Merge branch 'master' into clara.poncet/ai-guard-standalone Fix stale junit-utils import after test-junit-utils module rename in master Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>
1 parent 613b386 commit 1358e8e

8 files changed

Lines changed: 160 additions & 7 deletions

File tree

‎dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import datadog.communication.http.OkHttpUtils;
1414
import datadog.context.ContextScope;
1515
import datadog.trace.api.Config;
16+
import datadog.trace.api.ProductTraceSource;
1617
import datadog.trace.api.aiguard.AIGuard;
1718
import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError;
1819
import datadog.trace.api.aiguard.AIGuard.AIGuardClientError;
@@ -71,7 +72,7 @@ public BadConfigurationException(final String message) {
7172
static final String ACTION_TAG = "ai_guard.action";
7273
static final String REASON_TAG = "ai_guard.reason";
7374
static final String BLOCKED_TAG = "ai_guard.blocked";
74-
static final String EVENT_TAG = "ai_guard.event";
75+
7576
static final String META_STRUCT_TAG = "ai_guard";
7677
static final String META_STRUCT_MESSAGES = "messages";
7778
static final String META_STRUCT_CATEGORIES = "attack_categories";
@@ -275,7 +276,8 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
275276
final AgentSpan localRootSpan = span.getLocalRootSpan();
276277
if (localRootSpan != null) {
277278
localRootSpan.setTag(Tags.AI_GUARD_KEEP, true);
278-
localRootSpan.setTag(EVENT_TAG, true);
279+
localRootSpan.setTag(Tags.AI_GUARD_EVENT, true);
280+
localRootSpan.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
279281
applyClientIpTags(localRootSpan);
280282
// copyAnomalyDetectionTags MUST run after applyClientIpTags, to make
281283
// sure client IP tags were populated.

‎dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import com.fasterxml.jackson.databind.PropertyNamingStrategies
66
import com.squareup.moshi.Moshi
77
import datadog.common.version.VersionInfo
88
import datadog.trace.api.Config
9+
import datadog.trace.api.ProductTraceSource
910
import datadog.trace.api.aiguard.AIGuard
1011
import datadog.trace.api.gateway.RequestContext
1112
import datadog.trace.api.telemetry.WafMetricCollector
@@ -192,7 +193,8 @@ class AIGuardInternalTests extends DDSpecification {
192193
then:
193194
1 * span.setTag(AIGuardInternal.TARGET_TAG, suite.target)
194195
1 * localRootSpan.setTag(Tags.AI_GUARD_KEEP, true)
195-
1 * localRootSpan.setTag(AIGuardInternal.EVENT_TAG, true)
196+
1 * localRootSpan.setTag(Tags.AI_GUARD_EVENT, true)
197+
1 * localRootSpan.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD)
196198
if (suite.target == 'tool') {
197199
1 * span.setTag(AIGuardInternal.TOOL_TAG, 'calc')
198200
}

‎dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,12 +64,12 @@ public void setSamplingPriorityIfNecessary() {
6464
// Locks inside DDSpanContext ensure the correct behavior in the race case
6565
DDSpan rootSpan = getRootSpan();
6666
if (traceConfig.sampler instanceof PrioritySampler && rootSpan != null) {
67-
// Ignore the force-keep priority in the absence of propagated _dd.p.ts span tag marked for
68-
// ASM.
67+
// Skip sampler override when _dd.p.ts is marked for ASM or AI Guard.
6968
if ((!Config.get().isApmTracingEnabled()
7069
&& !ProductTraceSource.isProductMarked(
7170
rootSpan.spanContext().getPropagationTags().getTraceSource(),
72-
ProductTraceSource.ASM))
71+
ProductTraceSource.ASM,
72+
ProductTraceSource.AI_GUARD))
7373
|| rootSpan.spanContext().getSamplingPriority() == PrioritySampling.UNSET) {
7474
((PrioritySampler) traceConfig.sampler).setSamplingPriority(rootSpan);
7575
}

‎dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
package datadog.trace.common.sampling;
2+
3+
import static datadog.trace.api.config.AppSecConfig.APPSEC_ENABLED;
4+
import static datadog.trace.api.config.GeneralConfig.APM_TRACING_ENABLED;
5+
import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP;
6+
import static org.junit.jupiter.api.Assertions.assertEquals;
7+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
8+
9+
import datadog.trace.api.ProductTraceSource;
10+
import datadog.trace.bootstrap.instrumentation.api.Tags;
11+
import datadog.trace.common.writer.ListWriter;
12+
import datadog.trace.core.CoreTracer;
13+
import datadog.trace.core.DDCoreJavaSpecification;
14+
import datadog.trace.core.DDSpan;
15+
import datadog.trace.test.junit.utils.config.WithConfig;
16+
import java.util.Map;
17+
import org.junit.jupiter.api.BeforeEach;
18+
import org.junit.jupiter.api.Test;
19+
20+
/** Verifies that AI Guard traces are kept regardless of APM/ASM configuration. */
21+
public class AIGuardSamplingTest extends DDCoreJavaSpecification {
22+
23+
private ListWriter writer;
24+
private CoreTracer tracer;
25+
26+
@BeforeEach
27+
void setUp() {
28+
writer = new ListWriter();
29+
tracer = tracerBuilder().writer(writer).build();
30+
}
31+
32+
@Test
33+
void aiGuardTraceIsKeptWhenApmEnabled() throws Exception {
34+
DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start();
35+
span.setTag(Tags.AI_GUARD_KEEP, true);
36+
span.setTag(Tags.AI_GUARD_EVENT, true);
37+
span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
38+
span.finish();
39+
writer.waitForTraces(1);
40+
assertEquals(USER_KEEP, (int) span.getSamplingPriority());
41+
assertDecisionMakerIsAiGuard(span);
42+
}
43+
44+
@Test
45+
@WithConfig(key = APM_TRACING_ENABLED, value = "false")
46+
void aiGuardTraceIsKeptWhenApmAndAsmDisabled() throws Exception {
47+
DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start();
48+
span.setTag(Tags.AI_GUARD_KEEP, true);
49+
span.setTag(Tags.AI_GUARD_EVENT, true);
50+
span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
51+
span.finish();
52+
writer.waitForTraces(1);
53+
assertEquals(USER_KEEP, (int) span.getSamplingPriority());
54+
assertDecisionMakerIsAiGuard(span);
55+
}
56+
57+
@Test
58+
@WithConfig(key = APM_TRACING_ENABLED, value = "false")
59+
@WithConfig(key = APPSEC_ENABLED, value = "true")
60+
void aiGuardTraceIsKeptInAsmStandaloneMode() throws Exception {
61+
DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start();
62+
span.setTag(Tags.AI_GUARD_KEEP, true);
63+
span.setTag(Tags.AI_GUARD_EVENT, true);
64+
span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
65+
span.finish();
66+
writer.waitForTraces(1);
67+
assertEquals(USER_KEEP, (int) span.getSamplingPriority());
68+
assertDecisionMakerIsAiGuard(span);
69+
}
70+
71+
/** _dd.p.ts alone (without AI_GUARD_KEEP) must not bypass the sampler. */
72+
@Test
73+
@WithConfig(key = APM_TRACING_ENABLED, value = "false")
74+
@WithConfig(key = APPSEC_ENABLED, value = "true")
75+
void propagatedTraceSourceAloneDoesNotBypassSampler() throws Exception {
76+
DDSpan span = (DDSpan) tracer.buildSpan("datadog", "op").start();
77+
span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
78+
span.finish();
79+
writer.waitForTraces(1);
80+
81+
// _dd.p.ts alone must not force-keep the trace; only AI_GUARD_KEEP does that.
82+
assertNotEquals(USER_KEEP, (int) span.getSamplingPriority());
83+
}
84+
85+
/**
86+
* AI Guard traces must be kept even when AsmStandaloneSampler has exhausted its rate-limit slot.
87+
*/
88+
@Test
89+
@WithConfig(key = APM_TRACING_ENABLED, value = "false")
90+
@WithConfig(key = APPSEC_ENABLED, value = "true")
91+
void aiGuardTraceIsKeptAfterRateLimitSlotIsExhausted() throws Exception {
92+
// consume the first allowed slot
93+
DDSpan first = (DDSpan) tracer.buildSpan("datadog", "op").start();
94+
first.finish();
95+
writer.waitForTraces(1);
96+
97+
// AI Guard trace must still be force-kept even though the rate-limit slot is gone
98+
DDSpan aiGuard = (DDSpan) tracer.buildSpan("datadog", "op").start();
99+
aiGuard.setTag(Tags.AI_GUARD_KEEP, true);
100+
aiGuard.setTag(Tags.AI_GUARD_EVENT, true);
101+
aiGuard.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
102+
aiGuard.finish();
103+
writer.waitForTraces(2);
104+
105+
assertEquals(USER_KEEP, (int) aiGuard.getSamplingPriority());
106+
assertDecisionMakerIsAiGuard(aiGuard);
107+
}
108+
109+
private static void assertDecisionMakerIsAiGuard(DDSpan span) {
110+
Map<String, String> ptags = span.spanContext().getPropagationTags().createTagMap();
111+
assertEquals("-13", ptags.get("_dd.p.dm"), "_dd.p.dm must be -13 (AI Guard decision maker)");
112+
assertEquals("20", ptags.get("_dd.p.ts"), "_dd.p.ts must have AI Guard bit set (0x20)");
113+
}
114+
}

‎dd-trace-core/src/test/java/datadog/trace/core/taginterceptor/TagInterceptorTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,13 @@ static Stream<Arguments> setPrioritySamplingViaTagArguments() {
389389
arguments("asm.keep / 'true'", Tags.ASM_KEEP, "true", (int) PrioritySampling.USER_KEEP),
390390
arguments("asm.keep / 'false'", Tags.ASM_KEEP, "false", null),
391391
arguments("asm.keep / 'asdf'", Tags.ASM_KEEP, "asdf", null),
392+
arguments(
393+
"ai_guard.keep / true", Tags.AI_GUARD_KEEP, true, (int) PrioritySampling.USER_KEEP),
394+
arguments("ai_guard.keep / false", Tags.AI_GUARD_KEEP, false, null),
395+
arguments(
396+
"ai_guard.keep / 'true'", Tags.AI_GUARD_KEEP, "true", (int) PrioritySampling.USER_KEEP),
397+
arguments("ai_guard.keep / 'false'", Tags.AI_GUARD_KEEP, "false", null),
398+
arguments("ai_guard.keep / 'asdf'", Tags.AI_GUARD_KEEP, "asdf", null),
392399
arguments(
393400
"sampling.priority / -1", Tags.SAMPLING_PRIORITY, -1, (int) PrioritySampling.USER_DROP),
394401
arguments(

‎internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ public class ProductTraceSource {
2222
public static final int DSM = 0x04;
2323
public static final int DJM = 0x08;
2424
public static final int DBM = 0x10;
25+
public static final int AI_GUARD = 0x20;
2526

2627
/** Updates the bitfield by setting the bit corresponding to a specific product. */
2728
public static int updateProduct(int bitfield, int product) {
@@ -30,7 +31,12 @@ public static int updateProduct(int bitfield, int product) {
3031

3132
/** Checks if the bitfield is marked for a specific product. */
3233
public static boolean isProductMarked(final int bitfield, int product) {
33-
return (bitfield & product) != 0; // Check if the bit is set
34+
return (bitfield & product) != 0;
35+
}
36+
37+
/** Checks if the bitfield is marked for either of the two given products. */
38+
public static boolean isProductMarked(final int bitfield, int productA, int productB) {
39+
return (bitfield & (productA | productB)) != 0;
3440
}
3541

3642
/**

‎internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,9 @@ public class Tags {
183183
/** AI Guard force tracer to keep the trace */
184184
public static final String AI_GUARD_KEEP = "ai_guard.keep";
185185

186+
/** Marks a span as originating from an AI Guard evaluation */
187+
public static final String AI_GUARD_EVENT = "ai_guard.event";
188+
186189
public static final String PROPAGATED_TRACE_SOURCE = "_dd.p.ts";
187190
public static final String PROPAGATED_DEBUG = "_dd.p.debug";
188191

‎internal-api/src/test/groovy/datadog/trace/api/ProductTraceSourceTest.groovy‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,23 @@ class ProductTraceSourceTest extends DDSpecification {
2828
value | product | expected
2929
ProductTraceSource.ASM | ProductTraceSource.ASM | true
3030
ProductTraceSource.DSM | ProductTraceSource.ASM | false
31+
ProductTraceSource.AI_GUARD | ProductTraceSource.AI_GUARD | true
32+
ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | false
33+
}
34+
35+
void 'test isProductMarked with two products'(){
36+
when:
37+
final result = ProductTraceSource.isProductMarked(value, productA, productB)
38+
39+
then:
40+
result == expected
41+
42+
where:
43+
value | productA | productB | expected
44+
ProductTraceSource.ASM | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | true
45+
ProductTraceSource.AI_GUARD | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | true
46+
ProductTraceSource.DSM | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | false
47+
ProductTraceSource.UNSET | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | false
3148
}
3249

3350
void 'test getBitfieldHex'(){
@@ -41,6 +58,7 @@ class ProductTraceSourceTest extends DDSpecification {
4158
value | expected
4259
ProductTraceSource.UNSET | "00"
4360
ProductTraceSource.ASM | "02"
61+
ProductTraceSource.AI_GUARD | "20"
4462
}
4563

4664
void 'test parseBitfieldHex'(){
@@ -56,5 +74,6 @@ class ProductTraceSourceTest extends DDSpecification {
5674
null | ProductTraceSource.UNSET
5775
"" | ProductTraceSource.UNSET
5876
"02" | ProductTraceSource.ASM
77+
"20" | ProductTraceSource.AI_GUARD
5978
}
6079
}

0 commit comments

Comments
 (0)