From 449af83b85b2e1ba552c2de4348a2a946f110693 Mon Sep 17 00:00:00 2001
From: "alejandro.gonzalez" <alejandro.gonzalez@datadoghq.com>
Date: Tue, 18 Feb 2025 12:33:00 +0100
Subject: [PATCH] fix addVulnerabilityStackTrace in reportVulnerability when
 multiple vulnerabilities share location

---
 .../main/java/com/datadog/iast/Reporter.java  |  4 +-
 .../com/datadog/iast/ReporterTest.groovy      | 74 +++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

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
index 1b03c4ec121..370d111bc65 100644
--- 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
@@ -77,7 +77,9 @@ private void reportVulnerability(
     final VulnerabilityBatch batch = getOrCreateVulnerabilityBatch(span);
     if (batch != null) {
       batch.add(vulnerability);
-      if (Config.get().isIastStackTraceEnabled() && batch.getVulnerabilities() != null) {
+      if (Config.get().isIastStackTraceEnabled()
+          && batch.getVulnerabilities() != null
+          && vulnerability.getLocation().getStackId() == null) {
         String stackId =
             addVulnerabilityStackTrace(span, String.valueOf(batch.getVulnerabilities().size()));
         if (stackId != null) {
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
index bfa8e2d3e51..c162faadcfe 100644
--- 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
@@ -533,6 +533,80 @@ class ReporterTest extends DDSpecification {
     batch.vulnerabilities.size() == 2
   }
 
+  void 'two cookie vulnerabilities that share location share stackId and only generate one stack'() {
+    given:
+    final Reporter reporter = new Reporter()
+    final traceSegment = Mock(TraceSegment)
+    final ctx = new IastRequestContext(noOpTaintedObjects())
+    final reqCtx = Mock(RequestContext)
+    final spanId = 123456
+    VulnerabilityBatch batch = null
+    Map stackTraceBatch = new ConcurrentHashMap()
+
+    final span = Stub(AgentSpan)
+    span.getRequestContext() >> reqCtx
+    span.getSpanId() >> spanId
+
+    final location = Location.forSpanAndStack(span, new StackTraceElement("foo", "foo", "foo", 1))
+
+    final v1 = new Vulnerability(
+      VulnerabilityType.INSECURE_COOKIE,
+      location,
+      new Evidence("JSESSIONID")
+      )
+    final v2 = new Vulnerability(
+      VulnerabilityType.NO_SAMESITE_COOKIE,
+      location,
+      new Evidence("JSESSIONID")
+      )
+
+    when:
+    reporter.report(span, v1)
+    reporter.report(span, v2)
+
+    then:
+    1 * reqCtx.getData(RequestContextSlot.IAST) >> ctx
+    2 * reqCtx.getTraceSegment() >> traceSegment
+    // first vulnerability
+    1 * traceSegment.getDataTop('iast') >> null
+    1 * traceSegment.setDataTop('iast', _) >> { batch = it[1] as VulnerabilityBatch }
+    1 * reqCtx.getOrCreateMetaStructTop('_dd.stack', _)  >> { stackTraceBatch }
+    // second vulnerability
+    1 * traceSegment.getDataTop('iast') >> { return batch } // second vulnerability
+    JSONAssert.assertEquals('''{
+      "vulnerabilities": [
+        {
+          "evidence": { "value":"JSESSIONID" },
+          "hash":1156210466,
+          "location": {
+            "spanId":123456,
+            "line":1,
+            "path":"foo",
+            "method": "foo",
+            "stackId":"1"
+          },
+          "type":"INSECURE_COOKIE"
+        },
+        {
+          "evidence": {"value":"JSESSIONID"},
+          "hash":1090504969,
+          "location": {
+            "spanId":123456,
+            "line":1,
+            "path":"foo",
+            "method": "foo",
+            "stackId":"1"
+          },
+          "type":"NO_SAMESITE_COOKIE"
+        }
+      ]
+    }''', batch.toString(), true)
+    1 * traceSegment.setTagTop('asm.keep', true)
+    1 * traceSegment.setTagTop('_dd.p.ts', ProductTraceSource.ASM)
+    assert  stackTraceBatch.get("vulnerability").size() == 1
+    0 * _
+  }
+
   private AgentSpan spanWithBatch(final VulnerabilityBatch batch) {
     final traceSegment = Mock(TraceSegment) {
       getDataTop('iast') >> batch