Change approach to prevent double calls

DataDog · Sep 5, 2024 · a3ea4fc · a3ea4fc
1 parent d77f60a
commit a3ea4fc
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 25 deletions.
diff --git a/packages/dd-trace/src/appsec/channels.js b/packages/dd-trace/src/appsec/channels.js
@@ -23,7 +23,6 @@ module.exports = {
   responseSetHeader: dc.channel('datadog:http:server:response:set-header:start'),
   setUncaughtExceptionCaptureCallbackStart: dc.channel('datadog:process:setUncaughtExceptionCaptureCallback:start'),
   pgQueryStart: dc.channel('apm:pg:query:start'),
-  pgPoolQueryStart: dc.channel('datadog:pg:pool:query:start'),
-  pgPoolQueryFinish: dc.channel('datadog:pg:pool:query:finish')
+  pgPoolQueryStart: dc.channel('datadog:pg:pool:query:start')
 
 }
diff --git a/packages/dd-trace/src/appsec/rasp/sql_injection.js b/packages/dd-trace/src/appsec/rasp/sql_injection.js
@@ -1,25 +1,23 @@
 'use strict'
 
-const { pgQueryStart, pgPoolQueryStart, pgPoolQueryFinish } = require('../channels')
+const { pgQueryStart, pgPoolQueryStart } = require('../channels')
 const { storage } = require('../../../../datadog-core')
 const addresses = require('../addresses')
 const waf = require('../waf')
 const { RULE_TYPES, handleResult } = require('./utils')
 
 let config
-
+const reqQueryMap = new WeakMap() // WeakMap<Request, Set<querystring>>
 function enable (_config) {
   config = _config
 
   pgQueryStart.subscribe(analyzePgSqlInjection)
-  pgPoolQueryStart.subscribe(analyzePgSqlInjectionInPool)
-  pgPoolQueryFinish.subscribe(pgPoolFinish)
+  pgPoolQueryStart.subscribe(analyzePgSqlInjection)
 }
 
 function disable () {
   if (pgQueryStart.hasSubscribers) pgQueryStart.unsubscribe(analyzePgSqlInjection)
-  if (pgPoolQueryStart.hasSubscribers) pgPoolQueryStart.subscribe(analyzePgSqlInjectionInPool)
-  if (pgPoolQueryFinish.hasSubscribers) pgPoolQueryFinish.subscribe(pgPoolFinish)
+  if (pgPoolQueryStart.hasSubscribers) pgPoolQueryStart.subscribe(analyzePgSqlInjection)
 }
 
 function analyzePgSqlInjection (ctx) {
@@ -33,6 +31,17 @@ function analyzePgSqlInjection (ctx) {
 
   if (!req || raspSqlAnalyzed) return
 
+  let executedQueries = reqQueryMap.get(req)
+  if (executedQueries?.has(query)) return
+
+  // Do not waste time executing same query twice
+  // This also will prevent double calls in pg.Pool internal queries
+  if (!executedQueries) {
+    executedQueries = new Set()
+    reqQueryMap.set(req, executedQueries)
+  }
+  executedQueries.add(query)
+
   const persistent = {
     [addresses.DB_STATEMENT]: query,
     [addresses.DB_SYSTEM]: 'postgresql' // TODO: Extract to constant
@@ -43,21 +52,4 @@ function analyzePgSqlInjection (ctx) {
   handleResult(result, req, res, ctx.abortController, config)
 }
 
-function analyzePgSqlInjectionInPool (ctx) {
-  const parentStore = storage.getStore()
-  if (!parentStore) return
-
-  analyzePgSqlInjection(ctx)
-
-  storage.enterWith({ ...parentStore, raspSqlAnalyzed: true, raspSqlParentStore: parentStore })
-}
-
-function pgPoolFinish () {
-  const store = storage.getStore()
-  if (!store) return
-  if (!store.raspSqlParentStore) return
-
-  storage.enterWith(store.raspSqlParentStore)
-}
-
 module.exports = { enable, disable }
diff --git a/packages/dd-trace/test/appsec/rasp/sql_injection.pg.plugin.spec.js b/packages/dd-trace/test/appsec/rasp/sql_injection.pg.plugin.spec.js
@@ -188,6 +188,61 @@ describe('RASP - sql_injection', () => {
 
             assert.fail('Request should be blocked')
           })
+
+          describe('double calls', () => {
+            const WAFContextWrapper = require('../../../src/appsec/waf/waf_context_wrapper')
+            let run
+
+            beforeEach(() => {
+              run = sinon.spy(WAFContextWrapper.prototype, 'run')
+            })
+
+            afterEach(() => {
+              sinon.restore()
+            })
+
+            it('should call to waf only once for sql injection using pg Pool', async () => {
+              app = async (req, res) => {
+                try {
+                  await pool.query(`SELECT * FROM users WHERE id = '${req.query.param}'`)
+                } catch (err) {
+                  if (err?.name === 'DatadogRaspAbortError') {
+                    res.statusCode = 500
+                  }
+                }
+                res.end()
+              }
+
+              await axios.get('/?param=123')
+
+              assert.equal(run.args.filter(arg => arg[1] === 'sql_injection').length, 1)
+            })
+
+            it('should call to waf twice for sql injection with two different queries in pg Pool', async () => {
+              app = async (req, res) => {
+                try {
+                  await pool.query(`SELECT * FROM users WHERE id = '${req.query.param}'`)
+                } catch (err) {
+                  if (err?.name === 'DatadogRaspAbortError') {
+                    res.statusCode = 500
+                  }
+                }
+
+                try {
+                  await pool.query(`SELECT * FROM users2 WHERE id = '${req.query.param}'`)
+                } catch (err) {
+                  if (err?.name === 'DatadogRaspAbortError') {
+                    res.statusCode = 500
+                  }
+                }
+                res.end()
+              }
+
+              await axios.get('/?param=123')
+
+              assert.equal(run.args.filter(arg => arg[1] === 'sql_injection').length, 2)
+            })
+          })
         })
       })
     })