Set _dd.p.dm to -5 for IAST and AppSec spans (#7219)

What Does This Do Add new asm.keep tag to be able to set PrioritySampling.USER_KEEP with SamplingMechanism.APPSEC Motivation Current implementation is using manual.keep tag to force keep spans, this is not correct as is setting a SamplingMechanism.MANUAL (4) instead of SamplingMechanism.APPSEC(5)
DataDog · Jun 20, 2024 · f6d95fd · f6d95fd
1 parent 496f0ca
commit f6d95fd
Show file tree

Hide file tree

Showing 14 changed files with 41 additions and 23 deletions.
diff --git a/...main/java/datadog/trace/bootstrap/instrumentation/decorator/AppSecUserEventDecorator.java b/...main/java/datadog/trace/bootstrap/instrumentation/decorator/AppSecUserEventDecorator.java
@@ -3,11 +3,11 @@
 import static datadog.trace.api.UserEventTrackingMode.DISABLED;
 
 import datadog.trace.api.Config;
-import datadog.trace.api.DDTags;
 import datadog.trace.api.UserEventTrackingMode;
 import datadog.trace.api.internal.TraceSegment;
 import datadog.trace.bootstrap.ActiveSubsystems;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
 import java.util.Map;
 import javax.annotation.Nonnull;
 
@@ -74,7 +74,7 @@ public void onSignup(String userId, Map<String, String> metadata) {
 
   private void onEvent(@Nonnull TraceSegment segment, String eventName, Map<String, String> tags) {
     segment.setTagTop("appsec.events." + eventName + ".track", true, true);
-    segment.setTagTop(DDTags.MANUAL_KEEP, true);
+    segment.setTagTop(Tags.ASM_KEEP, true);
 
     // Report user event tracking mode ("safe" or "extended")
     UserEventTrackingMode mode = Config.get().getAppSecUserEventsTrackingMode();

diff --git a/...ovy/datadog/trace/bootstrap/instrumentation/decorator/AppSecUserEventDecoratorTest.groovy b/...ovy/datadog/trace/bootstrap/instrumentation/decorator/AppSecUserEventDecoratorTest.groovy
@@ -32,7 +32,7 @@ class AppSecUserEventDecoratorTest extends DDSpecification {
     then:
     1 * traceSegment.setTagTop('_dd.appsec.events.users.signup.auto.mode', modeTag)
     1 * traceSegment.setTagTop('appsec.events.users.signup.track', true, true)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop('usr.id', 'user')
     1 * traceSegment.setTagTop('appsec.events.users.signup', ['key1':'value1', 'key2':'value2'])
     0 * _
@@ -54,7 +54,7 @@ class AppSecUserEventDecoratorTest extends DDSpecification {
     then:
     1 * traceSegment.setTagTop('_dd.appsec.events.users.login.success.auto.mode', modeTag)
     1 * traceSegment.setTagTop('appsec.events.users.login.success.track', true, true)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop('usr.id', 'user')
     1 * traceSegment.setTagTop('appsec.events.users.login.success', ['key1':'value1', 'key2':'value2'])
     0 * _
@@ -76,7 +76,7 @@ class AppSecUserEventDecoratorTest extends DDSpecification {
     then:
     1 * traceSegment.setTagTop('_dd.appsec.events.users.login.failure.auto.mode', modeTag)
     1 * traceSegment.setTagTop('appsec.events.users.login.failure.track', true, true)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop('appsec.events.users.login.failure.usr.id', 'user')
     1 * traceSegment.setTagTop('appsec.events.users.login.failure', ['key1':'value1', 'key2':'value2'])
     0 * _

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/Reporter.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/Reporter.java
@@ -7,7 +7,6 @@
 import com.datadog.iast.model.VulnerabilityBatch;
 import com.datadog.iast.taint.TaintedObjects;
 import datadog.trace.api.Config;
-import datadog.trace.api.DDTags;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.api.internal.TraceSegment;
@@ -98,7 +97,7 @@ private VulnerabilityBatch getOrCreateVulnerabilityBatch(final AgentSpan span) {
       // Once we have added a vulnerability, try to override sampling and keep the trace.
       // TODO: We need to check if we can have an API with more fine-grained semantics on why traces
       // are kept.
-      segment.setTagTop(DDTags.MANUAL_KEEP, true);
+      segment.setTagTop(Tags.ASM_KEEP, true);
       return batch;
     }
 

diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/ReporterTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/ReporterTest.groovy
@@ -79,7 +79,7 @@ class ReporterTest extends DDSpecification {
         }
       ]
     }''', batch.toString(), true)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     0 * _
   }
 
@@ -145,7 +145,7 @@ class ReporterTest extends DDSpecification {
         }
       ]
     }''', batch.toString(), true)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     0 * _
   }
 
@@ -268,7 +268,7 @@ class ReporterTest extends DDSpecification {
     1 * reqCtx.getTraceSegment() >> traceSegment
     1 * traceSegment.getDataTop('iast') >> null
     1 * traceSegment.setDataTop('iast', _ as VulnerabilityBatch)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop('_dd.iast.enabled', 1)
     0 * _
   }

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
@@ -22,7 +22,6 @@
 import com.datadog.appsec.stack_trace.StackTraceCollection;
 import com.datadog.appsec.util.ObjectFlattener;
 import datadog.trace.api.Config;
-import datadog.trace.api.DDTags;
 import datadog.trace.api.function.TriConsumer;
 import datadog.trace.api.function.TriFunction;
 import datadog.trace.api.gateway.Events;
@@ -147,7 +146,7 @@ public void init() {
               if (!collectedEvents.isEmpty()) {
                 // Keep event related span, because it could be ignored in case of
                 // reduced datadog sampling rate.
-                traceSeg.setTagTop(DDTags.MANUAL_KEEP, true);
+                traceSeg.setTagTop(Tags.ASM_KEEP, true);
                 traceSeg.setTagTop("appsec.event", true);
                 traceSeg.setTagTop("network.client.ip", ctx.getPeerAddress());
 

diff --git a/...ppsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFInitializationResultReporter.java b/...ppsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFInitializationResultReporter.java
@@ -6,8 +6,8 @@
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.Moshi;
 import com.squareup.moshi.Types;
-import datadog.trace.api.DDTags;
 import datadog.trace.api.internal.TraceSegment;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
 import io.sqreen.powerwaf.Powerwaf;
 import io.sqreen.powerwaf.RuleSetInfo;
 import java.util.Collection;
@@ -50,6 +50,6 @@ public void processTraceSegment(
     segment.setTagTop(RULE_ERROR_COUNT, report.getNumRulesError());
     segment.setTagTop(WAF_VERSION, Powerwaf.LIB_VERSION);
 
-    segment.setTagTop(DDTags.MANUAL_KEEP, true);
+    segment.setTagTop(Tags.ASM_KEEP, true);
   }
 }
diff --git a/...agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/...agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy
@@ -139,7 +139,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * mockAppSecCtx.transferCollectedEvents() >> [event]
     1 * mockAppSecCtx.peerAddress >> '2001::1'
     1 * mockAppSecCtx.close()
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop("_dd.appsec.enabled", 1)
     1 * traceSegment.setTagTop("_dd.runtime_family", "jvm")
     1 * traceSegment.setTagTop('appsec.event', true)

diff --git a/...ent/appsec/src/test/groovy/com/datadog/appsec/powerwaf/PowerWAFModuleSpecification.groovy b/...ent/appsec/src/test/groovy/com/datadog/appsec/powerwaf/PowerWAFModuleSpecification.groovy
@@ -646,7 +646,7 @@ class PowerWAFModuleSpecification extends DDSpecification {
     1 * segment.setTagTop('_dd.appsec.event_rules.loaded', 115)
     1 * segment.setTagTop('_dd.appsec.event_rules.error_count', 1)
     1 * segment.setTagTop('_dd.appsec.event_rules.errors', { it =~ /\{"[^"]+":\["bad rule"\]\}/})
-    1 * segment.setTagTop('manual.keep', true)
+    1 * segment.setTagTop('asm.keep', true)
     0 * segment._(*_)
 
     when:

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/EventTracker.java b/dd-trace-api/src/main/java/datadog/trace/api/EventTracker.java
@@ -33,7 +33,7 @@ public void trackLoginSuccessEvent(String userId, Map<String, String> metadata)
     segment.setTagTop("_dd.appsec.events.users.login.success.sdk", true);
     segment.setTagTop("appsec.events.users.login.success.track", true);
     segment.setTagTop("usr.id", userId);
-    segment.setTagTop(DDTags.MANUAL_KEEP, true);
+    segment.setTagTop("asm.keep", true);
 
     if (metadata != null && !metadata.isEmpty()) {
       segment.setTagTop("appsec.events.users.login.success", metadata);
@@ -66,7 +66,7 @@ public void trackLoginFailureEvent(String userId, boolean exists, Map<String, St
     segment.setTagTop("appsec.events.users.login.failure.track", true);
     segment.setTagTop("appsec.events.users.login.failure.usr.id", userId);
     segment.setTagTop("appsec.events.users.login.failure.usr.exists", exists);
-    segment.setTagTop(DDTags.MANUAL_KEEP, true);
+    segment.setTagTop("asm.keep", true);
 
     if (metadata != null && !metadata.isEmpty()) {
       segment.setTagTop("appsec.events.users.login.failure", metadata);
@@ -96,7 +96,7 @@ public void trackCustomEvent(String eventName, Map<String, String> metadata) {
 
     segment.setTagTop("_dd.appsec.events." + eventName + ".sdk", true);
     segment.setTagTop("appsec.events." + eventName + ".track", true, true);
-    segment.setTagTop(DDTags.MANUAL_KEEP, true);
+    segment.setTagTop("asm.keep", true);
 
     if (metadata != null && !metadata.isEmpty()) {
       segment.setTagTop("appsec.events." + eventName, metadata, true);

diff --git a/dd-trace-api/src/test/groovy/datadog/trace/api/EventTrackerTest.groovy b/dd-trace-api/src/test/groovy/datadog/trace/api/EventTrackerTest.groovy
@@ -33,7 +33,7 @@ class EventTrackerTest extends DDSpecification {
     1 * traceSegment.setTagTop('appsec.events.users.login.success.track', true)
     1 * traceSegment.setTagTop('_dd.appsec.events.users.login.success.sdk', true)
     1 * traceSegment.setTagTop('usr.id', 'user1')
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop('appsec.events.users.login.success', ['key1':'value1', 'key2':'value2'])
     0 * _
   }
@@ -47,7 +47,7 @@ class EventTrackerTest extends DDSpecification {
     1 * traceSegment.setTagTop('_dd.appsec.events.users.login.failure.sdk', true)
     1 * traceSegment.setTagTop('appsec.events.users.login.failure.usr.id', 'user1')
     1 * traceSegment.setTagTop('appsec.events.users.login.failure.usr.exists', true)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop('appsec.events.users.login.failure', ['key1':'value1', 'key2':'value2'])
     0 * _
   }
@@ -59,7 +59,7 @@ class EventTrackerTest extends DDSpecification {
     then:
     1 * traceSegment.setTagTop('appsec.events.myevent.track', true, true)
     1 * traceSegment.setTagTop('_dd.appsec.events.myevent.sdk', true)
-    1 * traceSegment.setTagTop('manual.keep', true)
+    1 * traceSegment.setTagTop('asm.keep', true)
     1 * traceSegment.setTagTop('appsec.events.myevent', ['key1':'value1', 'key2':'value2'], true)
     0 * _
   }

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java
@@ -492,9 +492,14 @@ public void setSpanType(final CharSequence spanType) {
     this.spanType = spanType;
   }
 
+  /** Forces the local root span sampling decision to keep according manual mechanism. */
   public void forceKeep() {
+    forceKeep(SamplingMechanism.MANUAL);
+  }
+
+  public void forceKeep(byte samplingMechanism) {
     // set trace level sampling priority
-    getRootSpanContextOrThis().forceKeepThisSpan(SamplingMechanism.MANUAL);
+    getRootSpanContextOrThis().forceKeepThisSpan(samplingMechanism);
   }
 
   private void forceKeepThisSpan(byte samplingMechanism) {

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/taginterceptor/TagInterceptor.java b/dd-trace-core/src/main/java/datadog/trace/core/taginterceptor/TagInterceptor.java
@@ -100,6 +100,12 @@ public boolean interceptTag(DDSpanContext span, String tag, Object value) {
       case DDTags.MANUAL_DROP:
         return interceptSamplingPriority(
             FORCE_MANUAL_DROP, USER_DROP, SamplingMechanism.MANUAL, span, value);
+      case Tags.ASM_KEEP:
+        if (asBoolean(value)) {
+          span.forceKeep(SamplingMechanism.APPSEC);
+          return true;
+        }
+        return false;
       case Tags.SAMPLING_PRIORITY:
         return interceptSamplingPriority(span, value);
       case InstrumentationTags.SERVLET_CONTEXT:

diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/taginterceptor/TagInterceptorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/taginterceptor/TagInterceptorTest.groovy
@@ -421,6 +421,12 @@ class TagInterceptorTest extends DDCoreSpecification {
     DDTags.MANUAL_DROP     | "false" | null
     DDTags.MANUAL_DROP     | "asdf"  | null
 
+    Tags.ASM_KEEP          | true    | PrioritySampling.USER_KEEP
+    Tags.ASM_KEEP          | false   | null
+    Tags.ASM_KEEP          | "true"  | PrioritySampling.USER_KEEP
+    Tags.ASM_KEEP          | "false" | null
+    Tags.ASM_KEEP          | "asdf"  | null
+
     Tags.SAMPLING_PRIORITY | -1      | PrioritySampling.USER_DROP
     Tags.SAMPLING_PRIORITY | 0       | PrioritySampling.USER_DROP
     Tags.SAMPLING_PRIORITY | 1       | PrioritySampling.USER_KEEP

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
@@ -126,4 +126,7 @@ public class Tags {
   public static final String DD_ENV = "dd.env";
 
   public static final String ENV = "env";
+
+  /** ASM force tracer to keep the trace */
+  public static final String ASM_KEEP = "asm.keep";
 }