Move on the fly span creation for vulns out of req to addVulnerabilit…

…y method
DataDog · Nov 26, 2024 · 40da9e1 · 40da9e1
1 parent 8a16b64
commit 40da9e1
Showing 1 changed file with 20 additions and 22 deletions.
diff --git a/packages/dd-trace/src/appsec/iast/vulnerability-reporter.js b/packages/dd-trace/src/appsec/iast/vulnerability-reporter.js
@@ -22,14 +22,28 @@ function addVulnerability (iastContext, vulnerability) {
 
     VULNERABILITY_HASHES.set(`${vulnerability.type}${vulnerability.hash}`, true)
 
-    if (iastContext?.rootSpan) {
-      keepTrace(iastContext.rootSpan, SAMPLING_MECHANISM_APPSEC)
-      standalone.sample(iastContext.rootSpan)
+    let span = iastContext?.rootSpan
+
+    if (!span && tracer) {
+      span = tracer.startSpan('vulnerability', {
+        type: 'vulnerability'
+      })
+
+      vulnerability.location.spanId = span.context().toSpanId()
+
+      span.addTags({
+        [IAST_ENABLED_TAG_KEY]: 1
+      })
+    }
+
+    keepTrace(span, SAMPLING_MECHANISM_APPSEC)
+    standalone.sample(span)
 
+    if (iastContext?.rootSpan) {
       iastContext[VULNERABILITIES_KEY] = iastContext[VULNERABILITIES_KEY] || []
       iastContext[VULNERABILITIES_KEY].push(vulnerability)
     } else {
-      sendVulnerabilities([vulnerability])
+      sendVulnerabilities([vulnerability], span, false)
     }
   }
 }
@@ -40,24 +54,8 @@ function isValidVulnerability (vulnerability) {
     vulnerability.location && vulnerability.location.spanId
 }
 
-function sendVulnerabilities (vulnerabilities, rootSpan) {
+function sendVulnerabilities (vulnerabilities, span, isRootSpan = true) {
   if (vulnerabilities && vulnerabilities.length) {
-    let span = rootSpan
-    if (!span && tracer) {
-      span = tracer.startSpan('vulnerability', {
-        type: 'vulnerability'
-      })
-      vulnerabilities.forEach((vulnerability) => {
-        vulnerability.location.spanId = span.context().toSpanId()
-      })
-      span.addTags({
-        [IAST_ENABLED_TAG_KEY]: 1
-      })
-
-      keepTrace(span, SAMPLING_MECHANISM_APPSEC)
-      standalone.sample(span)
-    }
-
     if (span && span.addTags) {
       const validatedVulnerabilities = vulnerabilities.filter(isValidVulnerability)
       const jsonToSend = vulnerabilitiesFormatter.toJson(validatedVulnerabilities)
@@ -68,7 +66,7 @@ function sendVulnerabilities (vulnerabilities, rootSpan) {
         tags[IAST_JSON_TAG_KEY] = JSON.stringify(jsonToSend)
         span.addTags(tags)
 
-        if (!rootSpan) span.finish()
+        if (!isRootSpan) span.finish()
       }
     }
   }