fix: safer wrapFunction for ESM instrumentation (#6761)

* Added prevention into wrapFunction

* Modify shimmer tests to comply with new behavior

* Added express-mongo-sanitize-test

Author: pabloerhard

.github/workflows/apm-integrations.yml

@@ -384,6 +384,16 @@ jobs:
       - uses: ./.github/actions/install
       - run: yarn test:plugins:ci
 
+  express-mongo-sanitize:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: express-mongo-sanitize
+    steps:
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8
+      - uses: ./.github/actions/plugins/test
+        with:
+          dd_api_key: ${{ secrets.DD_API_KEY }}
+
   express-session:
     runs-on: ubuntu-latest
     env:

packages/datadog-instrumentations/src/cookie-parser.js

@@ -25,8 +25,6 @@ addHook({
   name: 'cookie-parser',
   versions: ['>=1.0.0']
 }, cookieParser => {
-  // This prevents the non default export from entering the wrapping process
-  if (cookieParser.default) return cookieParser
   return shimmer.wrapFunction(cookieParser, cookieParser => function () {
     const cookieMiddleware = cookieParser.apply(this, arguments)
 

packages/datadog-instrumentations/src/express-session.js

@@ -37,6 +37,5 @@ addHook({
   name: 'express-session',
   versions: ['>=1.5.0']
 }, session => {
-  if (session.default) return session
   return shimmer.wrapFunction(session, wrapSession)
 })

packages/datadog-plugin-express-mongo-sanitize/test/integration-test/client.spec.js

@@ -0,0 +1,43 @@
+'use strict'
+
+const {
+  createSandbox, varySandbox,
+  FakeAgent, spawnPluginIntegrationTestProc
+} = require('../../../../integration-tests/helpers')
+const { assert } = require('chai')
+const { withVersions } = require('../../../dd-trace/test/setup/mocha')
+const axios = require('axios')
+withVersions('express-mongo-sanitize', 'express-mongo-sanitize', version => {
+  describe('ESM', () => {
+    let sandbox, variants, proc, agent
+
+    before(async function () {
+      this.timeout(50000)
+      sandbox = await createSandbox([`'express-mongo-sanitize@${version}'`, 'express@<=4.0.0'], false,
+        ['./packages/datadog-plugin-express-mongo-sanitize/test/integration-test/*'])
+      variants = varySandbox(sandbox, 'server.mjs', 'expressMongoSanitize', undefined, 'express-mongo-sanitize')
+    })
+
+    after(async function () {
+      this.timeout(50000)
+      await sandbox.remove()
+    })
+
+    beforeEach(async () => {
+      agent = await new FakeAgent().start()
+    })
+
+    afterEach(async () => {
+      proc?.kill()
+      await agent.stop()
+    })
+
+    for (const variant of varySandbox.VARIANTS) {
+      it(`is instrumented loaded with ${variant}`, async () => {
+        const proc = await spawnPluginIntegrationTestProc(sandbox.folder, variants[variant], agent.port)
+        const response = await axios.get(`${proc.url}/?param=paramvalue`)
+        assert.equal(response.headers['x-counter'], '1')
+      })
+    }
+  })
+})

packages/datadog-plugin-express-mongo-sanitize/test/integration-test/server.mjs

@@ -0,0 +1,24 @@
+import 'dd-trace/init.js'
+import express from 'express'
+import expressMongoSanitize from 'express-mongo-sanitize'
+import dc from 'dc-polyfill'
+const app = express()
+
+const sanitizeMiddlewareFinished = dc.channel('datadog:express-mongo-sanitize:filter:finish')
+
+let counter = 0
+
+sanitizeMiddlewareFinished.subscribe(() => {
+  counter += 1
+})
+
+app.use(expressMongoSanitize())
+app.all('/', (req, res) => {
+  res.setHeader('X-Counter', counter)
+  res.end()
+})
+
+const server = app.listen(0, () => {
+  const port = server.address().port
+  process.send({ port })
+})

packages/datadog-shimmer/src/shimmer.js

@@ -73,6 +73,8 @@ function copyObjectProperties (original, wrapped, skipKey) {
  * @returns {Function} The wrapped function.
  */
 function wrapFunction (original, wrapper) {
+  if (typeof original === 'object' && original !== null) return original
+
   const wrapped = wrapper(original)
 
   if (typeof original === 'function') {

packages/datadog-shimmer/test/shimmer.spec.js

@@ -372,12 +372,18 @@ describe('shimmer', () => {
       expect(() => shimmer.wrap(() => {}, () => {})).to.throw()
     })
 
-    it('should work without a function', () => {
-      const a = { b: 1 }
+    it('should work with null instead of function', () => {
+      const a = null
       const wrapped = shimmer.wrapFunction(a, x => () => x)
       expect(wrapped()).to.equal(a)
     })
 
+    it('should not work with an object', () => {
+      const a = { b: 1 }
+      const wrapped = shimmer.wrapFunction(a, x => () => x)
+      expect(typeof wrapped).to.not.equal('function')
+    })
+
     it('should wrap the function', () => {
       const count = inc => inc
 

packages/dd-trace/test/plugins/plugin-structure.spec.js

@@ -20,7 +20,8 @@ const missingPlugins = [
   'datadog-plugin-limitd-client', // limitd-client instrumentation handles trace context propagation, no tracing is done
   'datadog-plugin-mongoose', // mongoose tracing is done through mongodb-core instrumentation
   'datadog-plugin-cookie-parser', // cookie-parser does not produce spans
-  'datadog-plugin-express-session' // express-session does not produce spans
+  'datadog-plugin-express-session', // express-session does not produce spans
+  'datadog-plugin-express-mongo-sanitize' // express-mongo-sanitize does not produce spans
 ]
 
 // instrumentations that do not have a hook, but are still instrumented