From d77f60a40eb2986de34deaab32289c7d9c368471 Mon Sep 17 00:00:00 2001
From: Ugaitz Urien <ugaitz.urien@datadoghq.com>
Date: Tue, 3 Sep 2024 16:53:46 +0200
Subject: [PATCH] Apply some PR suggestions

---
 packages/datadog-instrumentations/src/pg.js        | 11 +++++++----
 packages/dd-trace/src/appsec/rasp/index.js         | 10 +---------
 packages/dd-trace/src/appsec/rasp/sql_injection.js |  3 +--
 packages/dd-trace/src/appsec/rasp/utils.js         | 12 ++++++------
 4 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/packages/datadog-instrumentations/src/pg.js b/packages/datadog-instrumentations/src/pg.js
index f60cdaf4953..55642d82e96 100644
--- a/packages/datadog-instrumentations/src/pg.js
+++ b/packages/datadog-instrumentations/src/pg.js
@@ -71,14 +71,17 @@ function wrapQuery (query) {
 
       if (abortController.signal.aborted) {
         const error = abortController.signal.reason || new Error('Aborted')
+
+        // eslint-disable-next-line max-len
+        // Based on: https://github.com/brianc/node-postgres/blob/54eb0fa216aaccd727765641e7d1cf5da2bc483d/packages/pg/lib/client.js#L510
         const reusingQuery = typeof pgQuery.submit === 'function'
         const callback = arguments[arguments.length - 1]
 
         finish(error)
 
         if (reusingQuery) {
-          if (typeof callback === 'function') {
-            pgQuery.callback = pgQuery.callback || callback
+          if (!pgQuery.callback && typeof callback === 'function') {
+            pgQuery.callback = callback
           }
 
           if (pgQuery.callback) {
@@ -96,9 +99,9 @@ function wrapQuery (query) {
           callback(error)
 
           return
-        } else {
-          return Promise.reject(error)
         }
+
+        return Promise.reject(error)
       }
 
       arguments[0] = pgQuery
diff --git a/packages/dd-trace/src/appsec/rasp/index.js b/packages/dd-trace/src/appsec/rasp/index.js
index fb0954cead7..801608e54d8 100644
--- a/packages/dd-trace/src/appsec/rasp/index.js
+++ b/packages/dd-trace/src/appsec/rasp/index.js
@@ -3,11 +3,10 @@
 const web = require('../../plugins/util/web')
 const { setUncaughtExceptionCaptureCallbackStart } = require('../channels')
 const { block } = require('../blocking')
-const log = require('../../log')
 const ssrf = require('./ssrf')
 const sqli = require('./sql_injection')
 
-const { setAbortOnUncaughtException, DatadogRaspAbortError } = require('./utils')
+const { DatadogRaspAbortError } = require('./utils')
 
 function removeAllListeners (emitter, event) {
   const listeners = emitter.listeners(event)
@@ -88,13 +87,6 @@ function enable (config) {
   sqli.enable(config)
 
   process.on('uncaughtExceptionMonitor', handleUncaughtExceptionMonitor)
-
-  const abortOnUncaughtException = process.execArgv?.includes('--abort-on-uncaught-exception')
-  setAbortOnUncaughtException(abortOnUncaughtException)
-
-  if (abortOnUncaughtException) {
-    log.warn('The --abort-on-uncaught-exception flag is enabled. The RASP module will not block operations.')
-  }
 }
 
 function disable () {
diff --git a/packages/dd-trace/src/appsec/rasp/sql_injection.js b/packages/dd-trace/src/appsec/rasp/sql_injection.js
index bf1b8a48d81..a1efeb0946b 100644
--- a/packages/dd-trace/src/appsec/rasp/sql_injection.js
+++ b/packages/dd-trace/src/appsec/rasp/sql_injection.js
@@ -29,7 +29,7 @@ function analyzePgSqlInjection (ctx) {
   const store = storage.getStore()
   if (!store) return
 
-  const { raspSqlAnalyzed, req } = store
+  const { raspSqlAnalyzed, req, res } = store
 
   if (!req || raspSqlAnalyzed) return
 
@@ -40,7 +40,6 @@ function analyzePgSqlInjection (ctx) {
 
   const result = waf.run({ persistent }, req, RULE_TYPES.SQL_INJECTION)
 
-  const res = store?.res
   handleResult(result, req, res, ctx.abortController, config)
 }
 
diff --git a/packages/dd-trace/src/appsec/rasp/utils.js b/packages/dd-trace/src/appsec/rasp/utils.js
index 3a0e20c18f6..3404ecce4fe 100644
--- a/packages/dd-trace/src/appsec/rasp/utils.js
+++ b/packages/dd-trace/src/appsec/rasp/utils.js
@@ -3,8 +3,13 @@
 const web = require('../../plugins/util/web')
 const { reportStackTrace } = require('../stack_trace')
 const { getBlockingAction } = require('../blocking')
+const log = require('../../log')
 
-let abortOnUncaughtException = false
+const abortOnUncaughtException = process.execArgv?.includes('--abort-on-uncaught-exception')
+
+if (abortOnUncaughtException) {
+  log.warn('The --abort-on-uncaught-exception flag is enabled. The RASP module will not block operations.')
+}
 
 const RULE_TYPES = {
   SSRF: 'ssrf',
@@ -55,13 +60,8 @@ function handleResult (actions, req, res, abortController, config) {
   }
 }
 
-function setAbortOnUncaughtException (newAbortOnUncaughtException) {
-  abortOnUncaughtException = newAbortOnUncaughtException
-}
-
 module.exports = {
   handleResult,
-  setAbortOnUncaughtException,
   RULE_TYPES,
   DatadogRaspAbortError
 }