From c44d5392f3f1ee5aa8465f9d78c0de66125626ea Mon Sep 17 00:00:00 2001 From: ishabi Date: Wed, 27 Nov 2024 15:13:26 +0100 Subject: [PATCH] clean router handle instrumentation --- packages/datadog-instrumentations/src/express.js | 16 ++++++++-------- packages/datadog-instrumentations/src/router.js | 9 +++------ .../src/appsec/iast/taint-tracking/plugin.js | 2 +- .../appsec/iast/taint-tracking/plugin.spec.js | 2 +- .../taint-tracking.express.plugin.spec.js | 6 ++---- packages/dd-trace/test/plugins/externals.json | 2 +- 6 files changed, 16 insertions(+), 21 deletions(-) diff --git a/packages/datadog-instrumentations/src/express.js b/packages/datadog-instrumentations/src/express.js index a4de43e3f2..a74d3c2e54 100644 --- a/packages/datadog-instrumentations/src/express.js +++ b/packages/datadog-instrumentations/src/express.js @@ -58,6 +58,8 @@ function wrapResponseRender (render) { } addHook({ name: 'express', versions: ['>=4'] }, express => { + shimmer.wrap(express.application, 'handle', wrapHandle) + shimmer.wrap(express.response, 'json', wrapResponseJson) shimmer.wrap(express.response, 'jsonp', wrapResponseJson) shimmer.wrap(express.response, 'render', wrapResponseRender) @@ -66,7 +68,6 @@ addHook({ name: 'express', versions: ['>=4'] }, express => { }) addHook({ name: 'express', versions: ['>=4.0.0 <5.0.0'] }, express => { - shimmer.wrap(express.application, 'handle', wrapHandle) shimmer.wrap(express.Router, 'use', wrapRouterMethod) shimmer.wrap(express.Router, 'route', wrapRouterMethod) @@ -74,22 +75,21 @@ addHook({ name: 'express', versions: ['>=4.0.0 <5.0.0'] }, express => { }) addHook({ name: 'express', versions: ['>=5.0.0'] }, express => { - shimmer.wrap(express.application, 'handle', wrapHandle) shimmer.wrap(express.Router.prototype, 'use', wrapRouterMethod) shimmer.wrap(express.Router.prototype, 'route', wrapRouterMethod) return express }) -const queryReaderReadCh = channel('datadog:query:read:finish') +const queryParserReadCh = channel('datadog:query:read:finish') function publishQueryParsedAndNext (req, res, next) { return shimmer.wrapFunction(next, next => function () { - if (queryReaderReadCh.hasSubscribers && req) { + if (queryParserReadCh.hasSubscribers && req) { const abortController = new AbortController() const query = req.query - queryReaderReadCh.publish({ req, res, query, abortController }) + queryParserReadCh.publish({ req, res, query, abortController }) if (abortController.signal.aborted) return } @@ -146,7 +146,7 @@ addHook({ name: 'express', versions: ['>=4.3.0 <5.0.0'] }, express => { return express }) -const queryParserReadCh = channel('datadog:query:parse:finish') +const queryReadCh = channel('datadog:express:query:finish') addHook({ name: 'express', file: ['lib/request.js'], versions: ['>=5.0.0'] }, request => { const requestDescriptor = Object.getOwnPropertyDescriptor(request, 'query') @@ -155,8 +155,8 @@ addHook({ name: 'express', file: ['lib/request.js'], versions: ['>=5.0.0'] }, re return function wrappedGet () { const query = originalGet.apply(this, arguments) - if (queryParserReadCh.hasSubscribers && query) { - queryParserReadCh.publish({ query }) + if (queryReadCh.hasSubscribers && query) { + queryReadCh.publish({ query }) } return query diff --git a/packages/datadog-instrumentations/src/router.js b/packages/datadog-instrumentations/src/router.js index 5491034a7c..b6f3a583c7 100644 --- a/packages/datadog-instrumentations/src/router.js +++ b/packages/datadog-instrumentations/src/router.js @@ -183,11 +183,8 @@ addHook({ name: 'router', versions: ['>=2'] }, Router => { return function wrappedMethod (options) { const router = originalRouter.call(this, options) - if (!router._queryParsingWrapped) { - router._queryParsingWrapped = true - const originalHandle = router.handle - - router.handle = function (req, res, next) { + shimmer.wrap(router, 'handle', function wrapHandle (originalHandle) { + return function wrappedHandle (req, res, next) { const abortController = new AbortController() if (queryParserReadCh.hasSubscribers && req) { @@ -198,7 +195,7 @@ addHook({ name: 'router', versions: ['>=2'] }, Router => { return originalHandle.apply(this, arguments) } - } + }) return router } diff --git a/packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js b/packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js index b4b8a01f59..62fdd46d02 100644 --- a/packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js +++ b/packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js @@ -51,7 +51,7 @@ class TaintTrackingPlugin extends SourceIastPlugin { ) this.addSub( - { channelName: 'datadog:query:parse:finish', tag: HTTP_REQUEST_PARAMETER }, + { channelName: 'datadog:express:query:finish', tag: HTTP_REQUEST_PARAMETER }, ({ query }) => this._taintTrackingHandler(HTTP_REQUEST_PARAMETER, query) ) diff --git a/packages/dd-trace/test/appsec/iast/taint-tracking/plugin.spec.js b/packages/dd-trace/test/appsec/iast/taint-tracking/plugin.spec.js index d6ee9ed567..5f9c4f4860 100644 --- a/packages/dd-trace/test/appsec/iast/taint-tracking/plugin.spec.js +++ b/packages/dd-trace/test/appsec/iast/taint-tracking/plugin.spec.js @@ -47,7 +47,7 @@ describe('IAST Taint tracking plugin', () => { expect(taintTrackingPlugin._subscriptions[0]._channel.name).to.equals('datadog:body-parser:read:finish') expect(taintTrackingPlugin._subscriptions[1]._channel.name).to.equals('datadog:multer:read:finish') expect(taintTrackingPlugin._subscriptions[2]._channel.name).to.equals('datadog:query:read:finish') - expect(taintTrackingPlugin._subscriptions[3]._channel.name).to.equals('datadog:query:parse:finish') + expect(taintTrackingPlugin._subscriptions[3]._channel.name).to.equals('datadog:express:query:finish') expect(taintTrackingPlugin._subscriptions[4]._channel.name).to.equals('apm:express:middleware:next') expect(taintTrackingPlugin._subscriptions[5]._channel.name).to.equals('datadog:cookie:parse:finish') expect(taintTrackingPlugin._subscriptions[6]._channel.name).to.equals('datadog:express:process_params:start') diff --git a/packages/dd-trace/test/appsec/iast/taint-tracking/sources/taint-tracking.express.plugin.spec.js b/packages/dd-trace/test/appsec/iast/taint-tracking/sources/taint-tracking.express.plugin.spec.js index e8bb7e497c..ebc35bfaba 100644 --- a/packages/dd-trace/test/appsec/iast/taint-tracking/sources/taint-tracking.express.plugin.spec.js +++ b/packages/dd-trace/test/appsec/iast/taint-tracking/sources/taint-tracking.express.plugin.spec.js @@ -189,8 +189,7 @@ describe('Path params sourcing with express', () => { res.status(200).send() }) - app.param('parameter1', checkParamIsTaintedAndNext) - app.param('parameter2', checkParamIsTaintedAndNext) + app.param(['parameter1', 'parameter2'], checkParamIsTaintedAndNext) appListener = app.listen(0, 'localhost', () => { const port = appListener.address().port @@ -209,8 +208,7 @@ describe('Path params sourcing with express', () => { res.status(200).send() }) - app.param('parameter1', checkParamIsTaintedAndNext) - app.param('parameter2', checkParamIsTaintedAndNext) + app.param(['parameter1', 'parameter2'], checkParamIsTaintedAndNext) appListener = app.listen(0, 'localhost', () => { const port = appListener.address().port diff --git a/packages/dd-trace/test/plugins/externals.json b/packages/dd-trace/test/plugins/externals.json index d045e07e1d..7e1c0f9163 100644 --- a/packages/dd-trace/test/plugins/externals.json +++ b/packages/dd-trace/test/plugins/externals.json @@ -332,7 +332,7 @@ }, { "name": "express", - "versions": [">=4", ">=4.0.0 <4.3.0", ">=4.0.0 <5.0.0", ">=4.3.0 <5.0.0", ">=5.0.0"] + "versions": [">=4", ">=4.0.0 <4.3.0", ">=4.0.0 <5.0.0", ">=4.3.0 <5.0.0", ">=5.0.0"] }, { "name": "body-parser",