Skip to content

Commit

Permalink
Sql injection Exploit Prevention implementation for mysql2 library (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
uurien authored and tlhunter committed Oct 7, 2024
1 parent 86bc5d4 commit 8dd0bc2
Show file tree
Hide file tree
Showing 8 changed files with 1,269 additions and 7 deletions.
221 changes: 220 additions & 1 deletion packages/datadog-instrumentations/src/mysql2.js
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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) {
Expand Down Expand Up @@ -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
})
Loading

0 comments on commit 8dd0bc2

Please sign in to comment.