Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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";
Expand Down Expand Up @@ -275,7 +276,8 @@ public Evaluation evaluate(final List<Message> 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);
Comment thread
claponcet marked this conversation as resolved.
localRootSpan.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
applyClientIpTags(localRootSpan);
// copyAnomalyDetectionTags MUST run after applyClientIpTags, to make
// sure client IP tags were populated.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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')
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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<String, String> 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)");
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,13 @@ static Stream<Arguments> 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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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;
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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'(){
Expand All @@ -41,6 +58,7 @@ class ProductTraceSourceTest extends DDSpecification {
value | expected
ProductTraceSource.UNSET | "00"
ProductTraceSource.ASM | "02"
ProductTraceSource.AI_GUARD | "20"
}

void 'test parseBitfieldHex'(){
Expand All @@ -56,5 +74,6 @@ class ProductTraceSourceTest extends DDSpecification {
null | ProductTraceSource.UNSET
"" | ProductTraceSource.UNSET
"02" | ProductTraceSource.ASM
"20" | ProductTraceSource.AI_GUARD
}
}