dd-trace 5.17.0 → 5.18.0

package/packages/datadog-plugin-aws-sdk/src/services/sns.js

@@ -59,7 +59,12 @@ class Sns extends BaseAwsSdkPlugin {
         break
       case 'publishBatch':
         for (let i = 0; i < params.PublishBatchRequestEntries.length; i++) {
-          this.injectToMessage(span, params.PublishBatchRequestEntries[i], params.TopicArn, i === 0)
+          this.injectToMessage(
+            span,
+            params.PublishBatchRequestEntries[i],
+            params.TopicArn,
+            i === 0 || (this.config.sns && this.config.sns.batchPropagationEnabled)
+          )
         }
         break
     }

package/packages/datadog-plugin-aws-sdk/src/services/sqs.js

@@ -23,7 +23,7 @@ class Sqs extends BaseAwsSdkPlugin {
       const plugin = this
       const contextExtraction = this.responseExtract(request.params, request.operation, response)
       let span
-      let parsedMessageAttributes
+      let parsedMessageAttributes = null
       if (contextExtraction && contextExtraction.datadogContext) {
         obj.needsFinish = true
         const options = {
@@ -39,8 +39,9 @@ class Sqs extends BaseAwsSdkPlugin {
         this.enter(span, store)
       }
       // extract DSM context after as we might not have a parent-child but may have a DSM context
+
       this.responseExtractDSMContext(
-        request.operation, request.params, response, span || null, parsedMessageAttributes || null
+        request.operation, request.params, response, span || null, { parsedMessageAttributes }
       )
     })
 
@@ -165,7 +166,8 @@ class Sqs extends BaseAwsSdkPlugin {
     }
   }
 
-  responseExtractDSMContext (operation, params, response, span, parsedAttributes) {
+  responseExtractDSMContext (operation, params, response, span, kwargs = {}) {
+    let { parsedAttributes } = kwargs
     if (!this.config.dsmEnabled) return
     if (operation !== 'receiveMessage') return
     if (!response || !response.Messages || !response.Messages[0]) return
@@ -188,7 +190,7 @@ class Sqs extends BaseAwsSdkPlugin {
             // SQS to SQS
           }
         }
-        if (message.MessageAttributes && message.MessageAttributes._datadog) {
+        if (!parsedAttributes && message.MessageAttributes && message.MessageAttributes._datadog) {
           parsedAttributes = this.parseDatadogAttributes(message.MessageAttributes._datadog)
         }
       }
@@ -216,7 +218,23 @@ class Sqs extends BaseAwsSdkPlugin {
         break
       case 'sendMessageBatch':
         for (let i = 0; i < params.Entries.length; i++) {
-          this.injectToMessage(span, params.Entries[i], params.QueueUrl, i === 0)
+          this.injectToMessage(
+            span,
+            params.Entries[i],
+            params.QueueUrl,
+            i === 0 || (this.config.sqs && this.config.sqs.batchPropagationEnabled)
+          )
+        }
+        break
+      case 'receiveMessage':
+        if (!params.MessageAttributeNames) {
+          params.MessageAttributeNames = ['_datadog']
+        } else if (
+          !params.MessageAttributeNames.includes('_datadog') &&
+          !params.MessageAttributeNames.includes('.*') &&
+          !params.MessageAttributeNames.includes('All')
+        ) {
+          params.MessageAttributeNames.push('_datadog')
         }
         break
     }

package/packages/datadog-plugin-child_process/src/index.js

@@ -54,7 +54,7 @@ class ChildProcessPlugin extends TracingPlugin {
     }
 
     this.startSpan('command_execution', {
-      service: this.config.service,
+      service: this.config.service || this._tracerConfig.service,
       resource: (shell === true) ? 'sh' : cmdFields[0],
       type: 'system',
       meta

package/packages/datadog-plugin-cucumber/src/index.js

@@ -195,6 +195,17 @@ class CucumberPlugin extends CiPlugin {
       this.enter(testSpan, store)
     })
 
+    this.addSub('ci:cucumber:test:retry', (isFlakyRetry) => {
+      const store = storage.getStore()
+      const span = store.span
+      if (isFlakyRetry) {
+        span.setTag(TEST_IS_RETRY, 'true')
+      }
+      span.setTag(TEST_STATUS, 'fail')
+      span.finish()
+      finishAllTraceSpans(span)
+    })
+
     this.addSub('ci:cucumber:test-step:start', ({ resource }) => {
       const store = storage.getStore()
       const childOf = store ? store.span : store
@@ -239,7 +250,15 @@ class CucumberPlugin extends CiPlugin {
       })
     })
 
-    this.addSub('ci:cucumber:test:finish', ({ isStep, status, skipReason, errorMessage, isNew, isEfdRetry }) => {
+    this.addSub('ci:cucumber:test:finish', ({
+      isStep,
+      status,
+      skipReason,
+      errorMessage,
+      isNew,
+      isEfdRetry,
+      isFlakyRetry
+    }) => {
       const span = storage.getStore().span
       const statusTag = isStep ? 'step.status' : TEST_STATUS
 
@@ -260,6 +279,10 @@ class CucumberPlugin extends CiPlugin {
         span.setTag(ERROR_MESSAGE, errorMessage)
       }
 
+      if (isFlakyRetry > 0) {
+        span.setTag(TEST_IS_RETRY, 'true')
+      }
+
       span.finish()
       if (!isStep) {
         this.telemetry.ciVisEvent(

package/packages/datadog-plugin-mocha/src/index.js

@@ -175,13 +175,15 @@ class MochaPlugin extends CiPlugin {
       this.tracer._exporter.flush()
     })
 
-    this.addSub('ci:mocha:test:finish', (status) => {
+    this.addSub('ci:mocha:test:finish', ({ status, hasBeenRetried }) => {
       const store = storage.getStore()
       const span = store?.span
 
       if (span) {
         span.setTag(TEST_STATUS, status)
-
+        if (hasBeenRetried) {
+          span.setTag(TEST_IS_RETRY, 'true')
+        }
         span.finish()
         this.telemetry.ciVisEvent(
           TELEMETRY_EVENT_FINISHED,
@@ -204,8 +206,8 @@ class MochaPlugin extends CiPlugin {
 
     this.addSub('ci:mocha:test:error', (err) => {
       const store = storage.getStore()
-      if (err && store && store.span) {
-        const span = store.span
+      const span = store?.span
+      if (err && span) {
         if (err.constructor.name === 'Pending' && !this.forbidPending) {
           span.setTag(TEST_STATUS, 'skip')
         } else {
@@ -215,6 +217,25 @@ class MochaPlugin extends CiPlugin {
       }
     })
 
+    this.addSub('ci:mocha:test:retry', (isFirstAttempt) => {
+      const store = storage.getStore()
+      const span = store?.span
+      if (span) {
+        span.setTag(TEST_STATUS, 'fail')
+        if (!isFirstAttempt) {
+          span.setTag(TEST_IS_RETRY, 'true')
+        }
+
+        span.finish()
+        this.telemetry.ciVisEvent(
+          TELEMETRY_EVENT_FINISHED,
+          'test',
+          { hasCodeOwners: !!span.context()._tags[TEST_CODE_OWNERS] }
+        )
+        finishAllTraceSpans(span)
+      }
+    })
+
     this.addSub('ci:mocha:test:parameterize', ({ title, params }) => {
       this._testTitleToParams[title] = params
     })

package/packages/datadog-plugin-openai/src/index.js

@@ -7,6 +7,7 @@ const { storage } = require('../../datadog-core')
 const services = require('./services')
 const Sampler = require('../../dd-trace/src/sampler')
 const { MEASURED } = require('../../../ext/tags')
+const { estimateTokens } = require('./token-estimator')
 
 // String#replaceAll unavailable on Node.js@v14 (dd-trace@<=v3)
 const RE_NEWLINE = /\n/g
@@ -15,14 +16,17 @@ const RE_TAB = /\t/g
 // TODO: In the future we should refactor config.js to make it requirable
 let MAX_TEXT_LEN = 128
 
-let encodingForModel
-try {
-  // eslint-disable-next-line import/no-extraneous-dependencies
-  encodingForModel = require('tiktoken').encoding_for_model
-} catch {
-  // we will use token count estimations in this case
+function safeRequire (path) {
+  try {
+    // eslint-disable-next-line import/no-extraneous-dependencies
+    return require(path)
+  } catch {
+    return null
+  }
 }
 
+const encodingForModel = safeRequire('tiktoken')?.encoding_for_model
+
 class OpenApiPlugin extends TracingPlugin {
   static get id () { return 'openai' }
   static get operation () { return 'request' }
@@ -305,6 +309,7 @@ class OpenApiPlugin extends TracingPlugin {
   }
 
   sendLog (methodName, span, tags, store, error) {
+    if (!store) return
     if (!Object.keys(store).length) return
     if (!this.sampler.isSampled()) return
 
@@ -325,9 +330,22 @@ function countPromptTokens (methodName, payload, model) {
     const messages = payload.messages
     for (const message of messages) {
       const content = message.content
-      const { tokens, estimated } = countTokens(content, model)
-      promptTokens += tokens
-      promptEstimated = estimated
+      if (typeof content === 'string') {
+        const { tokens, estimated } = countTokens(content, model)
+        promptTokens += tokens
+        promptEstimated = estimated
+      } else if (Array.isArray(content)) {
+        for (const c of content) {
+          if (c.type === 'text') {
+            const { tokens, estimated } = countTokens(c.text, model)
+            promptTokens += tokens
+            promptEstimated = estimated
+          }
+          // unsupported token computation for image_url
+          // as even though URL is a string, its true token count
+          // is based on the image itself, something onerous to do client-side
+        }
+      }
     }
   } else if (methodName === 'completions.create') {
     let prompt = payload.prompt
@@ -382,25 +400,6 @@ function countTokens (content, model) {
   }
 }
 
-// If model is unavailable or tiktoken is not imported, then provide a very rough estimate of the number of tokens
-// Approximate using the following assumptions:
-//    * English text
-//    * 1 token ~= 4 chars
-//    * 1 token ~= ¾ words
-function estimateTokens (content) {
-  let estimatedTokens = 0
-  if (typeof content === 'string') {
-    const estimation1 = content.length / 4
-
-    const matches = content.match(/[\w']+|[.,!?;~@#$%^&*()+/-]/g)
-    const estimation2 = matches ? matches.length * 0.75 : 0 // in the case of an empty string
-    estimatedTokens = Math.round((1.5 * estimation1 + 0.5 * estimation2) / 2)
-  } else if (Array.isArray(content) && typeof content[0] === 'number') {
-    estimatedTokens = content.length
-  }
-  return estimatedTokens
-}
-
 function createEditRequestExtraction (tags, payload, store) {
   const instruction = payload.instruction
   tags['openai.request.instruction'] = instruction
@@ -418,7 +417,7 @@ function createChatCompletionRequestExtraction (tags, payload, store) {
   store.messages = payload.messages
   for (let i = 0; i < payload.messages.length; i++) {
     const message = payload.messages[i]
-    tags[`openai.request.messages.${i}.content`] = truncateText(message.content)
+    tagChatCompletionRequestContent(message.content, i, tags)
     tags[`openai.request.messages.${i}.role`] = message.role
     tags[`openai.request.messages.${i}.name`] = message.name
     tags[`openai.request.messages.${i}.finish_reason`] = message.finish_reason
@@ -707,7 +706,7 @@ function commonCreateResponseExtraction (tags, body, store, methodName) {
   for (let choiceIdx = 0; choiceIdx < body.choices.length; choiceIdx++) {
     const choice = body.choices[choiceIdx]
 
-    // logprobs can be nullm and we still want to tag it as 'returned' even when set to 'null'
+    // logprobs can be null and we still want to tag it as 'returned' even when set to 'null'
     const specifiesLogProb = Object.keys(choice).indexOf('logprobs') !== -1
 
     tags[`openai.response.choices.${choiceIdx}.finish_reason`] = choice.finish_reason
@@ -781,6 +780,7 @@ function truncateApiKey (apiKey) {
  */
 function truncateText (text) {
   if (!text) return
+  if (typeof text !== 'string' || !text || (typeof text === 'string' && text.length === 0)) return
 
   text = text
     .replace(RE_NEWLINE, '\\n')
@@ -793,6 +793,28 @@ function truncateText (text) {
   return text
 }
 
+function tagChatCompletionRequestContent (contents, messageIdx, tags) {
+  if (typeof contents === 'string') {
+    tags[`openai.request.messages.${messageIdx}.content`] = contents
+  } else if (Array.isArray(contents)) {
+    // content can also be an array of objects
+    // which represent text input or image url
+    for (const contentIdx in contents) {
+      const content = contents[contentIdx]
+      const type = content.type
+      tags[`openai.request.messages.${messageIdx}.content.${contentIdx}.type`] = content.type
+      if (type === 'text') {
+        tags[`openai.request.messages.${messageIdx}.content.${contentIdx}.text`] = truncateText(content.text)
+      } else if (type === 'image_url') {
+        tags[`openai.request.messages.${messageIdx}.content.${contentIdx}.image_url.url`] =
+          truncateText(content.image_url.url)
+      }
+      // unsupported type otherwise, won't be tagged
+    }
+  }
+  // unsupported type otherwise, won't be tagged
+}
+
 // The server almost always responds with JSON
 function coerceResponseBody (body, methodName) {
   switch (methodName) {

package/packages/datadog-plugin-openai/src/token-estimator.js

@@ -0,0 +1,20 @@
+'use strict'
+
+// If model is unavailable or tiktoken is not imported, then provide a very rough estimate of the number of tokens
+// Approximate using the following assumptions:
+//    * English text
+//    * 1 token ~= 4 chars
+//    * 1 token ~= ¾ words
+module.exports.estimateTokens = function (content) {
+  let estimatedTokens = 0
+  if (typeof content === 'string') {
+    const estimation1 = content.length / 4
+
+    const matches = content.match(/[\w']+|[.,!?;~@#$%^&*()+/-]/g)
+    const estimation2 = matches ? matches.length * 0.75 : 0 // in the case of an empty string
+    estimatedTokens = Math.round((1.5 * estimation1 + 0.5 * estimation2) / 2)
+  } else if (Array.isArray(content) && typeof content[0] === 'number') {
+    estimatedTokens = content.length
+  }
+  return estimatedTokens
+}

package/packages/datadog-plugin-vitest/src/index.js

@@ -0,0 +1,156 @@
+const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
+const { storage } = require('../../datadog-core')
+
+const {
+  TEST_STATUS,
+  finishAllTraceSpans,
+  getTestSuitePath,
+  getTestSuiteCommonTags,
+  TEST_SOURCE_FILE
+} = require('../../dd-trace/src/plugins/util/test')
+const { COMPONENT } = require('../../dd-trace/src/constants')
+
+// Milliseconds that we subtract from the error test duration
+// so that they do not overlap with the following test
+// This is because there's some loss of resolution.
+const MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION = 5
+
+class VitestPlugin extends CiPlugin {
+  static get id () {
+    return 'vitest'
+  }
+
+  constructor (...args) {
+    super(...args)
+
+    this.taskToFinishTime = new WeakMap()
+
+    this.addSub('ci:vitest:test:start', ({ testName, testSuiteAbsolutePath }) => {
+      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+      const store = storage.getStore()
+      const span = this.startTestSpan(
+        testName,
+        testSuite,
+        this.testSuiteSpan,
+        {
+          [TEST_SOURCE_FILE]: testSuite
+        }
+      )
+
+      this.enter(span, store)
+    })
+
+    this.addSub('ci:vitest:test:finish-time', ({ status, task }) => {
+      const store = storage.getStore()
+      const span = store?.span
+
+      // we store the finish time to finish at a later hook
+      // this is because the test might fail at a `afterEach` hook
+      if (span) {
+        span.setTag(TEST_STATUS, status)
+        this.taskToFinishTime.set(task, span._getTime())
+      }
+    })
+
+    this.addSub('ci:vitest:test:pass', ({ task }) => {
+      const store = storage.getStore()
+      const span = store?.span
+
+      if (span) {
+        span.setTag(TEST_STATUS, 'pass')
+        span.finish(this.taskToFinishTime.get(task))
+        finishAllTraceSpans(span)
+      }
+    })
+
+    this.addSub('ci:vitest:test:error', ({ duration, error }) => {
+      const store = storage.getStore()
+      const span = store?.span
+
+      if (span) {
+        span.setTag(TEST_STATUS, 'fail')
+
+        if (error) {
+          span.setTag('error', error)
+        }
+        span.finish(span._startTime + duration - MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION) // milliseconds
+        finishAllTraceSpans(span)
+      }
+    })
+
+    this.addSub('ci:vitest:test:skip', ({ testName, testSuiteAbsolutePath }) => {
+      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+      this.startTestSpan(
+        testName,
+        testSuite,
+        this.testSuiteSpan,
+        {
+          [TEST_SOURCE_FILE]: testSuite,
+          [TEST_STATUS]: 'skip'
+        }
+      ).finish()
+    })
+
+    this.addSub('ci:vitest:test-suite:start', (testSuiteAbsolutePath) => {
+      const testSessionSpanContext = this.tracer.extract('text_map', {
+        'x-datadog-trace-id': process.env.DD_CIVISIBILITY_TEST_SESSION_ID,
+        'x-datadog-parent-id': process.env.DD_CIVISIBILITY_TEST_MODULE_ID
+      })
+
+      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+      const testSuiteMetadata = getTestSuiteCommonTags(
+        this.command,
+        this.frameworkVersion,
+        testSuite,
+        'vitest'
+      )
+      const testSuiteSpan = this.tracer.startSpan('vitest.test_suite', {
+        childOf: testSessionSpanContext,
+        tags: {
+          [COMPONENT]: this.constructor.id,
+          ...this.testEnvironmentMetadata,
+          ...testSuiteMetadata
+        }
+      })
+      const store = storage.getStore()
+      this.enter(testSuiteSpan, store)
+      this.testSuiteSpan = testSuiteSpan
+    })
+
+    this.addSub('ci:vitest:test-suite:finish', ({ status, onFinish }) => {
+      const store = storage.getStore()
+      const span = store?.span
+      if (span) {
+        span.setTag(TEST_STATUS, status)
+        span.finish()
+        finishAllTraceSpans(span)
+      }
+      // TODO: too frequent flush - find for method in worker to decrease frequency
+      this.tracer._exporter.flush(onFinish)
+    })
+
+    this.addSub('ci:vitest:test-suite:error', ({ error }) => {
+      const store = storage.getStore()
+      const span = store?.span
+      if (span && error) {
+        span.setTag('error', error)
+        span.setTag(TEST_STATUS, 'fail')
+      }
+    })
+
+    this.addSub('ci:vitest:session:finish', ({ status, onFinish, error }) => {
+      this.testSessionSpan.setTag(TEST_STATUS, status)
+      this.testModuleSpan.setTag(TEST_STATUS, status)
+      if (error) {
+        this.testModuleSpan.setTag('error', error)
+        this.testSessionSpan.setTag('error', error)
+      }
+      this.testModuleSpan.finish()
+      this.testSessionSpan.finish()
+      finishAllTraceSpans(this.testSessionSpan)
+      this.tracer._exporter.flush(onFinish)
+    })
+  }
+}
+
+module.exports = VitestPlugin

package/packages/dd-trace/src/appsec/iast/path-line.js

@@ -3,6 +3,7 @@
 const path = require('path')
 const process = require('process')
 const { calculateDDBasePath } = require('../../util')
+const { getCallSiteList } = require('../stack_trace')
 const pathLine = {
   getFirstNonDDPathAndLine,
   getNodeModulesPaths,
@@ -24,24 +25,6 @@ const EXCLUDED_PATH_PREFIXES = [
   'async_hooks'
 ]
 
-function getCallSiteInfo () {
-  const previousPrepareStackTrace = Error.prepareStackTrace
-  const previousStackTraceLimit = Error.stackTraceLimit
-  let callsiteList
-  Error.stackTraceLimit = 100
-  try {
-    Error.prepareStackTrace = function (_, callsites) {
-      callsiteList = callsites
-    }
-    const e = new Error()
-    e.stack
-  } finally {
-    Error.prepareStackTrace = previousPrepareStackTrace
-    Error.stackTraceLimit = previousStackTraceLimit
-  }
-  return callsiteList
-}
-
 function getFirstNonDDPathAndLineFromCallsites (callsites, externallyExcludedPaths) {
   if (callsites) {
     for (let i = 0; i < callsites.length; i++) {
@@ -91,7 +74,7 @@ function isExcluded (callsite, externallyExcludedPaths) {
 }
 
 function getFirstNonDDPathAndLine (externallyExcludedPaths) {
-  return getFirstNonDDPathAndLineFromCallsites(getCallSiteInfo(), externallyExcludedPaths)
+  return getFirstNonDDPathAndLineFromCallsites(getCallSiteList(), externallyExcludedPaths)
 }
 
 function getNodeModulesPaths (...paths) {

package/packages/dd-trace/src/appsec/iast/vulnerability-reporter.js

@@ -4,6 +4,7 @@ const { MANUAL_KEEP } = require('../../../../../ext/tags')
 const LRU = require('lru-cache')
 const vulnerabilitiesFormatter = require('./vulnerabilities-formatter')
 const { IAST_ENABLED_TAG_KEY, IAST_JSON_TAG_KEY } = require('./tags')
+const standalone = require('../standalone')
 
 const VULNERABILITIES_KEY = 'vulnerabilities'
 const VULNERABILITY_HASHES_MAX_SIZE = 1000
@@ -57,6 +58,9 @@ function sendVulnerabilities (vulnerabilities, rootSpan) {
         tags[IAST_JSON_TAG_KEY] = JSON.stringify(jsonToSend)
         tags[MANUAL_KEEP] = 'true'
         span.addTags(tags)
+
+        standalone.sample(span)
+
         if (!rootSpan) span.finish()
       }
     }

package/packages/dd-trace/src/appsec/index.js

@@ -40,7 +40,7 @@ function enable (_config) {
     graphql.enable()
 
     if (_config.appsec.rasp.enabled) {
-      rasp.enable()
+      rasp.enable(_config)
     }
 
     setTemplates(_config)

package/packages/dd-trace/src/appsec/rasp.js

@@ -1,11 +1,20 @@
 'use strict'
 
 const { storage } = require('../../../datadog-core')
+const web = require('./../plugins/util/web')
 const addresses = require('./addresses')
 const { httpClientRequestStart } = require('./channels')
+const { reportStackTrace } = require('./stack_trace')
 const waf = require('./waf')
 
-function enable () {
+const RULE_TYPES = {
+  SSRF: 'ssrf'
+}
+
+let config
+
+function enable (_config) {
+  config = _config
   httpClientRequestStart.subscribe(analyzeSsrf)
 }
 
@@ -24,12 +33,30 @@ function analyzeSsrf (ctx) {
     [addresses.HTTP_OUTGOING_URL]: url
   }
   // TODO: Currently this is only monitoring, we should
-  //     block the request if SSRF attempt and
-  //     generate stack traces
-  waf.run({ persistent }, req)
+  //     block the request if SSRF attempt
+  const result = waf.run({ persistent }, req, RULE_TYPES.SSRF)
+  handleResult(result, req)
+}
+
+function getGenerateStackTraceAction (actions) {
+  return actions?.generate_stack
+}
+
+function handleResult (actions, req) {
+  const generateStackTraceAction = getGenerateStackTraceAction(actions)
+  if (generateStackTraceAction && config.appsec.stackTrace.enabled) {
+    const rootSpan = web.root(req)
+    reportStackTrace(
+      rootSpan,
+      generateStackTraceAction.stack_id,
+      config.appsec.stackTrace.maxDepth,
+      config.appsec.stackTrace.maxStackTraces
+    )
+  }
 }
 
 module.exports = {
   enable,
-  disable
+  disable,
+  handleResult
 }