[IAST] Add manual.keep tag when a vulnerability is reported

DataDog · Sep 23, 2022 · 15469d6 · 15469d6
1 parent c81fa9b
commit 15469d6
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/dd-java-agent/iast/src/main/java/com/datadog/iast/Reporter.java b/dd-java-agent/iast/src/main/java/com/datadog/iast/Reporter.java
@@ -2,6 +2,8 @@
 
 import com.datadog.iast.model.Vulnerability;
 import com.datadog.iast.model.VulnerabilityBatch;
+import datadog.trace.api.DDTags;
+import datadog.trace.api.TraceSegment;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -26,7 +28,12 @@ public void report(final AgentSpan span, final Vulnerability vulnerability) {
     final VulnerabilityBatch batch = ctx.getVulnerabilityBatch();
     batch.add(vulnerability);
     if (!ctx.getAndSetSpanDataIsSet()) {
-      reqCtx.getTraceSegment().setDataTop("iast", batch);
+      final TraceSegment segment = reqCtx.getTraceSegment();
+      segment.setDataTop("iast", batch);
+      // 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);
     }
   }
 }
diff --git a/dd-java-agent/iast/src/test/groovy/com/datadog/iast/ReporterTest.groovy b/dd-java-agent/iast/src/test/groovy/com/datadog/iast/ReporterTest.groovy
@@ -38,6 +38,7 @@ class ReporterTest extends DDSpecification {
     then:
     1 * traceSegment.setDataTop('iast', _) >> { batch = it[1] as VulnerabilityBatch }
     batch.toString() == '{"vulnerabilities":[{"evidence":{"value":"MD5"},"location":{"line":1,"path":"foo"},"type":"WEAK_HASH"}]}'
+    1 * traceSegment.setTagTop('manual.keep', true)
     0 * _
   }
 
@@ -72,6 +73,7 @@ class ReporterTest extends DDSpecification {
     then:
     1 * traceSegment.setDataTop('iast', _) >> { batch = it[1] as VulnerabilityBatch }
     batch.toString() == '{"vulnerabilities":[{"evidence":{"value":"MD5"},"location":{"line":1,"path":"foo"},"type":"WEAK_HASH"},{"evidence":{"value":"MD4"},"location":{"line":1,"path":"foo"},"type":"WEAK_HASH"}]}'
+    1 * traceSegment.setTagTop('manual.keep', true)
     0 * _
   }