diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java index 11ede90b51c..d921be448d8 100644 --- a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java @@ -13,6 +13,7 @@ import datadog.communication.http.OkHttpUtils; import datadog.context.ContextScope; import datadog.trace.api.Config; +import datadog.trace.api.ProductTraceSource; import datadog.trace.api.aiguard.AIGuard; import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError; import datadog.trace.api.aiguard.AIGuard.AIGuardClientError; @@ -71,7 +72,7 @@ public BadConfigurationException(final String message) { static final String ACTION_TAG = "ai_guard.action"; static final String REASON_TAG = "ai_guard.reason"; static final String BLOCKED_TAG = "ai_guard.blocked"; - static final String EVENT_TAG = "ai_guard.event"; + static final String META_STRUCT_TAG = "ai_guard"; static final String META_STRUCT_MESSAGES = "messages"; static final String META_STRUCT_CATEGORIES = "attack_categories"; @@ -275,7 +276,8 @@ public Evaluation evaluate(final List messages, final Options options) final AgentSpan localRootSpan = span.getLocalRootSpan(); if (localRootSpan != null) { localRootSpan.setTag(Tags.AI_GUARD_KEEP, true); - localRootSpan.setTag(EVENT_TAG, true); + localRootSpan.setTag(Tags.AI_GUARD_EVENT, true); + localRootSpan.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD); applyClientIpTags(localRootSpan); // copyAnomalyDetectionTags MUST run after applyClientIpTags, to make // sure client IP tags were populated. diff --git a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy index e355b2eec3c..114e83da3f1 100644 --- a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy +++ b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy @@ -6,6 +6,7 @@ import com.fasterxml.jackson.databind.PropertyNamingStrategies import com.squareup.moshi.Moshi import datadog.common.version.VersionInfo import datadog.trace.api.Config +import datadog.trace.api.ProductTraceSource import datadog.trace.api.aiguard.AIGuard import datadog.trace.api.gateway.RequestContext import datadog.trace.api.telemetry.WafMetricCollector @@ -192,7 +193,8 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setTag(AIGuardInternal.TARGET_TAG, suite.target) 1 * localRootSpan.setTag(Tags.AI_GUARD_KEEP, true) - 1 * localRootSpan.setTag(AIGuardInternal.EVENT_TAG, true) + 1 * localRootSpan.setTag(Tags.AI_GUARD_EVENT, true) + 1 * localRootSpan.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD) if (suite.target == 'tool') { 1 * span.setTag(AIGuardInternal.TOOL_TAG, 'calc') } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java b/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java index 2d97e895fb0..4cb9a3edb42 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java @@ -64,12 +64,12 @@ public void setSamplingPriorityIfNecessary() { // Locks inside DDSpanContext ensure the correct behavior in the race case DDSpan rootSpan = getRootSpan(); if (traceConfig.sampler instanceof PrioritySampler && rootSpan != null) { - // Ignore the force-keep priority in the absence of propagated _dd.p.ts span tag marked for - // ASM. + // Skip sampler override when _dd.p.ts is marked for ASM or AI Guard. if ((!Config.get().isApmTracingEnabled() && !ProductTraceSource.isProductMarked( rootSpan.spanContext().getPropagationTags().getTraceSource(), - ProductTraceSource.ASM)) + ProductTraceSource.ASM, + ProductTraceSource.AI_GUARD)) || rootSpan.spanContext().getSamplingPriority() == PrioritySampling.UNSET) { ((PrioritySampler) traceConfig.sampler).setSamplingPriority(rootSpan); } diff --git a/dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java b/dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java new file mode 100644 index 00000000000..fde358ec70f --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java @@ -0,0 +1,119 @@ +package datadog.trace.common.sampling; + +import static datadog.trace.api.config.AppSecConfig.APPSEC_ENABLED; +import static datadog.trace.api.config.GeneralConfig.APM_TRACING_ENABLED; +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP; +import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.trace.api.ProductTraceSource; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.common.writer.ListWriter; +import datadog.trace.core.CoreTracer; +import datadog.trace.core.DDCoreJavaSpecification; +import datadog.trace.core.DDSpan; +import datadog.trace.junit.utils.config.WithConfig; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Verifies that AI Guard traces are kept regardless of APM/ASM configuration. */ +public class AIGuardSamplingTest extends DDCoreJavaSpecification { + + private ListWriter writer; + private CoreTracer tracer; + + @BeforeEach + void setUp() { + writer = new ListWriter(); + tracer = tracerBuilder().writer(writer).build(); + } + + @Test + void aiGuardTraceIsKeptWhenApmEnabled() throws Exception { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start(); + span.setTag(Tags.AI_GUARD_KEEP, true); + span.setTag(Tags.AI_GUARD_EVENT, true); + span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD); + span.finish(); + writer.waitForTraces(1); + assertEquals(USER_KEEP, (int) span.getSamplingPriority()); + assertDecisionMakerIsAiGuard(span); + } + + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + void aiGuardTraceIsKeptWhenApmAndAsmDisabled() throws Exception { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start(); + span.setTag(Tags.AI_GUARD_KEEP, true); + span.setTag(Tags.AI_GUARD_EVENT, true); + span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD); + span.finish(); + writer.waitForTraces(1); + assertEquals(USER_KEEP, (int) span.getSamplingPriority()); + assertDecisionMakerIsAiGuard(span); + } + + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @WithConfig(key = APPSEC_ENABLED, value = "true") + void aiGuardTraceIsKeptInAsmStandaloneMode() throws Exception { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start(); + span.setTag(Tags.AI_GUARD_KEEP, true); + span.setTag(Tags.AI_GUARD_EVENT, true); + span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD); + span.finish(); + writer.waitForTraces(1); + assertEquals(USER_KEEP, (int) span.getSamplingPriority()); + assertDecisionMakerIsAiGuard(span); + } + + /** _dd.p.ts alone (without AI_GUARD_KEEP) must not bypass the sampler. */ + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @WithConfig(key = APPSEC_ENABLED, value = "true") + void propagatedTraceSourceAloneDoesNotBypassSampler() throws Exception { + // consume the first allowed slot + DDSpan first = (DDSpan) tracer.buildSpan("datadog", "op").start(); + first.finish(); + writer.waitForTraces(1); + + // second trace sets _dd.p.ts but NOT AI_GUARD_KEEP — should still be dropped + DDSpan second = (DDSpan) tracer.buildSpan("datadog", "op").start(); + second.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD); + second.finish(); + writer.waitForTraces(2); + + assertEquals(SAMPLER_DROP, (int) second.getSamplingPriority()); + } + + /** + * AI Guard traces must be kept even when AsmStandaloneSampler has exhausted its rate-limit slot. + */ + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @WithConfig(key = APPSEC_ENABLED, value = "true") + void aiGuardTraceIsKeptAfterRateLimitSlotIsExhausted() throws Exception { + // consume the first allowed slot + DDSpan first = (DDSpan) tracer.buildSpan("datadog", "op").start(); + first.finish(); + writer.waitForTraces(1); + + // AI Guard trace must still be force-kept even though the rate-limit slot is gone + DDSpan aiGuard = (DDSpan) tracer.buildSpan("datadog", "op").start(); + aiGuard.setTag(Tags.AI_GUARD_KEEP, true); + aiGuard.setTag(Tags.AI_GUARD_EVENT, true); + aiGuard.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD); + aiGuard.finish(); + writer.waitForTraces(2); + + assertEquals(USER_KEEP, (int) aiGuard.getSamplingPriority()); + assertDecisionMakerIsAiGuard(aiGuard); + } + + private static void assertDecisionMakerIsAiGuard(DDSpan span) { + Map ptags = span.spanContext().getPropagationTags().createTagMap(); + assertEquals("-13", ptags.get("_dd.p.dm"), "_dd.p.dm must be -13 (AI Guard decision maker)"); + assertEquals("20", ptags.get("_dd.p.ts"), "_dd.p.ts must have AI Guard bit set (0x20)"); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/taginterceptor/TagInterceptorTest.java b/dd-trace-core/src/test/java/datadog/trace/core/taginterceptor/TagInterceptorTest.java index 6f1e714334e..6c2a822971d 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/taginterceptor/TagInterceptorTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/taginterceptor/TagInterceptorTest.java @@ -389,6 +389,13 @@ static Stream setPrioritySamplingViaTagArguments() { arguments("asm.keep / 'true'", Tags.ASM_KEEP, "true", (int) PrioritySampling.USER_KEEP), arguments("asm.keep / 'false'", Tags.ASM_KEEP, "false", null), arguments("asm.keep / 'asdf'", Tags.ASM_KEEP, "asdf", null), + arguments( + "ai_guard.keep / true", Tags.AI_GUARD_KEEP, true, (int) PrioritySampling.USER_KEEP), + arguments("ai_guard.keep / false", Tags.AI_GUARD_KEEP, false, null), + arguments( + "ai_guard.keep / 'true'", Tags.AI_GUARD_KEEP, "true", (int) PrioritySampling.USER_KEEP), + arguments("ai_guard.keep / 'false'", Tags.AI_GUARD_KEEP, "false", null), + arguments("ai_guard.keep / 'asdf'", Tags.AI_GUARD_KEEP, "asdf", null), arguments( "sampling.priority / -1", Tags.SAMPLING_PRIORITY, -1, (int) PrioritySampling.USER_DROP), arguments( diff --git a/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java b/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java index 792829dde3a..cd4b0d2e8f4 100644 --- a/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java +++ b/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java @@ -22,6 +22,7 @@ public class ProductTraceSource { public static final int DSM = 0x04; public static final int DJM = 0x08; public static final int DBM = 0x10; + public static final int AI_GUARD = 0x20; /** Updates the bitfield by setting the bit corresponding to a specific product. */ public static int updateProduct(int bitfield, int product) { @@ -30,7 +31,15 @@ public static int updateProduct(int bitfield, int product) { /** Checks if the bitfield is marked for a specific product. */ public static boolean isProductMarked(final int bitfield, int product) { - return (bitfield & product) != 0; // Check if the bit is set + return (bitfield & product) != 0; + } + + /** Checks if the bitfield is marked for any of the given products. */ + public static boolean isProductMarked(final int bitfield, int... products) { + for (int product : products) { + if ((bitfield & product) != 0) return true; + } + return false; } /** diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java index 14496e8b243..782e63fa1ec 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java @@ -177,6 +177,9 @@ public class Tags { /** AI Guard force tracer to keep the trace */ public static final String AI_GUARD_KEEP = "ai_guard.keep"; + /** Marks a span as originating from an AI Guard evaluation */ + public static final String AI_GUARD_EVENT = "ai_guard.event"; + public static final String PROPAGATED_TRACE_SOURCE = "_dd.p.ts"; public static final String PROPAGATED_DEBUG = "_dd.p.debug"; diff --git a/internal-api/src/test/groovy/datadog/trace/api/ProductTraceSourceTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/ProductTraceSourceTest.groovy index bc376154217..2e5f1d54843 100644 --- a/internal-api/src/test/groovy/datadog/trace/api/ProductTraceSourceTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/api/ProductTraceSourceTest.groovy @@ -28,6 +28,23 @@ class ProductTraceSourceTest extends DDSpecification { value | product | expected ProductTraceSource.ASM | ProductTraceSource.ASM | true ProductTraceSource.DSM | ProductTraceSource.ASM | false + ProductTraceSource.AI_GUARD | ProductTraceSource.AI_GUARD | true + ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | false + } + + void 'test isProductMarked with varargs'(){ + when: + final result = ProductTraceSource.isProductMarked(value, products as int[]) + + then: + result == expected + + where: + value | products | expected + ProductTraceSource.ASM | [ProductTraceSource.ASM, ProductTraceSource.AI_GUARD] | true + ProductTraceSource.AI_GUARD | [ProductTraceSource.ASM, ProductTraceSource.AI_GUARD] | true + ProductTraceSource.DSM | [ProductTraceSource.ASM, ProductTraceSource.AI_GUARD] | false + ProductTraceSource.UNSET | [ProductTraceSource.ASM, ProductTraceSource.AI_GUARD] | false } void 'test getBitfieldHex'(){ @@ -41,6 +58,7 @@ class ProductTraceSourceTest extends DDSpecification { value | expected ProductTraceSource.UNSET | "00" ProductTraceSource.ASM | "02" + ProductTraceSource.AI_GUARD | "20" } void 'test parseBitfieldHex'(){ @@ -56,5 +74,6 @@ class ProductTraceSourceTest extends DDSpecification { null | ProductTraceSource.UNSET "" | ProductTraceSource.UNSET "02" | ProductTraceSource.ASM + "20" | ProductTraceSource.AI_GUARD } }