dd-trace 5.25.0 → 5.27.0

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

@@ -0,0 +1,89 @@
+'use strict'
+
+const { MEASURED } = require('../../../ext/tags')
+const { storage } = require('../../datadog-core')
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+
+const API_KEY = 'langchain.request.api_key'
+const MODEL = 'langchain.request.model'
+const PROVIDER = 'langchain.request.provider'
+const TYPE = 'langchain.request.type'
+
+const LangChainHandler = require('./handlers/default')
+const LangChainChatModelHandler = require('./handlers/language_models/chat_model')
+const LangChainLLMHandler = require('./handlers/language_models/llm')
+const LangChainChainHandler = require('./handlers/chain')
+const LangChainEmbeddingHandler = require('./handlers/embedding')
+
+class LangChainPlugin extends TracingPlugin {
+  static get id () { return 'langchain' }
+  static get operation () { return 'invoke' }
+  static get system () { return 'langchain' }
+  static get prefix () {
+    return 'tracing:apm:langchain:invoke'
+  }
+
+  constructor () {
+    super(...arguments)
+
+    const langchainConfig = this._tracerConfig.langchain || {}
+    this.handlers = {
+      chain: new LangChainChainHandler(langchainConfig),
+      chat_model: new LangChainChatModelHandler(langchainConfig),
+      llm: new LangChainLLMHandler(langchainConfig),
+      embedding: new LangChainEmbeddingHandler(langchainConfig),
+      default: new LangChainHandler(langchainConfig)
+    }
+  }
+
+  bindStart (ctx) {
+    const { resource, type } = ctx
+    const handler = this.handlers[type]
+
+    const instance = ctx.instance
+    const apiKey = handler.extractApiKey(instance)
+    const provider = handler.extractProvider(instance)
+    const model = handler.extractModel(instance)
+
+    const tags = handler.getSpanStartTags(ctx, provider) || []
+
+    if (apiKey) tags[API_KEY] = apiKey
+    if (provider) tags[PROVIDER] = provider
+    if (model) tags[MODEL] = model
+    if (type) tags[TYPE] = type
+
+    const span = this.startSpan('langchain.request', {
+      service: this.config.service,
+      resource,
+      kind: 'client',
+      meta: {
+        [MEASURED]: 1,
+        ...tags
+      }
+    }, false)
+
+    const store = storage.getStore() || {}
+    ctx.currentStore = { ...store, span }
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    const span = ctx.currentStore.span
+
+    const { type } = ctx
+
+    const handler = this.handlers[type]
+    const tags = handler.getSpanEndTags(ctx) || {}
+
+    span.addTags(tags)
+
+    span.finish()
+  }
+
+  getHandler (type) {
+    return this.handlers[type] || this.handlers.default
+  }
+}
+
+module.exports = LangChainPlugin

package/packages/datadog-plugin-langchain/src/tokens.js

@@ -0,0 +1,35 @@
+'use strict'
+
+function getTokensFromLlmOutput (result) {
+  const tokens = {
+    input: 0,
+    output: 0,
+    total: 0
+  }
+  const { llmOutput } = result
+  if (!llmOutput) return tokens
+
+  const tokenUsage = llmOutput.tokenUsage || llmOutput.usage_metadata || llmOutput.usage_metadata
+  if (!tokenUsage) return tokens
+
+  for (const tokenNames of [['input', 'prompt'], ['output', 'completion'], ['total']]) {
+    let token = 0
+    for (const tokenName of tokenNames) {
+      const underScore = `${tokenName}_tokens`
+      const camelCase = `${tokenName}Tokens`
+
+      token = tokenUsage[underScore] || tokenUsage[camelCase] || token
+    }
+
+    tokens[tokenNames[0]] = token
+  }
+
+  // assign total_tokens again in case it was improperly set the first time, or was not on tokenUsage
+  tokens.total = tokens.total || tokens.input + tokens.output
+
+  return tokens
+}
+
+module.exports = {
+  getTokensFromLlmOutput
+}

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

@@ -85,7 +85,7 @@ class MochaPlugin extends CiPlugin {
       }
 
       const relativeCoverageFiles = [...coverageFiles, suiteFile]
-        .map(filename => getTestSuitePath(filename, this.sourceRoot))
+        .map(filename => getTestSuitePath(filename, this.repositoryRoot || this.sourceRoot))
 
       const { _traceId, _spanId } = testSuiteSpan.context()
 

package/packages/datadog-plugin-moleculer/src/server.js

@@ -9,7 +9,6 @@ class MoleculerServerPlugin extends ServerPlugin {
 
   start ({ action, ctx, broker }) {
     const followsFrom = this.tracer.extract('text_map', ctx.meta)
-
     this.startSpan(this.operationName(), {
       childOf: followsFrom || this.activeSpan,
       service: this.config.service || this.serviceName(),

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

@@ -1,61 +1,84 @@
 'use strict'
 
+const TTLCache = require('@isaacs/ttlcache')
+const web = require('../plugins/util/web')
 const log = require('../log')
+const { AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
+
+const MAX_SIZE = 4096
 
 let enabled
-let requestSampling
+let sampledRequests
 
-const sampledRequests = new WeakSet()
+class NoopTTLCache {
+  clear () { }
+  set (key) { return undefined }
+  has (key) { return false }
+}
 
 function configure ({ apiSecurity }) {
   enabled = apiSecurity.enabled
-  setRequestSampling(apiSecurity.requestSampling)
+  sampledRequests = apiSecurity.sampleDelay === 0
+    ? new NoopTTLCache()
+    : new TTLCache({ max: MAX_SIZE, ttl: apiSecurity.sampleDelay * 1000 })
 }
 
 function disable () {
   enabled = false
+  sampledRequests?.clear()
 }
 
-function setRequestSampling (sampling) {
-  requestSampling = parseRequestSampling(sampling)
-}
+function sampleRequest (req, res, force = false) {
+  if (!enabled) return false
 
-function parseRequestSampling (requestSampling) {
-  let parsed = parseFloat(requestSampling)
+  const key = computeKey(req, res)
+  if (!key || isSampled(key)) return false
 
-  if (isNaN(parsed)) {
-    log.warn(`Incorrect API Security request sampling value: ${requestSampling}`)
+  const rootSpan = web.root(req)
+  if (!rootSpan) return false
 
-    parsed = 0
-  } else {
-    parsed = Math.min(1, Math.max(0, parsed))
+  let priority = getSpanPriority(rootSpan)
+  if (!priority) {
+    rootSpan._prioritySampler?.sample(rootSpan)
+    priority = getSpanPriority(rootSpan)
   }
 
-  return parsed
-}
-
-function sampleRequest (req) {
-  if (!enabled || !requestSampling) {
+  if (priority === AUTO_REJECT || priority === USER_REJECT) {
     return false
   }
 
-  const shouldSample = Math.random() <= requestSampling
-
-  if (shouldSample) {
-    sampledRequests.add(req)
+  if (force) {
+    sampledRequests.set(key)
   }
 
-  return shouldSample
+  return true
+}
+
+function isSampled (key) {
+  return sampledRequests.has(key)
+}
+
+function computeKey (req, res) {
+  const route = web.getContext(req)?.paths?.join('') || ''
+  const method = req.method
+  const status = res.statusCode
+
+  if (!method || !status) {
+    log.warn('Unsupported groupkey for API security')
+    return null
+  }
+  return method + route + status
 }
 
-function isSampled (req) {
-  return sampledRequests.has(req)
+function getSpanPriority (span) {
+  const spanContext = span.context?.()
+  return spanContext._sampling?.priority
 }
 
 module.exports = {
   configure,
   disable,
-  setRequestSampling,
   sampleRequest,
-  isSampled
+  isSampled,
+  computeKey
 }

package/packages/dd-trace/src/appsec/iast/analyzers/analyzers.js

@@ -15,6 +15,7 @@ module.exports = {
   PATH_TRAVERSAL_ANALYZER: require('./path-traversal-analyzer'),
   SQL_INJECTION_ANALYZER: require('./sql-injection-analyzer'),
   SSRF: require('./ssrf-analyzer'),
+  TEMPLATE_INJECTION_ANALYZER: require('./template-injection-analyzer'),
   UNVALIDATED_REDIRECT_ANALYZER: require('./unvalidated-redirect-analyzer'),
   WEAK_CIPHER_ANALYZER: require('./weak-cipher-analyzer'),
   WEAK_HASH_ANALYZER: require('./weak-hash-analyzer'),

package/packages/dd-trace/src/appsec/iast/analyzers/header-injection-analyzer.js

@@ -6,7 +6,6 @@ const { getNodeModulesPaths } = require('../path-line')
 const { HEADER_NAME_VALUE_SEPARATOR } = require('../vulnerabilities-formatter/constants')
 const { getRanges } = require('../taint-tracking/operations')
 const {
-  HTTP_REQUEST_COOKIE_NAME,
   HTTP_REQUEST_COOKIE_VALUE,
   HTTP_REQUEST_HEADER_VALUE
 } = require('../taint-tracking/source-types')
@@ -45,13 +44,7 @@ class HeaderInjectionAnalyzer extends InjectionAnalyzer {
     if (this.isExcludedHeaderName(lowerCasedHeaderName) || typeof value !== 'string') return
 
     const ranges = getRanges(iastContext, value)
-    if (ranges?.length > 0) {
-      return !(this.isCookieExclusion(lowerCasedHeaderName, ranges) ||
-        this.isSameHeaderExclusion(lowerCasedHeaderName, ranges) ||
-        this.isAccessControlAllowExclusion(lowerCasedHeaderName, ranges))
-    }
-
-    return false
+    return ranges?.length > 0 && !this.shouldIgnoreHeader(lowerCasedHeaderName, ranges)
   }
 
   _getEvidence (headerInfo, iastContext) {
@@ -75,28 +68,52 @@ class HeaderInjectionAnalyzer extends InjectionAnalyzer {
     return EXCLUDED_HEADER_NAMES.includes(name)
   }
 
-  isCookieExclusion (name, ranges) {
-    if (name === 'set-cookie') {
-      return ranges
-        .every(range => range.iinfo.type === HTTP_REQUEST_COOKIE_VALUE || range.iinfo.type === HTTP_REQUEST_COOKIE_NAME)
-    }
+  isAllRangesFromHeader (ranges, headerName) {
+    return ranges
+      .every(range =>
+        range.iinfo.type === HTTP_REQUEST_HEADER_VALUE && range.iinfo.parameterName?.toLowerCase() === headerName
+      )
+  }
 
-    return false
+  isAllRangesFromSource (ranges, source) {
+    return ranges
+      .every(range => range.iinfo.type === source)
   }
 
+  /**
+   * Exclude access-control-allow-*: when the header starts with access-control-allow- and the
+   * source of the tainted range is a request header
+   */
   isAccessControlAllowExclusion (name, ranges) {
     if (name?.startsWith('access-control-allow-')) {
-      return ranges
-        .every(range => range.iinfo.type === HTTP_REQUEST_HEADER_VALUE)
+      return this.isAllRangesFromSource(ranges, HTTP_REQUEST_HEADER_VALUE)
     }
 
     return false
   }
 
+  /** Exclude when the header is reflected from the request */
   isSameHeaderExclusion (name, ranges) {
     return ranges.length === 1 && name === ranges[0].iinfo.parameterName?.toLowerCase()
   }
 
+  shouldIgnoreHeader (headerName, ranges) {
+    switch (headerName) {
+      case 'set-cookie':
+        /** Exclude set-cookie header if the source of all the tainted ranges are cookies */
+        return this.isAllRangesFromSource(ranges, HTTP_REQUEST_COOKIE_VALUE)
+      case 'pragma':
+        /** Ignore pragma headers when the source is the cache control header. */
+        return this.isAllRangesFromHeader(ranges, 'cache-control')
+      case 'transfer-encoding':
+      case 'content-encoding':
+        /** Ignore transfer and content encoding headers when the source is the accept encoding header. */
+        return this.isAllRangesFromHeader(ranges, 'accept-encoding')
+    }
+
+    return this.isAccessControlAllowExclusion(headerName, ranges) || this.isSameHeaderExclusion(headerName, ranges)
+  }
+
   _getExcludedPaths () {
     return EXCLUDED_PATHS
   }

package/packages/dd-trace/src/appsec/iast/analyzers/template-injection-analyzer.js

@@ -0,0 +1,18 @@
+'use strict'
+
+const InjectionAnalyzer = require('./injection-analyzer')
+const { TEMPLATE_INJECTION } = require('../vulnerabilities')
+
+class TemplateInjectionAnalyzer extends InjectionAnalyzer {
+  constructor () {
+    super(TEMPLATE_INJECTION)
+  }
+
+  onConfigure () {
+    this.addSub('datadog:handlebars:compile:start', ({ source }) => this.analyze(source))
+    this.addSub('datadog:handlebars:register-partial:start', ({ partial }) => this.analyze(partial))
+    this.addSub('datadog:pug:compile:start', ({ source }) => this.analyze(source))
+  }
+}
+
+module.exports = new TemplateInjectionAnalyzer()

package/packages/dd-trace/src/appsec/iast/vulnerabilities-formatter/evidence-redaction/sensitive-handler.js

@@ -5,13 +5,13 @@ const vulnerabilities = require('../../vulnerabilities')
 
 const { contains, intersects, remove } = require('./range-utils')
 
-const codeInjectionSensitiveAnalyzer = require('./sensitive-analyzers/code-injection-sensitive-analyzer')
 const commandSensitiveAnalyzer = require('./sensitive-analyzers/command-sensitive-analyzer')
 const hardcodedPasswordAnalyzer = require('./sensitive-analyzers/hardcoded-password-analyzer')
 const headerSensitiveAnalyzer = require('./sensitive-analyzers/header-sensitive-analyzer')
 const jsonSensitiveAnalyzer = require('./sensitive-analyzers/json-sensitive-analyzer')
 const ldapSensitiveAnalyzer = require('./sensitive-analyzers/ldap-sensitive-analyzer')
 const sqlSensitiveAnalyzer = require('./sensitive-analyzers/sql-sensitive-analyzer')
+const taintedRangeBasedSensitiveAnalyzer = require('./sensitive-analyzers/tainted-range-based-sensitive-analyzer')
 const urlSensitiveAnalyzer = require('./sensitive-analyzers/url-sensitive-analyzer')
 
 const { DEFAULT_IAST_REDACTION_NAME_PATTERN, DEFAULT_IAST_REDACTION_VALUE_PATTERN } = require('./sensitive-regex')
@@ -24,7 +24,8 @@ class SensitiveHandler {
     this._valuePattern = new RegExp(DEFAULT_IAST_REDACTION_VALUE_PATTERN, 'gmi')
 
     this._sensitiveAnalyzers = new Map()
-    this._sensitiveAnalyzers.set(vulnerabilities.CODE_INJECTION, codeInjectionSensitiveAnalyzer)
+    this._sensitiveAnalyzers.set(vulnerabilities.CODE_INJECTION, taintedRangeBasedSensitiveAnalyzer)
+    this._sensitiveAnalyzers.set(vulnerabilities.TEMPLATE_INJECTION, taintedRangeBasedSensitiveAnalyzer)
     this._sensitiveAnalyzers.set(vulnerabilities.COMMAND_INJECTION, commandSensitiveAnalyzer)
     this._sensitiveAnalyzers.set(vulnerabilities.NOSQL_MONGODB_INJECTION, jsonSensitiveAnalyzer)
     this._sensitiveAnalyzers.set(vulnerabilities.LDAP_INJECTION, ldapSensitiveAnalyzer)

package/packages/dd-trace/src/appsec/iast/vulnerabilities.js

@@ -13,6 +13,7 @@ module.exports = {
   PATH_TRAVERSAL: 'PATH_TRAVERSAL',
   SQL_INJECTION: 'SQL_INJECTION',
   SSRF: 'SSRF',
+  TEMPLATE_INJECTION: 'TEMPLATE_INJECTION',
   UNVALIDATED_REDIRECT: 'UNVALIDATED_REDIRECT',
   WEAK_CIPHER: 'WEAK_CIPHER',
   WEAK_HASH: 'WEAK_HASH',

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

@@ -145,10 +145,6 @@ function incomingHttpStartTranslator ({ req, res, abortController }) {
     persistent[addresses.HTTP_CLIENT_IP] = clientIp
   }
 
-  if (apiSecuritySampler.sampleRequest(req)) {
-    persistent[addresses.WAF_CONTEXT_PROCESSOR] = { 'extract-schema': true }
-  }
-
   const actions = waf.run({ persistent }, req)
 
   handleResults(actions, req, res, rootSpan, abortController)
@@ -172,6 +168,10 @@ function incomingHttpEndTranslator ({ req, res }) {
     persistent[addresses.HTTP_INCOMING_QUERY] = req.query
   }
 
+  if (apiSecuritySampler.sampleRequest(req, res, true)) {
+    persistent[addresses.WAF_CONTEXT_PROCESSOR] = { 'extract-schema': true }
+  }
+
   if (Object.keys(persistent).length) {
     waf.run({ persistent }, req)
   }
@@ -228,9 +228,9 @@ function onRequestProcessParams ({ req, res, abortController, params }) {
   handleResults(results, req, res, rootSpan, abortController)
 }
 
-function onResponseBody ({ req, body }) {
+function onResponseBody ({ req, res, body }) {
   if (!body || typeof body !== 'object') return
-  if (!apiSecuritySampler.isSampled(req)) return
+  if (!apiSecuritySampler.sampleRequest(req, res)) return
 
   // we don't support blocking at this point, so no results needed
   waf.run({