dd-trace/packages/datadog-instrumentations/src/helpers/hook.js

Datadog APM tracing client for JavaScript

'use strict'

const path = require('path')

const iitm = require('../../../dd-trace/src/iitm')
const ritm = require('../../../dd-trace/src/ritm')
const log = require('../../../dd-trace/src/log')
const requirePackageJson = require('../../../dd-trace/src/require-package-json')

/**
 * @param {string} moduleBaseDir
 * @returns {string|undefined}
 */
function getVersion (moduleBaseDir) {
  if (moduleBaseDir) {
    return requirePackageJson(moduleBaseDir, /** @type {import('module').Module} */ (module)).version
  }

  return process.version
}

/**
 * This is called for every package/internal-module that dd-trace supports instrumentation for
 * In practice, `modules` is always an array with a single entry.
 *
 * @overload
 * @param {string[]} modules list of modules to hook into
 * @param {object} hookOptions hook options
 * @param {Function} onrequire callback to be executed upon encountering module
 */
/**
 * @overload
 * @param {string[]} modules list of modules to hook into
 * @param {object} hookOptions hook options
 * @param {Function} onrequire callback to be executed upon encountering module
 */
function Hook (modules, hookOptions, onrequire) {
  // TODO: Rewrite this to use class syntax. The same should be done for ritm.
  if (!(this instanceof Hook)) return new Hook(modules, hookOptions, onrequire)

  if (typeof hookOptions === 'function') {
    onrequire = hookOptions
    hookOptions = {}
  }

  this._patched = Object.create(null)
  const patched = new WeakMap()

  const safeHook = (moduleExports, moduleName, moduleBaseDir, moduleVersion, isIitm) => {
    const parts = [moduleBaseDir, moduleName].filter(Boolean)
    const filename = path.join(...parts)

    let defaultWrapResult

    const wrappedOnrequire = (moduleExports, ...args) => {
      if (this._patched[filename] && patched.has(moduleExports)) {
        return patched.get(moduleExports)
      }

      const result = onrequire(moduleExports, ...args)
      if (result && (typeof result === 'object' || typeof result === 'function')) {
        patched.set(moduleExports, result)
        patched.set(result, result)
      }

      return result
    }

    try {
      moduleVersion ||= getVersion(moduleBaseDir)
    } catch (error) {
      log.error('Error getting version for "%s": %s', moduleName, error.message, error)
      return
    }

    if (
      isIitm &&
      moduleExports.default &&
      (typeof moduleExports.default === 'object' ||
      typeof moduleExports.default === 'function')
    ) {
      defaultWrapResult = wrappedOnrequire(moduleExports.default, moduleName, moduleBaseDir, moduleVersion, isIitm)
    }

    const newExports = wrappedOnrequire(moduleExports, moduleName, moduleBaseDir, moduleVersion, isIitm)

    if (defaultWrapResult) newExports.default = defaultWrapResult

    this._patched[filename] = true

    return newExports
  }

  this._ritmHook = ritm(modules, {}, safeHook)
  this._iitmHook = iitm(modules, hookOptions, (moduleExports, moduleName, moduleBaseDir) => {
    return safeHook(moduleExports, moduleName, moduleBaseDir, null, true)
  })
}

module.exports = Hook