Support tainted strings coming from database for SQLi, SSTi and Code injection (#4904)

Author: uurien

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)
   },