Safer shimming

packages/datadog-instrumentations/src/aerospike.js

@@ -43,5 +43,5 @@ addHook({
   versions: ['^3.16.2', '4', '5']
 },
 commandFactory => {
-  return shimmer.wrap(commandFactory, wrapCreateCommand(commandFactory))
+  return shimmer.wrapFunction(commandFactory, f => wrapCreateCommand(f))
 })

packages/datadog-instrumentations/src/apollo-server.js

@@ -55,7 +55,7 @@ function apolloExpress4Hook (express4) {
     return function expressMiddleware (server, options) {
       const originalMiddleware = originalExpressMiddleware.apply(this, arguments)
 
-      return shimmer.wrap(originalMiddleware, function (req, res, next) {
+      return shimmer.wrapFunction(originalMiddleware, originalMiddleware => function (req, res, next) {
         if (!graphqlMiddlewareChannel.start.hasSubscribers) {
           return originalMiddleware.apply(this, arguments)
         }

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

@@ -25,7 +25,7 @@ addHook({
   file: 'lib/read.js',
   versions: ['>=1.4.0 <1.20.0']
 }, read => {
-  return shimmer.wrap(read, function (req, res, next) {
+  return shimmer.wrapFunction(read, read => function (req, res, next) {
     const nextResource = new AsyncResource('bound-anonymous-fn')
     arguments[2] = nextResource.bind(publishRequestBodyAndNext(req, res, next))
     return read.apply(this, arguments)
@@ -37,7 +37,7 @@ addHook({
   file: 'lib/read.js',
   versions: ['>=1.20.0']
 }, read => {
-  return shimmer.wrap(read, function (req, res, next) {
+  return shimmer.wrapFunction(read, read => function (req, res, next) {
     arguments[2] = publishRequestBodyAndNext(req, res, next)
     return read.apply(this, arguments)
   })

packages/datadog-instrumentations/src/child_process.js

@@ -133,8 +133,8 @@ function wrapChildProcessAsyncMethod (shell = false) {
 
     if (childProcessMethod[util.promisify.custom]) {
       const wrapedChildProcessCustomPromisifyMethod =
-        shimmer.wrap(childProcessMethod[util.promisify.custom],
-          wrapChildProcessCustomPromisifyMethod(childProcessMethod[util.promisify.custom]), shell)
+        shimmer.wrapFunction(childProcessMethod[util.promisify.custom],
+          promisify => wrapChildProcessCustomPromisifyMethod(promisify, shell))
 
       // should do it in this way because the original property is readonly
       const descriptor = Object.getOwnPropertyDescriptor(childProcessMethod, util.promisify.custom)

packages/datadog-instrumentations/src/connect.js

@@ -59,7 +59,7 @@ function wrapLayerHandle (layer) {
 
   const original = layer.handle
 
-  return shimmer.wrap(original, function () {
+  return shimmer.wrapFunction(original, original => function () {
     if (!enterChannel.hasSubscribers) return original.apply(this, arguments)
 
     const lastIndex = arguments.length - 1
@@ -103,7 +103,7 @@ function wrapNext (req, next) {
 }
 
 addHook({ name: 'connect', versions: ['>=3'] }, connect => {
-  return shimmer.wrap(connect, wrapConnect(connect))
+  return shimmer.wrapFunction(connect, connect => wrapConnect(connect))
 })
 
 addHook({ name: 'connect', versions: ['2.2.2'] }, connect => {

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

@@ -25,10 +25,10 @@ addHook({
   name: 'cookie-parser',
   versions: ['>=1.0.0']
 }, cookieParser => {
-  return shimmer.wrap(cookieParser, function () {
+  return shimmer.wrapFunction(cookieParser, cookieParser => function () {
     const cookieMiddleware = cookieParser.apply(this, arguments)
 
-    return shimmer.wrap(cookieMiddleware, function (req, res, next) {
+    return shimmer.wrapFunction(cookieMiddleware, cookieMiddleware => function (req, res, next) {
       arguments[2] = publishRequestCookieAndNext(req, res, next)
       return cookieMiddleware.apply(this, arguments)
     })

packages/datadog-instrumentations/src/couchbase.js

@@ -37,7 +37,7 @@ function wrapMaybeInvoke (_maybeInvoke) {
 
     return _maybeInvoke.apply(this, arguments)
   }
-  return shimmer.wrap(_maybeInvoke, wrapped)
+  return wrapped
 }
 
 function wrapQuery (query) {
@@ -51,7 +51,7 @@ function wrapQuery (query) {
     const res = query.apply(this, arguments)
     return res
   }
-  return shimmer.wrap(query, wrapped)
+  return wrapped
 }
 
 function wrap (prefix, fn) {
@@ -94,7 +94,7 @@ function wrap (prefix, fn) {
       }
     })
   }
-  return shimmer.wrap(fn, wrapped)
+  return wrapped
 }
 
 // semver >=3
@@ -166,8 +166,8 @@ addHook({ name: 'couchbase', file: 'lib/bucket.js', versions: ['^2.6.12'] }, Buc
   const finishCh = channel('apm:couchbase:query:finish')
   const errorCh = channel('apm:couchbase:query:error')
 
-  Bucket.prototype._maybeInvoke = wrapMaybeInvoke(Bucket.prototype._maybeInvoke)
-  Bucket.prototype.query = wrapQuery(Bucket.prototype.query)
+  shimmer.wrap(Bucket.prototype, '_maybeInvoke', maybeInvoke => wrapMaybeInvoke(maybeInvoke))
+  shimmer.wrap(Bucket.prototype, 'query', query => wrapQuery(query))
 
   shimmer.wrap(Bucket.prototype, '_n1qlReq', _n1qlReq => function (host, q, adhoc, emitter) {
     if (!startCh.hasSubscribers) {
@@ -203,15 +203,15 @@ addHook({ name: 'couchbase', file: 'lib/bucket.js', versions: ['^2.6.12'] }, Buc
   })
 
   wrapAllNames(['upsert', 'insert', 'replace', 'append', 'prepend'], name => {
-    Bucket.prototype[name] = wrap(`apm:couchbase:${name}`, Bucket.prototype[name])
+    shimmer.wrap(Bucket.prototype, name, fn => wrap(`apm:couchbase:${name}`, fn))
   })
 
   return Bucket
 })
 
 addHook({ name: 'couchbase', file: 'lib/cluster.js', versions: ['^2.6.12'] }, Cluster => {
-  Cluster.prototype._maybeInvoke = wrapMaybeInvoke(Cluster.prototype._maybeInvoke)
-  Cluster.prototype.query = wrapQuery(Cluster.prototype.query)
+  shimmer.wrap(Cluster.prototype, '_maybeInvoke', maybeInvoke => wrapMaybeInvoke(maybeInvoke))
+  shimmer.wrap(Cluster.prototype, 'query', query => wrapQuery(query))
 
   shimmer.wrap(Cluster.prototype, 'openBucket', openBucket => {
     return function () {

packages/datadog-instrumentations/src/dns.js

@@ -21,16 +21,16 @@ const rrtypeMap = new WeakMap()
 const names = ['dns', 'node:dns']
 
 addHook({ name: names }, dns => {
-  dns.lookup = wrap('apm:dns:lookup', dns.lookup, 2)
-  dns.lookupService = wrap('apm:dns:lookup_service', dns.lookupService, 3)
-  dns.resolve = wrap('apm:dns:resolve', dns.resolve, 2)
-  dns.reverse = wrap('apm:dns:reverse', dns.reverse, 2)
+  shimmer.wrap(dns, 'lookup', fn => wrap('apm:dns:lookup', fn, 2))
+  shimmer.wrap(dns, 'lookupService', fn => wrap('apm:dns:lookup_service', fn, 2))
+  shimmer.wrap(dns, 'resolve', fn => wrap('apm:dns:resolve', fn, 2))
+  shimmer.wrap(dns, 'reverse', fn => wrap('apm:dns:reverse', fn, 2))
 
   patchResolveShorthands(dns)
 
   if (dns.Resolver) {
-    dns.Resolver.prototype.resolve = wrap('apm:dns:resolve', dns.Resolver.prototype.resolve, 2)
-    dns.Resolver.prototype.reverse = wrap('apm:dns:reverse', dns.Resolver.prototype.reverse, 2)
+    shimmer.wrap(dns.Resolver.prototype, 'resolve', fn => wrap('apm:dns:resolve', fn, 2))
+    shimmer.wrap(dns.Resolver.prototype, 'reverse', fn => wrap('apm:dns:reverse', fn, 2))
 
     patchResolveShorthands(dns.Resolver.prototype)
   }
@@ -43,7 +43,7 @@ function patchResolveShorthands (prototype) {
     .filter(method => !!prototype[method])
     .forEach(method => {
       rrtypeMap.set(prototype[method], rrtypes[method])
-      prototype[method] = wrap('apm:dns:resolve', prototype[method], 2, rrtypes[method])
+      shimmer.wrap(prototype, method, fn => wrap('apm:dns:resolve', fn, 2, rrtypes[method]))
     })
 }
 
@@ -92,5 +92,5 @@ function wrap (prefix, fn, expectedArgs, rrtype) {
     })
   }
 
-  return shimmer.wrap(fn, wrapped)
+  return wrapped
 }

packages/datadog-instrumentations/src/express-mongo-sanitize.js

@@ -22,15 +22,15 @@ addHook({ name: 'express-mongo-sanitize', versions: ['>=1.0.0'] }, expressMongoS
     return sanitizedObject
   })
 
-  return shimmer.wrap(expressMongoSanitize, function () {
+  return shimmer.wrapFunction(expressMongoSanitize, expressMongoSanitize => function () {
     const middleware = expressMongoSanitize.apply(this, arguments)
 
-    return shimmer.wrap(middleware, function (req, res, next) {
+    return shimmer.wrapFunction(middleware, middleware => function (req, res, next) {
       if (!sanitizeMiddlewareFinished.hasSubscribers) {
         return middleware.apply(this, arguments)
       }
 
-      const wrappedNext = shimmer.wrap(next, function () {
+      const wrappedNext = shimmer.wrapFunction(next, next => function () {
         sanitizeMiddlewareFinished.publish({
           sanitizedProperties: propertiesToSanitize,
           req

packages/datadog-instrumentations/src/express.js

@@ -68,10 +68,10 @@ addHook({
   versions: ['>=4'],
   file: 'lib/middleware/query.js'
 }, query => {
-  return shimmer.wrap(query, function () {
+  return shimmer.wrapFunction(query, query => function () {
     const queryMiddleware = query.apply(this, arguments)
 
-    return shimmer.wrap(queryMiddleware, function (req, res, next) {
+    return shimmer.wrapFunction(queryMiddleware, queryMiddleware => function (req, res, next) {
       arguments[2] = publishQueryParsedAndNext(req, res, next)
       return queryMiddleware.apply(this, arguments)
     })

packages/datadog-instrumentations/src/fastify.js

@@ -39,7 +39,7 @@ function wrapAddHook (addHook) {
 
     if (typeof fn !== 'function') return addHook.apply(this, arguments)
 
-    arguments[arguments.length - 1] = shimmer.wrap(fn, function (request, reply, done) {
+    arguments[arguments.length - 1] = shimmer.wrapFunction(fn, fn => function (request, reply, done) {
       const req = getReq(request)
 
       try {
@@ -151,7 +151,7 @@ function publishError (error, req) {
 }
 
 addHook({ name: 'fastify', versions: ['>=3'] }, fastify => {
-  const wrapped = shimmer.wrap(fastify, wrapFastify(fastify, true))
+  const wrapped = shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, true))
 
   wrapped.fastify = wrapped
   wrapped.default = wrapped
@@ -160,9 +160,9 @@ addHook({ name: 'fastify', versions: ['>=3'] }, fastify => {
 })
 
 addHook({ name: 'fastify', versions: ['2'] }, fastify => {
-  return shimmer.wrap(fastify, wrapFastify(fastify, true))
+  return shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, true))
 })
 
 addHook({ name: 'fastify', versions: ['1'] }, fastify => {
-  return shimmer.wrap(fastify, wrapFastify(fastify, false))
+  return shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, false))
 })

packages/datadog-instrumentations/src/fetch.js

@@ -8,5 +8,5 @@ if (globalThis.fetch) {
   const ch = tracingChannel('apm:fetch:request')
   const wrapFetch = createWrapFetch(globalThis.Request, ch)
 
-  globalThis.fetch = shimmer.wrap(fetch, wrapFetch(fetch))
+  globalThis.fetch = shimmer.wrapFunction(fetch, fetch => wrapFetch(fetch))
 }

packages/datadog-instrumentations/src/jest.js

@@ -638,7 +638,7 @@ addHook({
 
 function jestAdapterWrapper (jestAdapter, jestVersion) {
   const adapter = jestAdapter.default ? jestAdapter.default : jestAdapter
-  const newAdapter = shimmer.wrap(adapter, function () {
+  const newAdapter = shimmer.wrapFunction(adapter, adapter => function () {
     const environment = arguments[2]
     if (!environment) {
       return adapter.apply(this, arguments)

packages/datadog-instrumentations/src/koa.js

@@ -71,7 +71,7 @@ function wrapStack (layer) {
 
     middleware = original || middleware
 
-    const handler = shimmer.wrap(middleware, wrapMiddleware(middleware, layer))
+    const handler = shimmer.wrapFunction(middleware, middleware => wrapMiddleware(middleware, layer))
 
     originals.set(handler, middleware)
 

packages/datadog-instrumentations/src/ldapjs.js

@@ -77,7 +77,7 @@ addHook({ name: 'ldapjs', versions: ['>=2'] }, ldapjs => {
     if (callbackIndex > -1) {
       const callback = arguments[callbackIndex]
       // eslint-disable-next-line n/handle-callback-err
-      arguments[callbackIndex] = shimmer.wrap(callback, function (err, corkedEmitter) {
+      arguments[callbackIndex] = shimmer.wrapFunction(callback, callback => function (err, corkedEmitter) {
         if (corkedEmitter !== null && typeof corkedEmitter === 'object' && typeof corkedEmitter.on === 'function') {
           wrapEmitter(corkedEmitter)
         }

packages/datadog-instrumentations/src/mariadb.js

@@ -119,7 +119,7 @@ function createWrapQueryCallback (options) {
   }
 }
 
-function wrapConnection (Connection, promiseMethod) {
+function wrapConnection (promiseMethod, Connection) {
   return function (options) {
     Connection.apply(this, arguments)
 
@@ -170,13 +170,13 @@ addHook({ name, file: 'lib/pool.js', versions: ['>=3'] }, (Pool) => {
 })
 
 addHook({ name, file: 'lib/connection.js', versions: ['>=2.5.2 <3'] }, (Connection) => {
-  return shimmer.wrap(Connection, wrapConnection(Connection, '_queryPromise'))
+  return shimmer.wrapFunction(Connection, wrapConnection.bind(null, '_queryPromise'))
 })
 
 addHook({ name, file: 'lib/connection.js', versions: ['>=2.0.4 <=2.5.1'] }, (Connection) => {
-  return shimmer.wrap(Connection, wrapConnection(Connection, 'query'))
+  return shimmer.wrapFunction(Connection, wrapConnection.bind(null, 'query'))
 })
 
 addHook({ name, file: 'lib/pool-base.js', versions: ['>=2.0.4 <3'] }, (PoolBase) => {
-  return shimmer.wrap(PoolBase, wrapPoolBase(PoolBase))
+  return shimmer.wrapFunction(PoolBase, wrapPoolBase)
 })

packages/datadog-instrumentations/src/microgateway-core.js

@@ -58,9 +58,9 @@ function wrapNext (req, res, next) {
 }
 
 addHook({ name, versions, file: 'lib/config-proxy-middleware.js' }, configProxyFactory => {
-  return shimmer.wrap(configProxyFactory, wrapConfigProxyFactory(configProxyFactory))
+  return shimmer.wrapFunction(configProxyFactory, wrapConfigProxyFactory)
 })
 
 addHook({ name, versions, file: 'lib/plugins-middleware.js' }, pluginsFactory => {
-  return shimmer.wrap(pluginsFactory, wrapPluginsFactory(pluginsFactory))
+  return shimmer.wrapFunction(pluginsFactory, wrapPluginsFactory)
 })

packages/datadog-instrumentations/src/mocha/common.js

@@ -15,7 +15,7 @@ addHook({
 
   patched.add(mochaEach)
 
-  return shimmer.wrap(mochaEach, function () {
+  return shimmer.wrapFunction(mochaEach, mochaEach => function () {
     const [params] = arguments
     const { it, ...rest } = mochaEach.apply(this, arguments)
     return {

packages/datadog-instrumentations/src/mongodb-core.js

@@ -92,7 +92,7 @@ function wrapUnifiedCommand (command, operation, name) {
     }
     return instrument(operation, command, this, arguments, server, ns, ops, { name })
   }
-  return shimmer.wrap(command, wrapped)
+  return wrapped
 }
 
 function wrapConnectionCommand (command, operation, name, instrumentFn = instrument) {
@@ -109,7 +109,7 @@ function wrapConnectionCommand (command, operation, name, instrumentFn = instrum
     ns = `${ns.db}.${ns.collection}`
     return instrumentFn(operation, command, this, arguments, topology, ns, ops, { name })
   }
-  return shimmer.wrap(command, wrapped)
+  return wrapped
 }
 
 function wrapQuery (query, operation, name) {
@@ -123,7 +123,7 @@ function wrapQuery (query, operation, name) {
     return instrument(operation, query, this, arguments, pool, ns, ops)
   }
 
-  return shimmer.wrap(query, wrapped)
+  return wrapped
 }
 
 function wrapCursor (cursor, operation, name) {
@@ -135,7 +135,7 @@ function wrapCursor (cursor, operation, name) {
     const ns = this.ns
     return instrument(operation, cursor, this, arguments, pool, ns, {}, { name })
   }
-  return shimmer.wrap(cursor, wrapped)
+  return wrapped
 }
 
 function wrapCommand (command, operation, name) {
@@ -145,7 +145,7 @@ function wrapCommand (command, operation, name) {
     }
     return instrument(operation, command, this, arguments, this, ns, ops, { name })
   }
-  return shimmer.wrap(command, wrapped)
+  return wrapped
 }
 
 function instrument (operation, command, ctx, args, server, ns, ops, options = {}) {

packages/datadog-instrumentations/src/mongoose.js

@@ -169,7 +169,7 @@ addHook({
   versions: ['6', '>=7'],
   file: 'lib/helpers/query/sanitizeFilter.js'
 }, sanitizeFilter => {
-  return shimmer.wrap(sanitizeFilter, function wrappedSanitizeFilter () {
+  return shimmer.wrapFunction(sanitizeFilter, sanitizeFilter => function wrappedSanitizeFilter () {
     const sanitizedObject = sanitizeFilter.apply(this, arguments)
 
     if (sanitizeFilterFinishCh.hasSubscribers) {

packages/datadog-instrumentations/src/mysql.js

@@ -92,7 +92,7 @@ addHook({ name: 'mysql', file: 'lib/Pool.js', versions: ['>=2'] }, Pool => {
 
       const cb = arguments[arguments.length - 1]
       if (typeof cb === 'function') {
-        arguments[arguments.length - 1] = shimmer.wrap(cb, function () {
+        arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => function () {
           finish()
           return cb.apply(this, arguments)
         })

packages/datadog-instrumentations/src/passport-http.js

@@ -9,7 +9,7 @@ addHook({
   file: 'lib/passport-http/strategies/basic.js',
   versions: ['>=0.3.0']
 }, BasicStrategy => {
-  return shimmer.wrap(BasicStrategy, function () {
+  return shimmer.wrapFunction(BasicStrategy, BasicStrategy => function () {
     const type = 'http'
 
     if (typeof arguments[0] === 'function') {

packages/datadog-instrumentations/src/passport-local.js

@@ -9,7 +9,7 @@ addHook({
   file: 'lib/strategy.js',
   versions: ['>=1.0.0']
 }, Strategy => {
-  return shimmer.wrap(Strategy, function () {
+  return shimmer.wrapFunction(Strategy, Strategy => function () {
     const type = 'local'
 
     if (typeof arguments[0] === 'function') {