Check if prepareStackTrace has been wrapped before restore it (#4138)

* Check if prepareStackTrace has been wrapped before restore it * Use an IAST general flag
DataDog · Mar 7, 2024 · 17d540e · 17d540e
1 parent ff980c6
commit 17d540e
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 13 deletions.
diff --git a/packages/dd-trace/src/appsec/iast/index.js b/packages/dd-trace/src/appsec/iast/index.js
@@ -21,7 +21,11 @@ const requestStart = dc.channel('dd-trace:incomingHttpRequestStart')
 const requestClose = dc.channel('dd-trace:incomingHttpRequestEnd')
 const iastResponseEnd = dc.channel('datadog:iast:response-end')
 
+let isEnabled = false
+
 function enable (config, _tracer) {
+  if (isEnabled) return
+
   iastTelemetry.configure(config, config.iast?.telemetryVerbosity)
   enableAllAnalyzers(config)
   enableTaintTracking(config.iast, iastTelemetry.verbosity)
@@ -30,9 +34,15 @@ function enable (config, _tracer) {
   overheadController.configure(config.iast)
   overheadController.startGlobalContext()
   vulnerabilityReporter.start(config, _tracer)
+
+  isEnabled = true
 }
 
 function disable () {
+  if (!isEnabled) return
+
+  isEnabled = false
+
   iastTelemetry.stop()
   disableAllAnalyzers()
   disableTaintTracking()

diff --git a/packages/dd-trace/src/appsec/iast/taint-tracking/rewriter.js b/packages/dd-trace/src/appsec/iast/taint-tracking/rewriter.js
@@ -65,16 +65,18 @@ function getRewriter (telemetryVerbosity) {
   return rewriter
 }
 
-let originalPrepareStackTrace = Error.prepareStackTrace
+let originalPrepareStackTrace
+let actualPrepareStackTrace
 function getPrepareStackTraceAccessor () {
-  let actual = getPrepareStackTrace(originalPrepareStackTrace)
+  originalPrepareStackTrace = Error.prepareStackTrace
+  actualPrepareStackTrace = getPrepareStackTrace(originalPrepareStackTrace)
   return {
     configurable: true,
     get () {
-      return actual
+      return actualPrepareStackTrace
     },
     set (value) {
-      actual = getPrepareStackTrace(value)
+      actualPrepareStackTrace = getPrepareStackTrace(value)
       originalPrepareStackTrace = value
     }
   }
@@ -121,7 +123,18 @@ function enableRewriter (telemetryVerbosity) {
 
 function disableRewriter () {
   shimmer.unwrap(Module.prototype, '_compile')
-  Error.prepareStackTrace = originalPrepareStackTrace
+
+  if (!actualPrepareStackTrace) return
+
+  try {
+    delete Error.prepareStackTrace
+
+    Error.prepareStackTrace = originalPrepareStackTrace
+
+    actualPrepareStackTrace = undefined
+  } catch (e) {
+    iastLog.warn(e)
+  }
 }
 
 function getOriginalPathAndLineFromSourceMap ({ path, line, column }) {

diff --git a/packages/dd-trace/test/appsec/iast/index.spec.js b/packages/dd-trace/test/appsec/iast/index.spec.js
@@ -143,9 +143,24 @@ describe('IAST Index', () => {
       })
 
       it('should finish global context refresher on iast disabled', () => {
+        mockIast.enable(config)
+
         mockIast.disable()
         expect(mockOverheadController.finishGlobalContext).to.have.been.calledOnce
       })
+
+      it('should start global context only once when calling enable multiple times', () => {
+        mockIast.enable(config)
+        mockIast.enable(config)
+
+        expect(mockOverheadController.startGlobalContext).to.have.been.calledOnce
+      })
+
+      it('should not finish global context if not enabled before ', () => {
+        mockIast.disable(config)
+
+        expect(mockOverheadController.finishGlobalContext).to.have.been.not.called
+      })
     })
 
     describe('managing vulnerability reporter', () => {
@@ -156,6 +171,8 @@ describe('IAST Index', () => {
       })
 
       it('should stop vulnerability reporter on iast disabled', () => {
+        mockIast.enable(config)
+
         mockIast.disable()
         expect(mockVulnerabilityReporter.stop).to.have.been.calledOnce
       })

diff --git a/packages/dd-trace/test/appsec/iast/taint-tracking/rewriter.spec.js b/packages/dd-trace/test/appsec/iast/taint-tracking/rewriter.spec.js
@@ -13,12 +13,7 @@ describe('IAST Rewriter', () => {
   })
 
   describe('Enabling rewriter', () => {
-    let rewriter, iastTelemetry
-
-    const shimmer = {
-      wrap: sinon.spy(),
-      unwrap: sinon.spy()
-    }
+    let rewriter, iastTelemetry, shimmer
 
     class Rewriter {
       rewrite (content, filename) {
@@ -35,28 +30,76 @@ describe('IAST Rewriter', () => {
       iastTelemetry = {
         add: sinon.spy()
       }
+
+      shimmer = {
+        wrap: sinon.spy(),
+        unwrap: sinon.spy()
+      }
+
       rewriter = proxyquire('../../../../src/appsec/iast/taint-tracking/rewriter', {
-        '@datadog/native-iast-rewriter': { Rewriter, getPrepareStackTrace: function () {} },
+        '@datadog/native-iast-rewriter': {
+          Rewriter,
+          getPrepareStackTrace: function (fn) {
+            return function testWrappedPrepareStackTrace (_, callsites) {
+              return fn(_, callsites)
+            }
+          }
+        },
         '../../../../../datadog-shimmer': shimmer,
         '../../telemetry': iastTelemetry
       })
     })
 
     afterEach(() => {
-      sinon.restore()
+      sinon.reset()
     })
 
     it('Should wrap module compile method on taint tracking enable', () => {
       rewriter.enableRewriter()
       expect(shimmer.wrap).to.be.calledOnce
       expect(shimmer.wrap.getCall(0).args[1]).eq('_compile')
+
+      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')
     })
+
+    it('Should keep original prepareStackTrace fn when calling enable and then disable', () => {
+      const orig = Error.prepareStackTrace
+
+      rewriter.enableRewriter()
+
+      const testPrepareStackTrace = (_, callsites) => {
+        // do nothing
+      }
+      Error.prepareStackTrace = testPrepareStackTrace
+
+      rewriter.disableRewriter()
+
+      expect(Error.prepareStackTrace).to.be.eq(testPrepareStackTrace)
+
+      Error.prepareStackTrace = orig
+    })
+
+    it('Should keep original prepareStackTrace fn when calling disable only', () => {
+      const orig = Error.prepareStackTrace
+
+      const testPrepareStackTrace = (_, callsites) => {
+        // do nothing
+      }
+      Error.prepareStackTrace = testPrepareStackTrace
+
+      rewriter.disableRewriter()
+
+      expect(Error.prepareStackTrace).to.be.eq(testPrepareStackTrace)
+
+      Error.prepareStackTrace = orig
+    })
   })
 
   describe('getOriginalPathAndLineFromSourceMap', () => {