dd-trace 4.40.0 → 4.42.0

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-undici/src/index.js

@@ -0,0 +1,12 @@
+'use strict'
+
+const FetchPlugin = require('../../datadog-plugin-fetch/src/index.js')
+
+class UndiciPlugin extends FetchPlugin {
+  static get id () { return 'undici' }
+  static get prefix () {
+    return 'tracing:apm:undici:fetch'
+  }
+}
+
+module.exports = UndiciPlugin

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/blocking.js

@@ -111,6 +111,10 @@ function block (req, res, rootSpan, abortController, actionParameters) {
 
   const { body, headers, statusCode } = getBlockingData(req, null, rootSpan, actionParameters)
 
+  for (const headerName of res.getHeaderNames()) {
+    res.removeHeader(headerName)
+  }
+
   res.writeHead(statusCode, headers).end(body)
 
   abortController?.abort()

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

@@ -18,5 +18,6 @@ module.exports = {
   nextBodyParsed: dc.channel('apm:next:body-parsed'),
   nextQueryParsed: dc.channel('apm:next:query-parsed'),
   responseBody: dc.channel('datadog:express:response:json:start'),
+  responseWriteHead: dc.channel('apm:http:server:response:writeHead:start'),
   httpClientRequestStart: dc.channel('apm:http:client:request:start')
 }

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

@@ -12,7 +12,8 @@ const {
   queryParser,
   nextBodyParsed,
   nextQueryParsed,
-  responseBody
+  responseBody,
+  responseWriteHead
 } = require('./channels')
 const waf = require('./waf')
 const addresses = require('./addresses')
@@ -39,7 +40,7 @@ function enable (_config) {
     graphql.enable()
 
     if (_config.appsec.rasp.enabled) {
-      rasp.enable()
+      rasp.enable(_config)
     }
 
     setTemplates(_config)
@@ -60,6 +61,7 @@ function enable (_config) {
     queryParser.subscribe(onRequestQueryParsed)
     cookieParser.subscribe(onRequestCookieParser)
     responseBody.subscribe(onResponseBody)
+    responseWriteHead.subscribe(onResponseWriteHead)
 
     if (_config.appsec.eventTracking.enabled) {
       passportVerify.subscribe(onPassportVerify)
@@ -110,14 +112,7 @@ function incomingHttpStartTranslator ({ req, res, abortController }) {
 }
 
 function incomingHttpEndTranslator ({ req, res }) {
-  // TODO: this doesn't support headers sent with res.writeHead()
-  const responseHeaders = Object.assign({}, res.getHeaders())
-  delete responseHeaders['set-cookie']
-
-  const persistent = {
-    [addresses.HTTP_INCOMING_RESPONSE_CODE]: '' + res.statusCode,
-    [addresses.HTTP_INCOMING_RESPONSE_HEADERS]: responseHeaders
-  }
+  const persistent = {}
 
   // we need to keep this to support other body parsers
   // TODO: no need to analyze it if it was already done by the body-parser hook
@@ -139,7 +134,9 @@ function incomingHttpEndTranslator ({ req, res }) {
     persistent[addresses.HTTP_INCOMING_QUERY] = req.query
   }
 
-  waf.run({ persistent }, req)
+  if (Object.keys(persistent).length) {
+    waf.run({ persistent }, req)
+  }
 
   waf.disposeContext(req)
 
@@ -225,12 +222,48 @@ function onPassportVerify ({ credentials, user }) {
   passportTrackEvent(credentials, user, rootSpan, config.appsec.eventTracking.mode)
 }
 
+const responseAnalyzedSet = new WeakSet()
+const responseBlockedSet = new WeakSet()
+
+function onResponseWriteHead ({ req, res, abortController, statusCode, responseHeaders }) {
+  // avoid "write after end" error
+  if (responseBlockedSet.has(res)) {
+    abortController?.abort()
+    return
+  }
+
+  // avoid double waf call
+  if (responseAnalyzedSet.has(res)) {
+    return
+  }
+
+  const rootSpan = web.root(req)
+  if (!rootSpan) return
+
+  responseHeaders = Object.assign({}, responseHeaders)
+  delete responseHeaders['set-cookie']
+
+  const results = waf.run({
+    persistent: {
+      [addresses.HTTP_INCOMING_RESPONSE_CODE]: '' + statusCode,
+      [addresses.HTTP_INCOMING_RESPONSE_HEADERS]: responseHeaders
+    }
+  }, req)
+
+  responseAnalyzedSet.add(res)
+
+  handleResults(results, req, res, rootSpan, abortController)
+}
+
 function handleResults (actions, req, res, rootSpan, abortController) {
   if (!actions || !req || !res || !rootSpan || !abortController) return
 
   const blockingAction = getBlockingAction(actions)
   if (blockingAction) {
     block(req, res, rootSpan, abortController, blockingAction)
+    if (!abortController.signal || abortController.signal.aborted) {
+      responseBlockedSet.add(res)
+    }
   }
 }
 
@@ -256,6 +289,7 @@ function disable () {
   if (cookieParser.hasSubscribers) cookieParser.unsubscribe(onRequestCookieParser)
   if (responseBody.hasSubscribers) responseBody.unsubscribe(onResponseBody)
   if (passportVerify.hasSubscribers) passportVerify.unsubscribe(onPassportVerify)
+  if (responseWriteHead.hasSubscribers) responseWriteHead.unsubscribe(onResponseWriteHead)
 }
 
 module.exports = {

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
 }