pg DB sources working for SQLi

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, result) {
         if (error) {
           errorCh.publish(error)
         }
-        finishCh.publish()
+        finishCh.publish({ result: result?.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', (result) => finish(null, result))
       } else {
-        newQuery.then(() => finish(), finish)
+        newQuery.then((res) => finish(null, res), finish)
       }
 
       try {

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

@@ -71,7 +71,12 @@ class TaintTrackingPlugin extends SourceIastPlugin {
 
     this.addSub(
       { channelName: 'datadog:sequelize:query:finish', tag: SQL_ROW_VALUE },
-      ({ result }) => this._taintSequelizeResult(result)
+      ({ result }) => this._taintDatabaseResult(result)
+    )
+
+    this.addSub(
+      { channelName: 'apm:pg:query:finish', tag: SQL_ROW_VALUE },
+      ({ result }) => this._taintDatabaseResult(result)
     )
 
     this.addSub(
@@ -177,25 +182,24 @@ class TaintTrackingPlugin extends SourceIastPlugin {
     this.taintUrl(req, iastContext)
   }
 
-  _taintSequelizeResult (result, iastContext = getIastContext(storage.getStore())) {
+  _taintDatabaseResult (result, iastContext = getIastContext(storage.getStore())) {
     if (!iastContext) return result
 
     const rowsToTaint = 1 // TODO fill this from config
     if (Array.isArray(result)) {
       for (let i = 0; i < result.length && i < rowsToTaint; i++) {
-        result[i] = this._taintSequelizeResult(result[i], iastContext)
+        result[i] = this._taintDatabaseResult(result[i], iastContext)
       }
     } else if (result && typeof result === 'object') {
-      if (result.dataValues) {
-        result.dataValues = this._taintSequelizeResult(result.dataValues, iastContext)
+      if (result.dataValues) { // TODO keep this in mind only with sequelize
+        result.dataValues = this._taintDatabaseResult(result.dataValues, iastContext)
       } else {
         Object.keys(result).forEach(key => {
-          result[key] = this._taintSequelizeResult(result[key], iastContext)
+          result[key] = this._taintDatabaseResult(result[key], iastContext)
         })
       }
     } else if (typeof result === 'string') {
       result = newTaintedString(iastContext, result, SQL_ROW_VALUE, SQL_ROW_VALUE)
-      // console.log('tainting?', result, getRanges(iastContext, result))
     }
 
     return result

packages/dd-trace/test/appsec/iast/taint-tracking/sources/sql_row.pg.plugin.spec.js

@@ -1 +1,113 @@
 'use strict'
+
+const { prepareTestServerForIast } = require('../../utils')
+
+const connectionData = {
+  host: '127.0.0.1',
+  user: 'postgres',
+  password: 'postgres',
+  database: 'postgres',
+  application_name: 'test'
+}
+
+describe('db sources with pg', () => {
+  let pg
+  withVersions('pg', 'pg', '>=8.0.3', version => {
+    let client
+    beforeEach(async () => {
+      pg = require(`../../../../../../../versions/pg@${version}`).get()
+      client = new pg.Client(connectionData)
+      await client.connect()
+
+      await client.query(`CREATE TABLE IF NOT EXISTS examples (
+                                      id INT,
+                                      name VARCHAR(50),
+                                      query VARCHAR(100),
+                                      command VARCHAR(50))`)
+
+      await client.query(`INSERT INTO examples (id, name, query, command)
+                                     VALUES (1, 'Item1', 'SELECT 1', 'ls'),
+                                            (2, 'Item2', 'SELECT 1', 'ls'),
+                                            (3, 'Item3', 'SELECT 1', 'ls')`)
+    })
+
+    afterEach(async () => {
+      await client.query('DROP TABLE examples')
+      client.end()
+    })
+
+    prepareTestServerForIast('sequelize', (testThatRequestHasVulnerability, testThatRequestHasNoVulnerability) => {
+      describe('using pg.Client', () => {
+        testThatRequestHasVulnerability(async (req, res) => {
+          const result = await client.query('SELECT * FROM examples')
+
+          const firstItem = result.rows[0]
+
+          await client.query(firstItem.query)
+
+          res.end()
+        }, 'SQL_INJECTION', { occurrences: 1 }, null, null,
+        'Should have SQL_INJECTION using the first row of the result')
+
+        testThatRequestHasNoVulnerability(async (req, res) => {
+          const result = await client.query('SELECT * FROM examples')
+
+          const secondItem = result.rows[1]
+
+          await client.query(secondItem.query)
+
+          res.end()
+        }, 'SQL_INJECTION', null, 'Should not taint the second row of a query with default configuration')
+
+        testThatRequestHasNoVulnerability(async (req, res) => {
+          const result = await client.query('SELECT * from examples')
+          const firstItem = result.rows[0]
+
+          const childProcess = require('child_process')
+          childProcess.execSync(firstItem.command)
+
+          res.end('OK')
+        }, 'COMMAND_INJECTION', null, 'Should not detect COMMAND_INJECTION with database source')
+      })
+
+      describe('using pg.Pool', () => {
+        let pool
+
+        beforeEach(() => {
+          pool = new pg.Pool(connectionData)
+        })
+
+        testThatRequestHasVulnerability(async (req, res) => {
+          const result = await pool.query('SELECT * FROM examples')
+
+          const firstItem = result.rows[0]
+
+          await client.query(firstItem.query)
+
+          res.end()
+        }, 'SQL_INJECTION', { occurrences: 1 }, null, null,
+        'Should have SQL_INJECTION using the first row of the result')
+
+        testThatRequestHasNoVulnerability(async (req, res) => {
+          const result = await pool.query('SELECT * FROM examples')
+
+          const secondItem = result.rows[1]
+
+          await client.query(secondItem.query)
+
+          res.end()
+        }, 'SQL_INJECTION', null, 'Should not taint the second row of a query with default configuration')
+
+        testThatRequestHasNoVulnerability(async (req, res) => {
+          const result = await pool.query('SELECT * from examples')
+          const firstItem = result.rows[0]
+
+          const childProcess = require('child_process')
+          childProcess.execSync(firstItem.command)
+
+          res.end('OK')
+        }, 'COMMAND_INJECTION', null, 'Should not detect COMMAND_INJECTION with database source')
+      })
+    })
+  })
+})

packages/dd-trace/test/appsec/iast/taint-tracking/sources/sql_row.sequelize.plugin.spec.js

@@ -14,12 +14,12 @@ describe('db sources with sequelize', () => {
           logging: false
         })
         await sequelize.query(`CREATE TABLE examples (
-  id INT,
-  name VARCHAR(50),
-  query VARCHAR(100),
-  command VARCHAR(50),
-  createdAt DATETIME DEFAULT CURRENT_TIMESTAMP,
-  updatedAt DATETIME DEFAULT CURRENT_TIMESTAMP )`)
+                                id INT,
+                                name VARCHAR(50),
+                                query VARCHAR(100),
+                                command VARCHAR(50),
+                                createdAt DATETIME DEFAULT CURRENT_TIMESTAMP,
+                                updatedAt DATETIME DEFAULT CURRENT_TIMESTAMP )`)
 
         await sequelize.query(`INSERT INTO examples (id, name, query, command)
                                VALUES (1, 'Item1', 'SELECT 1', 'ls'),
@@ -95,6 +95,7 @@ describe('db sources with sequelize', () => {
         testThatRequestHasNoVulnerability(async (req, res) => {
           const examples = await Example.findAll()
 
+          const childProcess = require('child_process')
           childProcess.execSync(examples[0].command)
 
           res.end('OK')