From 15469d6cb2e4ce8156d2312a57d9575a6e3153f2 Mon Sep 17 00:00:00 2001
From: Santiago Mola <santiago.mola@datadoghq.com>
Date: Fri, 16 Sep 2022 14:09:44 +0200
Subject: [PATCH] [IAST] Add manual.keep tag when a vulnerability is reported

---
 .../iast/src/main/java/com/datadog/iast/Reporter.java    | 9 ++++++++-
 .../src/test/groovy/com/datadog/iast/ReporterTest.groovy | 2 ++
 2 files changed, 10 insertions(+), 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
index 120ebd99a6a..2cdd861cde0 100644
--- 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
index 67446896af9..49a357ed0a7 100644
--- 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 * _
   }