fix shimmer slow shim by removing support for unwrap (#5557)

Author: rochdev

integration-tests/package-guardrails/index.js

@@ -11,5 +11,5 @@ try {
 } catch (e) {
   const fastify = require('fastify')
 
-  console.log(fastify.toString().startsWith('function shim'))
+  console.log(fastify.toString().startsWith('function fastifyWithTrace'))
 }

packages/datadog-shimmer/src/shimmer.js

@@ -2,9 +2,6 @@
 
 const log = require('../../dd-trace/src/log')
 
-// Use a weak map to avoid polluting the wrapped function/method.
-const unwrappers = new WeakMap()
-
 function copyProperties (original, wrapped) {
   // TODO getPrototypeOf is not fast. Should we instead do this in specific
   // instrumentations where needed?
@@ -27,21 +24,13 @@ function copyProperties (original, wrapped) {
 function wrapFunction (original, wrapper) {
   if (typeof original === 'function') assertNotClass(original)
 
-  let delegate = safeMode
+  const wrapped = safeMode
     ? safeWrapper(original, wrapper)
     : wrapper(original)
 
-  const shim = function shim () {
-    return delegate.apply(this, arguments)
-  }
-
-  unwrappers.set(shim, () => {
-    delegate = original
-  })
-
-  if (typeof original === 'function') copyProperties(original, shim)
+  if (typeof original === 'function') copyProperties(original, wrapped)
 
-  return shim
+  return wrapped
 }
 
 const wrapFn = function (original, delegate) {
@@ -99,8 +88,6 @@ function wrapMethod (target, name, wrapper, noAssert) {
   if (typeof original === 'function') copyProperties(original, wrapped)
 
   if (descriptor) {
-    unwrappers.set(wrapped, () => Object.defineProperty(target, name, descriptor))
-
     if (descriptor.get || descriptor.set) {
       attributes.get = () => wrapped
     } else {
@@ -114,7 +101,6 @@ function wrapMethod (target, name, wrapper, noAssert) {
       })
     }
   } else { // no descriptor means original was on the prototype
-    unwrappers.set(wrapped, () => delete target[name])
     attributes.value = wrapped
     attributes.writable = true
   }
@@ -218,18 +204,6 @@ function wrap (target, name, wrapper) {
     : wrapMethod(target, name, wrapper)
 }
 
-function unwrap (target, name) {
-  if (!target) return target // no target to unwrap
-
-  const unwrapper = unwrappers.get(name ? target[name] : target)
-
-  if (!unwrapper) return target // target is already unwrapped or isn't wrapped
-
-  unwrapper()
-
-  return target
-}
-
 function massWrap (targets, names, wrapper) {
   targets = toArray(targets)
   names = toArray(names)
@@ -241,17 +215,6 @@ function massWrap (targets, names, wrapper) {
   }
 }
 
-function massUnwrap (targets, names) {
-  targets = toArray(targets)
-  names = toArray(names)
-
-  for (const target of targets) {
-    for (const name of names) {
-      unwrap(target, name)
-    }
-  }
-}
-
 function toArray (maybeArray) {
   return Array.isArray(maybeArray) ? maybeArray : [maybeArray]
 }
@@ -294,7 +257,5 @@ module.exports = {
   wrap,
   wrapFunction,
   massWrap,
-  unwrap,
-  massUnwrap,
   setSafe
 }

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

@@ -166,28 +166,6 @@ describe('shimmer', () => {
       expect(count).to.have.property('enumerable', false)
     })
 
-    it('should unwrap a method', () => {
-      const count = inc => inc
-      const obj = { count }
-
-      shimmer.wrap(obj, 'count', count => inc => count(inc) + 1)
-      shimmer.unwrap(obj, 'count')
-
-      expect(obj.count(1)).to.equal(1)
-    })
-
-    it('should unwrap a method from the prototype', () => {
-      const count = inc => inc
-      const obj = {}
-
-      Object.setPrototypeOf(obj, { count })
-
-      shimmer.wrap(obj, 'count', count => inc => count(inc) + 1)
-      shimmer.unwrap(obj, 'count')
-
-      expect(obj).to.not.have.ownProperty('count')
-    })
-
     it('should validate that there is a target object', () => {
       expect(() => shimmer.wrap()).to.throw()
     })
@@ -212,22 +190,6 @@ describe('shimmer', () => {
       expect(() => shimmer.wrap({ a: () => {} }, 'a', 'notafunction')).to.throw()
     })
 
-    it('should not throw when unwrapping without a target', () => {
-      expect(() => shimmer.unwrap(null, 'a')).to.not.throw()
-    })
-
-    it('should not throw when unwrapping without a method', () => {
-      expect(() => shimmer.unwrap({}, 'a')).to.not.throw()
-    })
-
-    it('should not throw when unwrapping an invalid type', () => {
-      expect(() => shimmer.unwrap({ a: 'b' }, 'a')).to.not.throw()
-    })
-
-    it('should not throw when unwrapping a method that was not wrapped', () => {
-      expect(() => shimmer.unwrap({ a: () => {} }, 'a')).to.not.throw()
-    })
-
     describe('safe mode', () => {
       let obj
 
@@ -539,34 +501,6 @@ describe('shimmer', () => {
       expect(Object.getOwnPropertyNames(wrapped)).to.not.include('test')
     })
 
-    it('should unwrap a function', () => {
-      const count = inc => inc
-
-      const wrapped = shimmer.wrapFunction(count, count => inc => count(inc) + 1)
-
-      shimmer.unwrap(wrapped)
-
-      expect(wrapped(1)).to.equal(1)
-    })
-
-    it('should unwrap a constructor', () => {
-      const Counter = function (start) {
-        this.value = start
-      }
-
-      const WrappedCounter = shimmer.wrapFunction(Counter, Counter => function (...args) {
-        Counter.apply(this, arguments)
-        this.value++
-      })
-
-      shimmer.unwrap(WrappedCounter)
-
-      const counter = new WrappedCounter(1)
-
-      expect(counter.value).to.equal(1)
-      expect(counter).to.be.an.instanceof(Counter)
-    })
-
     it('should mass wrap methods on objects', () => {
       const foo = {
         a: () => 'original',
@@ -586,44 +520,12 @@ describe('shimmer', () => {
       expect(bar.b()).to.equal('wrapped')
     })
 
-    it('should mass wrap methods on objects', () => {
-      const foo = {
-        a: () => 'original',
-        b: () => 'original'
-      }
-
-      const bar = {
-        a: () => 'original',
-        b: () => 'original'
-      }
-
-      shimmer.massWrap([foo, bar], ['a', 'b'], () => () => 'wrapped')
-      shimmer.massUnwrap([foo, bar], ['a', 'b'])
-
-      expect(foo.a()).to.equal('original')
-      expect(foo.b()).to.equal('original')
-      expect(bar.a()).to.equal('original')
-      expect(bar.b()).to.equal('original')
-    })
-
     it('should validate that the function wrapper exists', () => {
       expect(() => shimmer.wrap(() => {})).to.throw()
     })
 
     it('should validate that the function wrapper is a function', () => {
       expect(() => shimmer.wrap(() => {}, 'a')).to.throw()
     })
-
-    it('should never throw when unwrapping', () => {
-      expect(() => shimmer.unwrap(() => {})).to.not.throw()
-    })
-
-    it('should not throw when unwrapping an invalid type', () => {
-      expect(() => shimmer.unwrap('foo')).to.not.throw()
-    })
-
-    it('should not throw when unwrapping a function that was not wrapped', () => {
-      expect(() => shimmer.unwrap(() => {})).to.not.throw()
-    })
   })
 })

packages/dd-trace/src/appsec/iast/taint-tracking/rewriter.js

@@ -15,6 +15,7 @@ const { LOG_MESSAGE, REWRITTEN_MESSAGE } = require('./constants')
 
 const hardcodedSecretCh = dc.channel('datadog:secrets:result')
 let rewriter
+let unwrapCompile = () => {}
 let getPrepareStackTrace, cacheRewrittenSourceMap
 let kSymbolPrepareStackTrace
 let esmRewriterEnabled = false
@@ -90,7 +91,8 @@ function getPrepareStackTraceAccessor () {
 
 function getCompileMethodFn (compileMethod) {
   const rewriteFn = getRewriteFunction(rewriter)
-  return function (content, filename) {
+
+  let delegate = function (content, filename) {
     try {
       if (isPrivateModule(filename) && isNotLibraryFile(filename)) {
         const rewritten = rewriteFn(content, filename)
@@ -108,6 +110,16 @@ function getCompileMethodFn (compileMethod) {
     }
     return compileMethod.apply(this, [content, filename])
   }
+
+  const shim = function () {
+    return delegate.apply(this, arguments)
+  }
+
+  unwrapCompile = function () {
+    delegate = compileMethod
+  }
+
+  return shim
 }
 
 function esmRewritePostProcess (rewritten, filename) {
@@ -198,7 +210,7 @@ function enableEsmRewriter (telemetryVerbosity) {
 }
 
 function disableRewriter () {
-  shimmer.unwrap(Module.prototype, '_compile')
+  unwrapCompile()
 
   if (!Error.prepareStackTrace?.[kSymbolPrepareStackTrace]) return
 

packages/dd-trace/test/appsec/iast/taint-tracking/rewriter.spec.js

@@ -107,12 +107,8 @@ describe('IAST Rewriter', () => {
       rewriter.disableRewriter()
     })
 
-    it('Should unwrap module compile method on taint tracking disable', () => {
-      rewriter.disableRewriter()
-
-      expect(shimmer.unwrap).to.be.calledOnce
-      expect(shimmer.unwrap.getCall(0).args[1]).eq('_compile')
-    })
+    // TODO: This cannot be tested with mocking.
+    it('Should unwrap module compile method on taint tracking disable') // eslint-disable-line mocha/no-pending-tests
 
     it('Should keep original prepareStackTrace fn when calling enable and then disable', () => {
       const orig = Error.prepareStackTrace