Support tainted strings coming from database for SQLi, SSTi and Code …

…injection (#4904)

docs/test.ts

@@ -131,6 +131,7 @@ tracer.init({
     requestSampling: 50,
     maxConcurrentRequests: 4,
     maxContextOperations: 30,
+    dbRowsToTaint: 12,
     deduplicationEnabled: true,
     redactionEnabled: true,
     redactionNamePattern: 'password',
@@ -147,6 +148,7 @@ tracer.init({
       requestSampling: 50,
       maxConcurrentRequests: 4,
       maxContextOperations: 30,
+      dbRowsToTaint: 6,
       deduplicationEnabled: true,
       redactionEnabled: true,
       redactionNamePattern: 'password',

index.d.ts

@@ -764,7 +764,7 @@ declare namespace tracer {
        */
       maxDepth?: number
     }
-    
+
     /**
      * Configuration enabling LLM Observability. Enablement is superceded by the DD_LLMOBS_ENABLED environment variable.
      */
@@ -2203,6 +2203,12 @@ declare namespace tracer {
      */
     cookieFilterPattern?: string,
 
+    /**
+     * Defines the number of rows to taint in data coming from databases
+     * @default 1
+     */
+    dbRowsToTaint?: number,
+
     /**
      * Whether to enable vulnerability deduplication
      */
@@ -2247,7 +2253,7 @@ declare namespace tracer {
        * Disable LLM Observability tracing.
        */
       disable (): void,
-      
+
       /**
        * Instruments a function by automatically creating a span activated on its
        * scope.
@@ -2289,10 +2295,10 @@ declare namespace tracer {
       /**
        * Decorate a function in a javascript runtime that supports function decorators.
        * Note that this is **not** supported in the Node.js runtime, but is in TypeScript.
-       * 
+       *
        * In TypeScript, this decorator is only supported in contexts where general TypeScript
        * function decorators are supported.
-       * 
+       *
        * @param options Optional LLM Observability span options.
        */
       decorate (options: llmobs.LLMObsNamelessSpanOptions): any
@@ -2309,7 +2315,7 @@ declare namespace tracer {
       /**
        * Sets inputs, outputs, tags, metadata, and metrics as provided for a given LLM Observability span.
        * Note that with the exception of tags, this method will override any existing values for the provided fields.
-       * 
+       *
        * For example:
        * ```javascript
        * llmobs.trace({ kind: 'llm', name: 'myLLM', modelName: 'gpt-4o', modelProvider: 'openai' }, () => {
@@ -2322,7 +2328,7 @@ declare namespace tracer {
        *  })
        * })
        * ```
-       * 
+       *
        * @param span The span to annotate (defaults to the current LLM Observability span if not provided)
        * @param options An object containing the inputs, outputs, tags, metadata, and metrics to set on the span.
        */
@@ -2498,14 +2504,14 @@ declare namespace tracer {
        * LLM Observability span kind. One of `agent`, `workflow`, `task`, `tool`, `retrieval`, `embedding`, or `llm`.
        */
       kind: llmobs.spanKind,
-  
+
       /**
        * The ID of the underlying user session. Required for tracking sessions.
        */
       sessionId?: string,
 
       /**
-       * The name of the ML application that the agent is orchestrating. 
+       * The name of the ML application that the agent is orchestrating.
        * If not provided, the default value will be set to mlApp provided during initalization, or `DD_LLMOBS_ML_APP`.
        */
       mlApp?: string,

packages/datadog-instrumentations/src/pg.js

@@ -62,11 +62,11 @@ function wrapQuery (query) {
         abortController
       })
 
-      const finish = asyncResource.bind(function (error) {
+      const finish = asyncResource.bind(function (error, res) {
         if (error) {
           errorCh.publish(error)
         }
-        finishCh.publish()
+        finishCh.publish({ result: res?.rows })
       })
 
       if (abortController.signal.aborted) {
@@ -119,15 +119,15 @@ function wrapQuery (query) {
       if (newQuery.callback) {
         const originalCallback = callbackResource.bind(newQuery.callback)
         newQuery.callback = function (err, res) {
-          finish(err)
+          finish(err, res)
           return originalCallback.apply(this, arguments)
         }
       } else if (newQuery.once) {
         newQuery
           .once('error', finish)
-          .once('end', () => finish())
+          .once('end', (res) => finish(null, res))
       } else {
-        newQuery.then(() => finish(), finish)
+        newQuery.then((res) => finish(null, res), finish)
       }
 
       try {

packages/datadog-instrumentations/src/sequelize.js

@@ -13,7 +13,7 @@ addHook({ name: 'sequelize', versions: ['>=4'] }, Sequelize => {
   const finishCh = channel('datadog:sequelize:query:finish')
 
   shimmer.wrap(Sequelize.prototype, 'query', query => {
-    return function (sql) {
+    return function (sql, options) {
       if (!startCh.hasSubscribers) {
         return query.apply(this, arguments)
       }
@@ -27,9 +27,14 @@ addHook({ name: 'sequelize', versions: ['>=4'] }, Sequelize => {
         dialect = this.dialect.name
       }
 
-      function onFinish () {
+      function onFinish (result) {
+        const type = options?.type || 'RAW'
+        if (type === 'RAW' && result?.length > 1) {
+          result = result[0]
+        }
+
         asyncResource.bind(function () {
-          finishCh.publish()
+          finishCh.publish({ result })
         }, this).apply(this)
       }
 
@@ -40,7 +45,7 @@ addHook({ name: 'sequelize', versions: ['>=4'] }, Sequelize => {
         })
 
         const promise = query.apply(this, arguments)
-        promise.then(onFinish, onFinish)
+        promise.then(onFinish, () => { onFinish() })
 
         return promise
       }, this).apply(this, arguments)

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

@@ -11,6 +11,10 @@ class CodeInjectionAnalyzer extends InjectionAnalyzer {
   onConfigure () {
     this.addSub('datadog:eval:call', ({ script }) => this.analyze(script))
   }
+
+  _areRangesVulnerable () {
+    return true
+  }
 }
 
 module.exports = new CodeInjectionAnalyzer()

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

@@ -1,19 +1,26 @@
 'use strict'
 const Analyzer = require('./vulnerability-analyzer')
-const { isTainted, getRanges } = require('../taint-tracking/operations')
+const { getRanges } = require('../taint-tracking/operations')
+const { SQL_ROW_VALUE } = require('../taint-tracking/source-types')
 
 class InjectionAnalyzer extends Analyzer {
   _isVulnerable (value, iastContext) {
-    if (value) {
-      return isTainted(iastContext, value)
+    const ranges = value && getRanges(iastContext, value)
+    if (ranges?.length > 0) {
+      return this._areRangesVulnerable(ranges)
     }
+
     return false
   }
 
   _getEvidence (value, iastContext) {
     const ranges = getRanges(iastContext, value)
     return { value, ranges }
   }
+
+  _areRangesVulnerable (ranges) {
+    return ranges?.some(range => range.iinfo.type !== SQL_ROW_VALUE)
+  }
 }
 
 module.exports = InjectionAnalyzer

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

@@ -82,6 +82,10 @@ class SqlInjectionAnalyzer extends InjectionAnalyzer {
       return knexDialect.toUpperCase()
     }
   }
+
+  _areRangesVulnerable () {
+    return true
+  }
 }
 
 module.exports = new SqlInjectionAnalyzer()

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

@@ -13,6 +13,10 @@ class TemplateInjectionAnalyzer extends InjectionAnalyzer {
     this.addSub('datadog:handlebars:register-partial:start', ({ partial }) => this.analyze(partial))
     this.addSub('datadog:pug:compile:start', ({ source }) => this.analyze(source))
   }
+
+  _areRangesVulnerable () {
+    return true
+  }
 }
 
 module.exports = new TemplateInjectionAnalyzer()

packages/dd-trace/src/appsec/iast/iast-plugin.js

@@ -98,7 +98,8 @@ class IastPlugin extends Plugin {
     }
   }
 
-  enable () {
+  enable (iastConfig) {
+    this.iastConfig = iastConfig
     this.configure(true)
   }
 

packages/dd-trace/src/appsec/iast/taint-tracking/index.js

@@ -18,10 +18,10 @@ module.exports = {
   enableTaintTracking (config, telemetryVerbosity) {
     enableRewriter(telemetryVerbosity)
     enableTaintOperations(telemetryVerbosity)
-    taintTrackingPlugin.enable()
+    taintTrackingPlugin.enable(config)
 
-    kafkaContextPlugin.enable()
-    kafkaConsumerPlugin.enable()
+    kafkaContextPlugin.enable(config)
+    kafkaConsumerPlugin.enable(config)
 
     setMaxTransactions(config.maxConcurrentRequests)
   },

packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js

@@ -12,7 +12,8 @@ const {
   HTTP_REQUEST_HEADER_NAME,
   HTTP_REQUEST_PARAMETER,
   HTTP_REQUEST_PATH_PARAM,
-  HTTP_REQUEST_URI
+  HTTP_REQUEST_URI,
+  SQL_ROW_VALUE
 } = require('./source-types')
 const { EXECUTED_SOURCE } = require('../telemetry/iast-metric')
 
@@ -26,6 +27,16 @@ class TaintTrackingPlugin extends SourceIastPlugin {
     this._taintedURLs = new WeakMap()
   }
 
+  configure (config) {
+    super.configure(config)
+
+    let rowsToTaint = this.iastConfig?.dbRowsToTaint
+    if (typeof rowsToTaint !== 'number') {
+      rowsToTaint = 1
+    }
+    this._rowsToTaint = rowsToTaint
+  }
+
   onConfigure () {
     const onRequestBody = ({ req }) => {
       const iastContext = getIastContext(storage.getStore())
@@ -73,6 +84,16 @@ class TaintTrackingPlugin extends SourceIastPlugin {
       ({ cookies }) => this._cookiesTaintTrackingHandler(cookies)
     )
 
+    this.addSub(
+      { channelName: 'datadog:sequelize:query:finish', tag: SQL_ROW_VALUE },
+      ({ result }) => this._taintDatabaseResult(result, 'sequelize')
+    )
+
+    this.addSub(
+      { channelName: 'apm:pg:query:finish', tag: SQL_ROW_VALUE },
+      ({ result }) => this._taintDatabaseResult(result, 'pg')
+    )
+
     this.addSub(
       { channelName: 'datadog:express:process_params:start', tag: HTTP_REQUEST_PATH_PARAM },
       ({ req }) => {
@@ -184,6 +205,32 @@ class TaintTrackingPlugin extends SourceIastPlugin {
     this.taintHeaders(req.headers, iastContext)
     this.taintUrl(req, iastContext)
   }
+
+  _taintDatabaseResult (result, dbOrigin, iastContext = getIastContext(storage.getStore()), name) {
+    if (!iastContext) return result
+
+    if (this._rowsToTaint === 0) return result
+
+    if (Array.isArray(result)) {
+      for (let i = 0; i < result.length && i < this._rowsToTaint; i++) {
+        const nextName = name ? `${name}.${i}` : '' + i
+        result[i] = this._taintDatabaseResult(result[i], dbOrigin, iastContext, nextName)
+      }
+    } else if (result && typeof result === 'object') {
+      if (dbOrigin === 'sequelize' && result.dataValues) {
+        result.dataValues = this._taintDatabaseResult(result.dataValues, dbOrigin, iastContext, name)
+      } else {
+        for (const key in result) {
+          const nextName = name ? `${name}.${key}` : key
+          result[key] = this._taintDatabaseResult(result[key], dbOrigin, iastContext, nextName)
+        }
+      }
+    } else if (typeof result === 'string') {
+      result = newTaintedString(iastContext, result, name, SQL_ROW_VALUE)
+    }
+
+    return result
+  }
 }
 
 module.exports = new TaintTrackingPlugin()

packages/dd-trace/src/appsec/iast/taint-tracking/source-types.js

@@ -11,5 +11,6 @@ module.exports = {
   HTTP_REQUEST_PATH_PARAM: 'http.request.path.parameter',
   HTTP_REQUEST_URI: 'http.request.uri',
   KAFKA_MESSAGE_KEY: 'kafka.message.key',
-  KAFKA_MESSAGE_VALUE: 'kafka.message.value'
+  KAFKA_MESSAGE_VALUE: 'kafka.message.value',
+  SQL_ROW_VALUE: 'sql.row.value'
 }

packages/dd-trace/src/config.js

@@ -485,6 +485,7 @@ class Config {
     this._setValue(defaults, 'headerTags', [])
     this._setValue(defaults, 'hostname', '127.0.0.1')
     this._setValue(defaults, 'iast.cookieFilterPattern', '.{32,}')
+    this._setValue(defaults, 'iast.dbRowsToTaint', 1)
     this._setValue(defaults, 'iast.deduplicationEnabled', true)
     this._setValue(defaults, 'iast.enabled', false)
     this._setValue(defaults, 'iast.maxConcurrentRequests', 2)
@@ -605,6 +606,7 @@ class Config {
       DD_GRPC_SERVER_ERROR_STATUSES,
       JEST_WORKER_ID,
       DD_IAST_COOKIE_FILTER_PATTERN,
+      DD_IAST_DB_ROWS_TO_TAINT,
       DD_IAST_DEDUPLICATION_ENABLED,
       DD_IAST_ENABLED,
       DD_IAST_MAX_CONCURRENT_REQUESTS,
@@ -757,6 +759,7 @@ class Config {
     this._setArray(env, 'headerTags', DD_TRACE_HEADER_TAGS)
     this._setString(env, 'hostname', coalesce(DD_AGENT_HOST, DD_TRACE_AGENT_HOSTNAME))
     this._setString(env, 'iast.cookieFilterPattern', DD_IAST_COOKIE_FILTER_PATTERN)
+    this._setValue(env, 'iast.dbRowsToTaint', maybeInt(DD_IAST_DB_ROWS_TO_TAINT))
     this._setBoolean(env, 'iast.deduplicationEnabled', DD_IAST_DEDUPLICATION_ENABLED)
     this._setBoolean(env, 'iast.enabled', DD_IAST_ENABLED)
     this._setValue(env, 'iast.maxConcurrentRequests', maybeInt(DD_IAST_MAX_CONCURRENT_REQUESTS))
@@ -932,6 +935,7 @@ class Config {
     this._setArray(opts, 'headerTags', options.headerTags)
     this._setString(opts, 'hostname', options.hostname)
     this._setString(opts, 'iast.cookieFilterPattern', options.iast?.cookieFilterPattern)
+    this._setValue(opts, 'iast.dbRowsToTaint', maybeInt(options.iast?.dbRowsToTaint))
     this._setBoolean(opts, 'iast.deduplicationEnabled', options.iast && options.iast.deduplicationEnabled)
     this._setBoolean(opts, 'iast.enabled',
       options.iast && (options.iast === true || options.iast.enabled === true))