dd-trace 5.67.0 → 5.69.0

package/packages/datadog-core/src/storage.js

@@ -10,12 +10,13 @@ const { AsyncLocalStorage } = require('async_hooks')
  * a "handle" object, which is used as a key in a WeakMap, where the values
  * are the real store objects.
  *
- * @template T
+ * @typedef {Record<string, unknown>} Store
  */
 class DatadogStorage extends AsyncLocalStorage {
   /**
    *
-   * @param store {DatadogStorage}
+   * @param store {Store}
+   * @override
    */
   enterWith (store) {
     const handle = {}
@@ -34,7 +35,7 @@ class DatadogStorage extends AsyncLocalStorage {
    *
    * TODO: Refactor the Scope class to use a span-only store and remove this.
    *
-   * @returns {{}}
+   * @returns {Store}
    */
   getHandle () {
     return super.getStore()
@@ -46,9 +47,9 @@ class DatadogStorage extends AsyncLocalStorage {
    * retrieved through `getHandle()` can also be passed in to be used as the
    * key. This is useful if you've stashed a handle somewhere and want to
    * retrieve the store with it.
-   *
-   * @param [handle] {{}}
-   * @returns {T | undefined}
+   * @param {{}} [handle]
+   * @returns {Store | undefined}
+   * @override
    */
   getStore (handle) {
     if (!handle) {
@@ -66,11 +67,12 @@ class DatadogStorage extends AsyncLocalStorage {
    * when dealing with the parent store, so that we don't have to access the
    * WeakMap.
    * @template R
-   * @template TArgs extends any[]
-   * @param store {DatadogStorage}
-   * @param fn {() => R}
-   * @param args {TArgs}
-   * @returns {void}
+   * @template TArgs = unknown[]
+   * @param {Store} store
+   * @param {() => R} fn
+   * @param {...TArgs} args
+   * @returns {R}
+   * @override
    */
   run (store, fn, ...args) {
     const prior = super.getStore()
@@ -85,8 +87,7 @@ class DatadogStorage extends AsyncLocalStorage {
 
 /**
  * This is the map from handles to real stores, used in the class above.
- * @template T
- * @type {WeakMap<WeakKey, T>}
+ * @type {WeakMap<WeakKey, Store>}
  */
 const stores = new WeakMap()
 

package/packages/datadog-esbuild/index.js

@@ -8,6 +8,8 @@ const extractPackageAndModulePath = require(
   '../datadog-instrumentations/src/helpers/extract-package-and-module-path.js'
 )
 
+let rewriter
+
 for (const hook of Object.values(hooks)) {
   if (typeof hook === 'object') {
     hook.fn()
@@ -33,6 +35,7 @@ const RAW_BUILTINS = require('module').builtinModules
 const CHANNEL = 'dd-trace:bundler:load'
 const path = require('path')
 const fs = require('fs')
+const { execSync } = require('child_process')
 
 const builtins = new Set()
 
@@ -42,6 +45,7 @@ for (const builtin of RAW_BUILTINS) {
 }
 
 const DEBUG = !!process.env.DD_TRACE_DEBUG
+const DD_IAST_ENABLED = process.env.DD_IAST_ENABLED?.toLowerCase() === 'true' || process.env.DD_IAST_ENABLED === '1'
 
 // We don't want to handle any built-in packages
 // Those packages will still be handled via RITM
@@ -61,11 +65,57 @@ function isESMBuild (build) {
   return format === 'esm' || outputFile?.endsWith('.mjs') || outExtension === '.mjs'
 }
 
+function getGitMetadata () {
+  const gitMetadata = {
+    repositoryURL: null,
+    commitSHA: null
+  }
+
+  try {
+    gitMetadata.repositoryURL = execSync('git config --get remote.origin.url', {
+      encoding: 'utf8',
+      stdio: ['pipe', 'pipe', 'ignore'],
+      cwd: process.cwd()
+    }).trim()
+  } catch (e) {
+    if (DEBUG) {
+      console.warn('Warning: failed to get git repository URL:', e.message)
+    }
+  }
+
+  try {
+    gitMetadata.commitSHA = execSync('git rev-parse HEAD', {
+      encoding: 'utf8',
+      stdio: ['pipe', 'pipe', 'ignore'],
+      cwd: process.cwd()
+    }).trim()
+  } catch (e) {
+    if (DEBUG) {
+      console.warn('Warning: failed to get git commit SHA:', e.message)
+    }
+  }
+
+  return gitMetadata
+}
+
 module.exports.setup = function (build) {
+  if (DD_IAST_ENABLED) {
+    const iastRewriter = require('../dd-trace/src/appsec/iast/taint-tracking/rewriter')
+    rewriter = iastRewriter.getRewriter()
+  }
+
+  const isSourceMapEnabled = !!build.initialOptions.sourcemap ||
+    ['internal', 'both'].includes(build.initialOptions.sourcemap)
   const externalModules = new Set(build.initialOptions.external || [])
+  build.initialOptions.banner ??= {}
+  build.initialOptions.banner.js ??= ''
+  if (DD_IAST_ENABLED) {
+    build.initialOptions.banner.js =
+      `globalThis.__DD_ESBUILD_IAST_${isSourceMapEnabled ? 'WITH_SM' : 'WITH_NO_SM'} = true;
+      ${isSourceMapEnabled ? `globalThis.__DD_ESBUILD_BASEPATH = '${require('../dd-trace/src/util').ddBasePath}';` : ''}
+${build.initialOptions.banner.js}`
+  }
   if (isESMBuild(build)) {
-    build.initialOptions.banner ??= {}
-    build.initialOptions.banner.js ??= ''
     if (!build.initialOptions.banner.js.includes('import { createRequire as $dd_createRequire } from \'module\'')) {
       build.initialOptions.banner.js = `import { createRequire as $dd_createRequire } from 'module';
 import { fileURLToPath as $dd_fileURLToPath } from 'url';
@@ -77,6 +127,28 @@ ${build.initialOptions.banner.js}`
     }
   }
 
+  // Get git metadata at build time and add it to the banner for both ESM and CommonJS builds
+  const gitMetadata = getGitMetadata()
+  if (gitMetadata.repositoryURL || gitMetadata.commitSHA) {
+    build.initialOptions.banner ??= {}
+    build.initialOptions.banner.js ??= ''
+
+    build.initialOptions.banner.js = `if (typeof process === 'object' && process !== null &&
+    process.env !== null && typeof process.env === 'object') {
+  ${gitMetadata.repositoryURL ? `process.env.DD_GIT_REPOSITORY_URL = '${gitMetadata.repositoryURL}';` : ''}
+  ${gitMetadata.commitSHA ? `process.env.DD_GIT_COMMIT_SHA = '${gitMetadata.commitSHA}';` : ''}
+}
+${build.initialOptions.banner.js}`
+
+    if (DEBUG) {
+      console.log('Info: automatically injected git metadata:')
+      console.log(`DD_GIT_REPOSITORY_URL: ${gitMetadata.repositoryURL || 'not available'}`)
+      console.log(`DD_GIT_COMMIT_SHA: ${gitMetadata.commitSHA || 'not available'}`)
+    }
+  } else if (DEBUG) {
+    console.warn('Warning: No git metadata available - skipping injection')
+  }
+
   build.onResolve({ filter: /.*/ }, args => {
     if (externalModules.has(args.path)) {
       // Internal Node.js packages will still be instrumented via require()
@@ -84,13 +156,6 @@ ${build.initialOptions.banner.js}`
       return
     }
 
-    // TODO: Should this also check for namespace === 'file'?
-    if (args.path.startsWith('.') && !args.importer.includes('node_modules/')) {
-      // This is local application code, not an instrumented package
-      if (DEBUG) console.log(`LOCAL: ${args.path}`)
-      return
-    }
-
     // TODO: Should this also check for namespace === 'file'?
     if (!modulesOfInterest.has(args.path) &&
         args.path.startsWith('@') &&
@@ -110,6 +175,21 @@ ${build.initialOptions.banner.js}`
       }
       return
     }
+
+    if (args.path.startsWith('.') && !args.importer.includes('node_modules/')) {
+      // It is local application code, not an instrumented package
+      if (DEBUG) console.log(`APP: ${args.path}`)
+
+      return {
+        path: fullPathToModule,
+        pluginData: {
+          path: args.path,
+          full: fullPathToModule,
+          applicationFile: true
+        }
+      }
+    }
+
     const extracted = extractPackageAndModulePath(fullPathToModule)
 
     const internal = builtins.has(args.path)
@@ -159,22 +239,19 @@ ${build.initialOptions.banner.js}`
   })
 
   build.onLoad({ filter: /.*/ }, args => {
-    if (!args.pluginData?.pkgOfInterest) {
-      return
-    }
-
-    const data = args.pluginData
+    if (args.pluginData?.pkgOfInterest) {
+      const data = args.pluginData
 
-    if (DEBUG) console.log(`LOAD: ${data.pkg}@${data.version}, pkg "${data.path}"`)
+      if (DEBUG) console.log(`LOAD: ${data.pkg}@${data.version}, pkg "${data.path}"`)
 
-    const pkgPath = data.raw !== data.pkg
-      ? `${data.pkg}/${data.path}`
-      : data.pkg
+      const pkgPath = data.raw !== data.pkg
+        ? `${data.pkg}/${data.path}`
+        : data.pkg
 
-    // Read the content of the module file of interest
-    const fileCode = fs.readFileSync(args.path, 'utf8')
+      // Read the content of the module file of interest
+      const fileCode = fs.readFileSync(args.path, 'utf8')
 
-    const contents = `
+      const contents = `
       (function() {
         ${fileCode}
       })(...arguments);
@@ -192,12 +269,27 @@ ${build.initialOptions.banner.js}`
         module.exports = payload.module;
     }
     `
+      // https://esbuild.github.io/plugins/#on-load-results
+      return {
+        contents,
+        loader: 'js',
+        resolveDir: path.dirname(args.path)
+      }
+    }
 
-    // https://esbuild.github.io/plugins/#on-load-results
-    return {
-      contents,
-      loader: 'js',
-      resolveDir: path.dirname(args.path)
+    if (DD_IAST_ENABLED && args.pluginData?.applicationFile) {
+      const ext = path.extname(args.path).toLowerCase()
+      const isJs = /^\.(js|mjs|cjs)$/.test(ext)
+      if (!isJs) return
+
+      if (DEBUG) console.log(`REWRITE: ${args.path}`)
+      const fileCode = fs.readFileSync(args.path, 'utf8')
+      const rewritten = rewriter.rewrite(fileCode, args.path, ['iast'])
+      return {
+        contents: rewritten.content,
+        loader: 'js',
+        resolveDir: path.dirname(args.path)
+      }
     }
   })
 }

package/packages/datadog-instrumentations/src/aws-sdk.js

@@ -61,7 +61,6 @@ function wrapSmithySend (send) {
 
     const startCh = channel(`apm:aws:request:start:${channelSuffix}`)
     const regionCh = channel(`apm:aws:request:region:${channelSuffix}`)
-    const completeChannel = channel(`apm:aws:request:complete:${channelSuffix}`)
     const responseStartChannel = channel(`apm:aws:response:start:${channelSuffix}`)
     const responseFinishChannel = channel(`apm:aws:response:finish:${channelSuffix}`)
 
@@ -87,7 +86,7 @@ function wrapSmithySend (send) {
         args[args.length - 1] = shimmer.wrapFunction(cb, cb => function (err, result) {
           addResponse(ctx, err, result)
 
-          completeChannel.publish(ctx)
+          handleCompletion(result, ctx, channelSuffix)
 
           const responseCtx = { request, response: ctx.response }
 
@@ -102,12 +101,12 @@ function wrapSmithySend (send) {
           .then(
             result => {
               addResponse(ctx, null, result)
-              completeChannel.publish(ctx)
+              handleCompletion(result, ctx, channelSuffix)
               return result
             },
             error => {
               addResponse(ctx, error)
-              completeChannel.publish(ctx)
+              handleCompletion(null, ctx, channelSuffix)
               throw error
             }
           )
@@ -118,6 +117,45 @@ function wrapSmithySend (send) {
   }
 }
 
+function handleCompletion (result, ctx, channelSuffix) {
+  const completeChannel = channel(`apm:aws:request:complete:${channelSuffix}`)
+  const streamedChunkChannel = channel(`apm:aws:response:streamed-chunk:${channelSuffix}`)
+
+  const iterator = result?.body?.[Symbol.asyncIterator]
+  if (!iterator) {
+    completeChannel.publish(ctx)
+    return
+  }
+
+  shimmer.wrap(result.body, Symbol.asyncIterator, function (asyncIterator) {
+    return function () {
+      const iterator = asyncIterator.apply(this, arguments)
+      shimmer.wrap(iterator, 'next', function (next) {
+        return function () {
+          return next.apply(this, arguments)
+            .then(result => {
+              const { done, value: chunk } = result
+              streamedChunkChannel.publish({ ctx, chunk, done })
+
+              if (done) {
+                completeChannel.publish(ctx)
+              }
+
+              return result
+            })
+            .catch(err => {
+              addResponse(ctx, err)
+              completeChannel.publish(ctx)
+              throw err
+            })
+        }
+      })
+
+      return iterator
+    }
+  })
+}
+
 function wrapCb (cb, serviceName, ctx) {
   // eslint-disable-next-line n/handle-callback-err
   return shimmer.wrapFunction(cb, cb => function wrappedCb (err, response) {

package/packages/datadog-instrumentations/src/azure-functions.js

@@ -8,7 +8,7 @@ const dc = require('dc-polyfill')
 
 const azureFunctionsChannel = dc.tracingChannel('datadog:azure:functions:invoke')
 
-addHook({ name: '@azure/functions', versions: ['>=4'] }, azureFunction => {
+addHook({ name: '@azure/functions', versions: ['>=4'], patchDefault: false }, (azureFunction) => {
   const { app } = azureFunction
 
   // Http triggers

package/packages/datadog-instrumentations/src/azure-service-bus.js

@@ -11,7 +11,7 @@ const producerStartCh = channel('apm:azure-service-bus:send:start')
 const producerErrorCh = channel('apm:azure-service-bus:send:error')
 const producerFinishCh = channel('apm:azure-service-bus:send:finish')
 
-addHook({ name: '@azure/service-bus', versions: ['>=7.9.2'] }, (obj) => {
+addHook({ name: '@azure/service-bus', versions: ['>=7.9.2'], patchDefault: false }, (obj) => {
   const ServiceBusClient = obj.ServiceBusClient
   shimmer.wrap(ServiceBusClient.prototype, 'createSender', createSender => function (queueOrTopicName) {
     const sender = createSender.apply(this, arguments)

package/packages/datadog-instrumentations/src/cassandra-driver.js

@@ -45,7 +45,7 @@ addHook({ name: 'cassandra-driver', versions: ['>=3.0.0'] }, cassandra => {
   return cassandra
 })
 
-addHook({ name: 'cassandra-driver', versions: ['>=4.4'] }, cassandra => {
+addHook({ name: 'cassandra-driver', versions: ['>=4.4'], patchDefault: false }, (cassandra) => {
   shimmer.wrap(cassandra.Client.prototype, '_execute', _execute => function (query, params, execOptions, callback) {
     if (!startCh.hasSubscribers) {
       return _execute.apply(this, arguments)
@@ -68,7 +68,7 @@ const isValid = (args) => {
   return args.length === 4 || typeof args[3] === 'function'
 }
 
-addHook({ name: 'cassandra-driver', versions: ['3 - 4.3'] }, cassandra => {
+addHook({ name: 'cassandra-driver', versions: ['3 - 4.3'], patchDefault: false }, (cassandra) => {
   shimmer.wrap(cassandra.Client.prototype, '_innerExecute', _innerExecute =>
     function (query, params, execOptions, callback) {
       if (!startCh.hasSubscribers) {

package/packages/datadog-instrumentations/src/connect.js

@@ -102,11 +102,15 @@ function wrapNext (req, next) {
   })
 }
 
-addHook({ name: 'connect', versions: ['>=3'] }, connect => {
+addHook({ name: 'connect', versions: ['>=3.4.0'] }, (connect) => {
   return shimmer.wrapFunction(connect, connect => wrapConnect(connect))
 })
 
-addHook({ name: 'connect', versions: ['2.2.2'] }, connect => {
+addHook({ name: 'connect', versions: ['>=3 <3.4.0'], file: 'lib/connect.js' }, (connect) => {
+  return shimmer.wrapFunction(connect, connect => wrapConnect(connect))
+})
+
+addHook({ name: 'connect', versions: ['2.2.2'], file: 'lib/connect.js' }, connect => {
   shimmer.wrap(connect.proto, 'use', wrapUse)
   shimmer.wrap(connect.proto, 'handle', wrapHandle)
 

package/packages/datadog-instrumentations/src/cucumber.js

@@ -82,10 +82,14 @@ let testManagementAttemptToFixRetries = 0
 let testManagementTests = {}
 let modifiedTests = {}
 let numTestRetries = 0
-let knownTests = []
+let knownTests = {}
 let skippedSuites = []
 let isSuitesSkipped = false
 
+function isValidKnownTests (receivedKnownTests) {
+  return !!receivedKnownTests.cucumber
+}
+
 function getSuiteStatusFromTestStatuses (testStatuses) {
   if (testStatuses.includes('fail')) {
     return 'fail'
@@ -123,7 +127,10 @@ function getStatusFromResultLatest (result) {
 }
 
 function isNewTest (testSuite, testName) {
-  const testsForSuite = knownTests.cucumber?.[testSuite] || []
+  if (!isValidKnownTests(knownTests)) {
+    return false
+  }
+  const testsForSuite = knownTests.cucumber[testSuite] || []
   return !testsForSuite.includes(testName)
 }
 
@@ -508,9 +515,9 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     pickleByFile = isCoordinator ? getPickleByFileNew(this) : getPickleByFile(this)
 
     if (isKnownTestsEnabled) {
-      const isFaulty = getIsFaultyEarlyFlakeDetection(
+      const isFaulty = !isValidKnownTests(knownTests) || getIsFaultyEarlyFlakeDetection(
         Object.keys(pickleByFile),
-        knownTests.cucumber || {},
+        knownTests.cucumber,
         earlyFlakeDetectionFaultyThreshold
       )
       if (isFaulty) {
@@ -592,6 +599,9 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
 // Handles EFD in both the main process and the worker process.
 function getWrappedRunTestCase (runTestCaseFunction, isNewerCucumberVersion = false, isWorker = false) {
   return async function () {
+    if (!testSuiteFinishCh.hasSubscribers) {
+      return runTestCaseFunction.apply(this, arguments)
+    }
     const pickle = isNewerCucumberVersion
       ? arguments[0].pickle
       : this.eventDataCollector.getPickle(arguments[0])
@@ -747,6 +757,9 @@ function getWrappedRunTestCase (runTestCaseFunction, isNewerCucumberVersion = fa
 
 function getWrappedParseWorkerMessage (parseWorkerMessageFunction, isNewVersion) {
   return function (worker, message) {
+    if (!testSuiteFinishCh.hasSubscribers) {
+      return parseWorkerMessageFunction.apply(this, arguments)
+    }
     // If the message is an array, it's a dd-trace message, so we need to stop cucumber processing,
     // or cucumber will throw an error
     // TODO: identify the message better
@@ -972,10 +985,17 @@ addHook({
   )
   // EFD in parallel mode only supported in >=11.0.0
   shimmer.wrap(adapterPackage.ChildProcessAdapter.prototype, 'startWorker', startWorker => function () {
-    if (isKnownTestsEnabled) {
+    if (isKnownTestsEnabled && isValidKnownTests(knownTests)) {
+      this.options.worldParameters._ddIsKnownTestsEnabled = true
       this.options.worldParameters._ddIsEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
       this.options.worldParameters._ddKnownTests = knownTests
       this.options.worldParameters._ddEarlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
+    } else {
+      isEarlyFlakeDetectionEnabled = false
+      isKnownTestsEnabled = false
+      this.options.worldParameters._ddIsEarlyFlakeDetectionEnabled = false
+      this.options.worldParameters._ddIsKnownTestsEnabled = false
+      this.options.worldParameters._ddEarlyFlakeDetectionNumRetries = 0
     }
 
     if (isImpactedTestsEnabled) {
@@ -1001,9 +1021,14 @@ addHook({
     'initialize',
     initialize => async function () {
       await initialize.apply(this, arguments)
-      isKnownTestsEnabled = !!this.options.worldParameters._ddKnownTests
+      isKnownTestsEnabled = !!this.options.worldParameters._ddIsKnownTestsEnabled
       if (isKnownTestsEnabled) {
         knownTests = this.options.worldParameters._ddKnownTests
+        // if for whatever reason the worker does not receive valid known tests, we disable EFD and known tests
+        if (!isValidKnownTests(knownTests)) {
+          isKnownTestsEnabled = false
+          knownTests = {}
+        }
       }
       isEarlyFlakeDetectionEnabled = !!this.options.worldParameters._ddIsEarlyFlakeDetectionEnabled
       if (isEarlyFlakeDetectionEnabled) {

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

@@ -2,8 +2,7 @@
 
 const { createWrapRouterMethod } = require('./router')
 const shimmer = require('../../datadog-shimmer')
-const { addHook, channel } = require('./helpers/instrument')
-const tracingChannel = require('dc-polyfill').tracingChannel
+const { addHook, channel, tracingChannel } = require('./helpers/instrument')
 
 const handleChannel = channel('apm:express:request:handle')
 
@@ -57,7 +56,7 @@ function wrapResponseRender (render) {
   }
 }
 
-addHook({ name: 'express', versions: ['>=4'] }, express => {
+addHook({ name: 'express', versions: ['>=4'], file: ['lib/express.js'] }, express => {
   shimmer.wrap(express.application, 'handle', wrapHandle)
 
   shimmer.wrap(express.response, 'json', wrapResponseJson)
@@ -70,7 +69,7 @@ addHook({ name: 'express', versions: ['>=4'] }, express => {
 // Express 5 does not rely on router in the same way as v4 and should not be instrumented anymore.
 // It would otherwise produce spans for router and express, and so duplicating them.
 // We now fall back to router instrumentation
-addHook({ name: 'express', versions: ['4'] }, express => {
+addHook({ name: 'express', versions: ['4'], file: 'lib/express.js' }, express => {
   shimmer.wrap(express.Router, 'use', wrapRouterMethod)
   shimmer.wrap(express.Router, 'route', wrapRouterMethod)
 
@@ -132,12 +131,12 @@ function wrapProcessParamsMethod (requestPositionInArguments) {
   }
 }
 
-addHook({ name: 'express', versions: ['>=4.0.0 <4.3.0'] }, express => {
+addHook({ name: 'express', versions: ['>=4.0.0 <4.3.0'], file: ['lib/express.js'] }, express => {
   shimmer.wrap(express.Router, 'process_params', wrapProcessParamsMethod(1))
   return express
 })
 
-addHook({ name: 'express', versions: ['>=4.3.0 <5.0.0'] }, express => {
+addHook({ name: 'express', versions: ['>=4.3.0 <5.0.0'], file: ['lib/express.js'] }, express => {
   shimmer.wrap(express.Router, 'process_params', wrapProcessParamsMethod(2))
   return express
 })

package/packages/datadog-instrumentations/src/fastify.js

@@ -265,7 +265,7 @@ function canPublishResponsePayload (payload) {
     !ArrayBuffer.isView(payload) // TypedArray
 }
 
-addHook({ name: 'fastify', versions: ['>=3'] }, fastify => {
+addHook({ name: 'fastify', versions: ['>=3'] }, (fastify) => {
   const wrapped = shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, true))
 
   wrapped.fastify = wrapped
@@ -274,11 +274,11 @@ addHook({ name: 'fastify', versions: ['>=3'] }, fastify => {
   return wrapped
 })
 
-addHook({ name: 'fastify', versions: ['2'] }, fastify => {
+addHook({ name: 'fastify', versions: ['2'] }, (fastify) => {
   return shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, true))
 })
 
-addHook({ name: 'fastify', versions: ['1'] }, fastify => {
+addHook({ name: 'fastify', versions: ['1'] }, (fastify) => {
   return shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, false))
 })
 

package/packages/datadog-instrumentations/src/helpers/hook.js

@@ -1,7 +1,6 @@
 'use strict'
-
-const path = require('path')
 const iitm = require('../../../dd-trace/src/iitm')
+const path = require('path')
 const ritm = require('../../../dd-trace/src/ritm')
 
 /**
@@ -20,29 +19,43 @@ function Hook (modules, hookOptions, onrequire) {
   }
 
   this._patched = Object.create(null)
+  const patched = new WeakMap()
 
-  const safeHook = (moduleExports, moduleName, moduleBaseDir, moduleVersion) => {
+  const safeHook = (moduleExports, moduleName, moduleBaseDir, moduleVersion, isIitm) => {
     const parts = [moduleBaseDir, moduleName].filter(Boolean)
     const filename = path.join(...parts)
 
-    if (this._patched[filename]) return moduleExports
+    if (this._patched[filename] && patched.has(moduleExports)) {
+      return patched.get(moduleExports)
+    }
 
-    this._patched[filename] = true
+    let defaultWrapResult
+
+    if (
+      isIitm &&
+      moduleExports.default &&
+      (typeof moduleExports.default === 'object' ||
+      typeof moduleExports.default === 'function')
+    ) {
+      defaultWrapResult = onrequire(moduleExports.default, moduleName, moduleBaseDir, moduleVersion, isIitm)
+    }
+
+    const newExports = onrequire(moduleExports, moduleName, moduleBaseDir, moduleVersion, isIitm)
+
+    if (defaultWrapResult) newExports.default = defaultWrapResult
 
-    return onrequire(moduleExports, moduleName, moduleBaseDir, moduleVersion)
+    this._patched[filename] = true
+    if (newExports &&
+      (typeof newExports === 'object' ||
+      typeof newExports === 'function')) {
+      patched.set(moduleExports, newExports)
+    }
+    return newExports
   }
 
   this._ritmHook = ritm(modules, {}, safeHook)
   this._iitmHook = iitm(modules, hookOptions, (moduleExports, moduleName, moduleBaseDir) => {
-    // TODO: Move this logic to import-in-the-middle and only do it for CommonJS
-    // modules and not ESM. In the meantime, all the modules we instrument are
-    // CommonJS modules for which the default export is always moved to
-    // `default` anyway.
-    if (moduleExports && moduleExports.default) {
-      moduleExports.default = safeHook(moduleExports.default, moduleName, moduleBaseDir)
-      return moduleExports
-    }
-    return safeHook(moduleExports, moduleName, moduleBaseDir)
+    return safeHook(moduleExports, moduleName, moduleBaseDir, null, true)
   })
 }
 

package/packages/datadog-instrumentations/src/helpers/hooks.js

@@ -16,6 +16,7 @@ module.exports = {
   '@graphql-tools/executor': () => require('../graphql'),
   '@grpc/grpc-js': () => require('../grpc'),
   '@hapi/hapi': () => require('../hapi'),
+  '@happy-dom/jest-environment': () => require('../jest'),
   '@jest/core': () => require('../jest'),
   '@jest/reporters': () => require('../jest'),
   '@jest/test-sequencer': () => require('../jest'),
@@ -130,6 +131,7 @@ module.exports = {
   sequelize: () => require('../sequelize'),
   sharedb: () => require('../sharedb'),
   tedious: () => require('../tedious'),
+  tinypool: { esmFirst: true, fn: () => require('../vitest') },
   undici: () => require('../undici'),
   url: () => require('../url'),
   vitest: { esmFirst: true, fn: () => require('../vitest') },