Sql injection Exploit Prevention implementation for mysql2 library

packages/datadog-instrumentations/src/mysql2.js

@@ -6,11 +6,14 @@ const {
   AsyncResource
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
+const semver = require('semver')
 
-addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connection => {
+addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Connection, version) => {
   const startCh = channel('apm:mysql2:query:start')
   const finishCh = channel('apm:mysql2:query:finish')
   const errorCh = channel('apm:mysql2:query:error')
+  const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
+  const shouldEmitEndAfterQueryAbort = semver.intersects(version, '>=1.3.3')
 
   shimmer.wrap(Connection.prototype, 'addCommand', addCommand => function (cmd) {
     if (!startCh.hasSubscribers) return addCommand.apply(this, arguments)
@@ -28,6 +31,76 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connec
     return asyncResource.bind(addCommand, this).apply(this, arguments)
   })
 
+  shimmer.wrap(Connection.prototype, 'query', query => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return query.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      const addCommand = this.addCommand
+      this.addCommand = function (cmd) { return cmd }
+
+      let queryCommand
+      try {
+        queryCommand = query.apply(this, arguments)
+      } finally {
+        this.addCommand = addCommand
+      }
+
+      cb = queryCommand.onResult
+
+      process.nextTick(() => {
+        if (cb) {
+          cb(abortController.signal.reason)
+        } else {
+          queryCommand.emit('error', abortController.signal.reason)
+        }
+
+        if (shouldEmitEndAfterQueryAbort) {
+          queryCommand.emit('end')
+        }
+      })
+
+      return queryCommand
+    }
+
+    return query.apply(this, arguments)
+  })
+
+  shimmer.wrap(Connection.prototype, 'execute', execute => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return execute.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      const addCommand = this.addCommand
+      this.addCommand = function (cmd) { return cmd }
+
+      let result
+      try {
+        result = execute.apply(this, arguments)
+      } finally {
+        this.addCommand = addCommand
+      }
+
+      result?.onResult(abortController.signal.reason)
+
+      return result
+    }
+
+    return execute.apply(this, arguments)
+  })
+
   return Connection
 
   function bindExecute (cmd, execute, asyncResource) {
@@ -79,3 +152,149 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connec
     }, cmd))
   }
 })
+
+addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['>=1'] }, (Pool, version) => {
+  const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
+  const shouldEmitEndAfterQueryAbort = semver.intersects(version, '>=1.3.3')
+
+  shimmer.wrap(Pool.prototype, 'query', query => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return query.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      const getConnection = this.getConnection
+      this.getConnection = function () {}
+
+      let queryCommand
+      try {
+        queryCommand = query.apply(this, arguments)
+      } finally {
+        this.getConnection = getConnection
+      }
+
+      process.nextTick(() => {
+        if (queryCommand.onResult) {
+          queryCommand.onResult(abortController.signal.reason)
+        } else {
+          queryCommand.emit('error', abortController.signal.reason)
+        }
+
+        if (shouldEmitEndAfterQueryAbort) {
+          queryCommand.emit('end')
+        }
+      })
+
+      return queryCommand
+    }
+
+    return query.apply(this, arguments)
+  })
+
+  shimmer.wrap(Pool.prototype, 'execute', execute => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return execute.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      if (typeof values === 'function') {
+        cb = values
+      }
+
+      process.nextTick(() => {
+        cb(abortController.signal.reason)
+      })
+      return
+    }
+
+    return execute.apply(this, arguments)
+  })
+
+  return Pool
+})
+
+// PoolNamespace.prototype.query does not exist in mysql2<2.3.0
+addHook({ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['>=2.3.0'] }, PoolCluster => {
+  const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
+  const wrappedPoolNamespaces = new WeakSet()
+
+  shimmer.wrap(PoolCluster.prototype, 'of', of => function () {
+    const poolNamespace = of.apply(this, arguments)
+
+    if (startOuterQueryCh.hasSubscribers && !wrappedPoolNamespaces.has(poolNamespace)) {
+      shimmer.wrap(poolNamespace, 'query', query => function (sql, values, cb) {
+        if (typeof sql === 'object') sql = sql?.sql
+
+        if (!sql) return query.apply(this, arguments)
+
+        const abortController = new AbortController()
+        startOuterQueryCh.publish({ sql, abortController })
+
+        if (abortController.signal.aborted) {
+          const getConnection = this.getConnection
+          this.getConnection = function () {}
+
+          let queryCommand
+          try {
+            queryCommand = query.apply(this, arguments)
+          } finally {
+            this.getConnection = getConnection
+          }
+
+          process.nextTick(() => {
+            if (queryCommand.onResult) {
+              queryCommand.onResult(abortController.signal.reason)
+            } else {
+              queryCommand.emit('error', abortController.signal.reason)
+            }
+
+            queryCommand.emit('end')
+          })
+
+          return queryCommand
+        }
+
+        return query.apply(this, arguments)
+      })
+
+      shimmer.wrap(poolNamespace, 'execute', execute => function (sql, values, cb) {
+        if (typeof sql === 'object') sql = sql?.sql
+
+        if (!sql) return execute.apply(this, arguments)
+
+        const abortController = new AbortController()
+        startOuterQueryCh.publish({ sql, abortController })
+
+        if (abortController.signal.aborted) {
+          if (typeof values === 'function') {
+            cb = values
+          }
+
+          process.nextTick(() => {
+            cb(abortController.signal.reason)
+          })
+
+          return
+        }
+
+        return execute.apply(this, arguments)
+      })
+
+      wrappedPoolNamespaces.add(poolNamespace)
+    }
+
+    return poolNamespace
+  })
+
+  return PoolCluster
+})