From c691443042fe764b77b8eca411e4bb4175c4a894 Mon Sep 17 00:00:00 2001 From: Ugaitz Urien Date: Fri, 20 Sep 2024 16:58:44 +0200 Subject: [PATCH] Use only one channel instead of 6 --- .../datadog-instrumentations/src/mysql2.js | 33 +++++++++---------- .../test/mysql2.spec.js | 25 +------------- packages/dd-trace/src/appsec/channels.js | 7 +--- .../dd-trace/src/appsec/rasp/sql_injection.js | 25 +++----------- 4 files changed, 21 insertions(+), 69 deletions(-) diff --git a/packages/datadog-instrumentations/src/mysql2.js b/packages/datadog-instrumentations/src/mysql2.js index c94e4d151c6..8cf3210106c 100644 --- a/packages/datadog-instrumentations/src/mysql2.js +++ b/packages/datadog-instrumentations/src/mysql2.js @@ -12,8 +12,7 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Conne const startCh = channel('apm:mysql2:query:start') const finishCh = channel('apm:mysql2:query:finish') const errorCh = channel('apm:mysql2:query:error') - const startConnectionQueryCh = channel('datadog:mysql2:connection:query:start') - const startConnectionExecuteCh = channel('datadog:mysql2:connection:execute:start') + const startOuterQueryCh = channel('datadog:mysql2:outerquery:start') const shouldEmitEndAfterQueryAbort = semver.intersects(version, '>=1.3.3') shimmer.wrap(Connection.prototype, 'addCommand', addCommand => function (cmd) { @@ -33,14 +32,14 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Conne }) shimmer.wrap(Connection.prototype, 'query', query => function (sql, values, cb) { - if (!startConnectionQueryCh.hasSubscribers) return query.apply(this, arguments) + if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments) const sqlIsString = typeof sql === 'string' const sqlString = sqlIsString ? sql : sql?.sql if (sqlString) { const abortController = new AbortController() - startConnectionQueryCh.publish({ sql: sqlString, abortController }) + startOuterQueryCh.publish({ sql: sqlString, abortController }) if (abortController.signal.aborted) { let queryCommand = sql @@ -70,12 +69,12 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Conne }) shimmer.wrap(Connection.prototype, 'execute', execute => function (sql, values, cb) { - if (!startConnectionExecuteCh.hasSubscribers) return execute.apply(this, arguments) + if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments) const sqlString = typeof sql === 'object' ? sql?.sql : sql if (sqlString) { const abortController = new AbortController() - startConnectionExecuteCh.publish({ sql, abortController }) + startOuterQueryCh.publish({ sql, abortController }) if (abortController.signal.aborted) { const addCommand = this.addCommand @@ -150,17 +149,16 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Conne }) addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['>=1'] }, (Pool, version) => { - const startPoolQueryCh = channel('datadog:mysql2:pool:query:start') - const startPoolExecuteCh = channel('datadog:mysql2:pool:execute:start') + 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 (!startPoolQueryCh.hasSubscribers) return query.apply(this, arguments) + if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments) const sqlString = typeof sql === 'object' ? sql?.sql : sql if (sqlString) { const abortController = new AbortController() - startPoolQueryCh.publish({ sql, abortController }) + startOuterQueryCh.publish({ sql, abortController }) if (abortController.signal.aborted) { const getConnection = this.getConnection @@ -193,10 +191,10 @@ addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['>=1'] }, (Pool, versi }) shimmer.wrap(Pool.prototype, 'execute', execute => function (sql, values, cb) { - if (!startPoolExecuteCh.hasSubscribers) return execute.apply(this, arguments) + if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments) const abortController = new AbortController() - startPoolExecuteCh.publish({ sql, abortController }) + startOuterQueryCh.publish({ sql, abortController }) if (abortController.signal.aborted) { if (typeof values === 'function') { @@ -217,18 +215,17 @@ addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['>=1'] }, (Pool, versi // 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 startPoolNamespaceQueryCh = channel('datadog:mysql2:poolnamespace:query:start') - const startPoolNamespaceExecuteCh = channel('datadog:mysql2:poolnamespace:execute:start') + const startOuterQueryCh = channel('datadog:mysql2:outerquery:start') shimmer.wrap(PoolCluster.prototype, 'of', of => function () { const poolNamespace = of.apply(this, arguments) - if (startPoolNamespaceQueryCh.hasSubscribers) { + if (startOuterQueryCh.hasSubscribers) { shimmer.wrap(poolNamespace, 'query', query => function (sql, values, cb) { if (typeof sql === 'object') sql = sql?.sql const abortController = new AbortController() - startPoolNamespaceQueryCh.publish({ sql, abortController }) + startOuterQueryCh.publish({ sql, abortController }) if (abortController.signal.aborted) { const getConnection = this.getConnection @@ -257,12 +254,12 @@ addHook({ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['>=2.3.0'] }, }) } - if (startPoolNamespaceExecuteCh.hasSubscribers) { + if (startOuterQueryCh.hasSubscribers) { shimmer.wrap(poolNamespace, 'execute', execute => function (sql, values, cb) { if (typeof sql === 'object') sql = sql?.sql const abortController = new AbortController() - startPoolNamespaceExecuteCh.publish({ sql, abortController }) + startOuterQueryCh.publish({ sql, abortController }) if (abortController.signal.aborted) { if (typeof values === 'function') { diff --git a/packages/datadog-instrumentations/test/mysql2.spec.js b/packages/datadog-instrumentations/test/mysql2.spec.js index 4bc616e09fe..b12ae0cf725 100644 --- a/packages/datadog-instrumentations/test/mysql2.spec.js +++ b/packages/datadog-instrumentations/test/mysql2.spec.js @@ -28,6 +28,7 @@ describe('mysql2 instrumentation', () => { let apmQueryStartChannel, apmQueryStart, mysql2Version before(() => { + startCh = channel('datadog:mysql2:outerquery:start') return agent.load(['mysql2']) }) @@ -65,10 +66,6 @@ describe('mysql2 instrumentation', () => { }) describe('Connection.prototype.query', () => { - beforeEach(() => { - startCh = channel('datadog:mysql2:connection:query:start') - }) - describe('with string as query', () => { describe('with callback', () => { it('should abort the query on abortController.abort()', (done) => { @@ -231,10 +228,6 @@ describe('mysql2 instrumentation', () => { }) describe('Connection.prototype.execute', () => { - beforeEach(() => { - startCh = channel('datadog:mysql2:connection:execute:start') - }) - describe('with the query in options', () => { it('should abort the query on abortController.abort()', (done) => { startCh.subscribe(abort) @@ -327,10 +320,6 @@ describe('mysql2 instrumentation', () => { }) describe('Pool.prototype.query', () => { - beforeEach(() => { - startCh = channel('datadog:mysql2:pool:query:start') - }) - describe('with callback', () => { it('should abort the query on abortController.abort()', (done) => { startCh.subscribe(abort) @@ -402,10 +391,6 @@ describe('mysql2 instrumentation', () => { }) describe('Pool.prototype.execute', () => { - beforeEach(() => { - startCh = channel('datadog:mysql2:pool:execute:start') - }) - describe('with callback', () => { it('should abort the query on abortController.abort()', (done) => { startCh.subscribe(abort) @@ -469,10 +454,6 @@ describe('mysql2 instrumentation', () => { }) describe('PoolNamespace.prototype.query', () => { - beforeEach(() => { - startCh = channel('datadog:mysql2:poolnamespace:query:start') - }) - it('should abort the query on abortController.abort()', (done) => { startCh.subscribe(abort) const namespace = poolCluster.of() @@ -510,10 +491,6 @@ describe('mysql2 instrumentation', () => { }) describe('PoolNamespace.prototype.execute', () => { - beforeEach(() => { - startCh = channel('datadog:mysql2:poolnamespace:execute:start') - }) - it('should abort the query on abortController.abort()', (done) => { startCh.subscribe(abort) diff --git a/packages/dd-trace/src/appsec/channels.js b/packages/dd-trace/src/appsec/channels.js index b6c36bbf3b2..a451b9ce145 100644 --- a/packages/dd-trace/src/appsec/channels.js +++ b/packages/dd-trace/src/appsec/channels.js @@ -24,11 +24,6 @@ module.exports = { setUncaughtExceptionCaptureCallbackStart: dc.channel('datadog:process:setUncaughtExceptionCaptureCallback:start'), pgQueryStart: dc.channel('apm:pg:query:start'), pgPoolQueryStart: dc.channel('datadog:pg:pool:query:start'), - mysql2ConnectionQueryStart: dc.channel('datadog:mysql2:connection:query:start'), - mysql2ConnectionExecuteStart: dc.channel('datadog:mysql2:connection:execute:start'), - mysql2PoolQueryStart: dc.channel('datadog:mysql2:pool:query:start'), - mysql2PoolExecuteStart: dc.channel('datadog:mysql2:pool:execute:start'), - mysql2PoolNamespaceQueryStart: dc.channel('datadog:mysql2:poolnamespace:query:start'), - mysql2PoolNamespaceExecuteStart: dc.channel('datadog:mysql2:poolnamespace:execute:start'), + mysql2OuterQueryStart: dc.channel('datadog:mysql2:outerquery:start'), wafRunFinished: dc.channel('datadog:waf:run:finish') } diff --git a/packages/dd-trace/src/appsec/rasp/sql_injection.js b/packages/dd-trace/src/appsec/rasp/sql_injection.js index 4eeebcf3b18..d4a165d8615 100644 --- a/packages/dd-trace/src/appsec/rasp/sql_injection.js +++ b/packages/dd-trace/src/appsec/rasp/sql_injection.js @@ -4,12 +4,7 @@ const { pgQueryStart, pgPoolQueryStart, wafRunFinished, - mysql2ConnectionQueryStart, - mysql2ConnectionExecuteStart, - mysql2PoolQueryStart, - mysql2PoolExecuteStart, - mysql2PoolNamespaceQueryStart, - mysql2PoolNamespaceExecuteStart + mysql2OuterQueryStart } = require('../channels') const { storage } = require('../../../../datadog-core') const addresses = require('../addresses') @@ -29,26 +24,14 @@ function enable (_config) { pgPoolQueryStart.subscribe(analyzePgSqlInjection) wafRunFinished.subscribe(clearQuerySet) - mysql2ConnectionQueryStart.subscribe(analyzeMysql2SqlInjection) - mysql2ConnectionExecuteStart.subscribe(analyzeMysql2SqlInjection) - mysql2PoolQueryStart.subscribe(analyzeMysql2SqlInjection) - mysql2PoolExecuteStart.subscribe(analyzeMysql2SqlInjection) - mysql2PoolNamespaceQueryStart.subscribe(analyzeMysql2SqlInjection) - mysql2PoolNamespaceExecuteStart.subscribe(analyzeMysql2SqlInjection) + mysql2OuterQueryStart.subscribe(analyzeMysql2SqlInjection) } function disable () { if (pgQueryStart.hasSubscribers) pgQueryStart.unsubscribe(analyzePgSqlInjection) if (pgPoolQueryStart.hasSubscribers) pgPoolQueryStart.unsubscribe(analyzePgSqlInjection) if (wafRunFinished.hasSubscribers) wafRunFinished.unsubscribe(clearQuerySet) - if (mysql2ConnectionQueryStart.hasSubscribers) mysql2ConnectionQueryStart.unsubscribe(analyzeMysql2SqlInjection) - if (mysql2ConnectionExecuteStart.hasSubscribers) mysql2ConnectionExecuteStart.unsubscribe(analyzeMysql2SqlInjection) - if (mysql2PoolQueryStart.hasSubscribers) mysql2PoolQueryStart.unsubscribe(analyzeMysql2SqlInjection) - if (mysql2PoolExecuteStart.hasSubscribers) mysql2PoolExecuteStart.unsubscribe(analyzeMysql2SqlInjection) - if (mysql2PoolNamespaceQueryStart.hasSubscribers) mysql2PoolNamespaceQueryStart.unsubscribe(analyzeMysql2SqlInjection) - if (mysql2PoolNamespaceExecuteStart.hasSubscribers) { - mysql2PoolNamespaceExecuteStart.unsubscribe(analyzeMysql2SqlInjection) - } + if (mysql2OuterQueryStart.hasSubscribers) mysql2OuterQueryStart.unsubscribe(analyzeMysql2SqlInjection) } function analyzeMysql2SqlInjection (ctx) { @@ -76,7 +59,7 @@ function analyzeSqlInjection (query, dbSystem, abortController) { let executedQueries = reqQueryMap.get(req) if (executedQueries?.has(query)) return - // Do not waste time executing same query twice + // Do not waste time checking same query twice // This also will prevent double calls in pg.Pool internal queries if (!executedQueries) { executedQueries = new Set()