Check if Error.prepareStackTrace has been wrapped by iast (#4158)

* Check Error.prepareStackTrace[kSymbolPrepareStackTrace] directly
DataDog · Mar 20, 2024 · 9796589 · 9796589
1 parent c4a6749
commit 9796589
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 20 deletions.
diff --git a/package.json b/package.json
@@ -70,7 +70,7 @@
   },
   "dependencies": {
     "@datadog/native-appsec": "7.1.0",
-    "@datadog/native-iast-rewriter": "2.2.3",
+    "@datadog/native-iast-rewriter": "2.3.0",
     "@datadog/native-iast-taint-tracking": "1.7.0",
     "@datadog/native-metrics": "^2.0.0",
     "@datadog/pprof": "5.1.0",

diff --git a/packages/dd-trace/src/appsec/iast/path-line.js b/packages/dd-trace/src/appsec/iast/path-line.js
@@ -29,13 +29,16 @@ function getCallSiteInfo () {
   const previousStackTraceLimit = Error.stackTraceLimit
   let callsiteList
   Error.stackTraceLimit = 100
-  Error.prepareStackTrace = function (_, callsites) {
-    callsiteList = callsites
+  try {
+    Error.prepareStackTrace = function (_, callsites) {
+      callsiteList = callsites
+    }
+    const e = new Error()
+    e.stack
+  } finally {
+    Error.prepareStackTrace = previousPrepareStackTrace
+    Error.stackTraceLimit = previousStackTraceLimit
   }
-  const e = new Error()
-  e.stack
-  Error.prepareStackTrace = previousPrepareStackTrace
-  Error.stackTraceLimit = previousStackTraceLimit
   return callsiteList
 }
 

diff --git a/packages/dd-trace/src/appsec/iast/taint-tracking/rewriter.js b/packages/dd-trace/src/appsec/iast/taint-tracking/rewriter.js
@@ -12,6 +12,7 @@ const dc = require('dc-polyfill')
 const hardcodedSecretCh = dc.channel('datadog:secrets:result')
 let rewriter
 let getPrepareStackTrace
+let kSymbolPrepareStackTrace
 
 let getRewriterOriginalPathAndLineFromSourceMap = function (path, line, column) {
   return { path, line, column }
@@ -44,6 +45,7 @@ function getRewriter (telemetryVerbosity) {
       const iastRewriter = require('@datadog/native-iast-rewriter')
       const Rewriter = iastRewriter.Rewriter
       getPrepareStackTrace = iastRewriter.getPrepareStackTrace
+      kSymbolPrepareStackTrace = iastRewriter.kSymbolPrepareStackTrace
 
       const chainSourceMap = isFlagPresent('--enable-source-maps')
       const getOriginalPathAndLineFromSourceMap = iastRewriter.getOriginalPathAndLineFromSourceMap
@@ -66,17 +68,16 @@ function getRewriter (telemetryVerbosity) {
 }
 
 let originalPrepareStackTrace
-let actualPrepareStackTrace
 function getPrepareStackTraceAccessor () {
   originalPrepareStackTrace = Error.prepareStackTrace
-  actualPrepareStackTrace = getPrepareStackTrace(originalPrepareStackTrace)
+  let actual = getPrepareStackTrace(originalPrepareStackTrace)
   return {
     configurable: true,
     get () {
-      return actualPrepareStackTrace
+      return actual
     },
     set (value) {
-      actualPrepareStackTrace = getPrepareStackTrace(value)
+      actual = getPrepareStackTrace(value)
       originalPrepareStackTrace = value
     }
   }
@@ -124,14 +125,12 @@ function enableRewriter (telemetryVerbosity) {
 function disableRewriter () {
   shimmer.unwrap(Module.prototype, '_compile')
 
-  if (!actualPrepareStackTrace) return
+  if (!Error.prepareStackTrace?.[kSymbolPrepareStackTrace]) return
 
   try {
     delete Error.prepareStackTrace
 
     Error.prepareStackTrace = originalPrepareStackTrace
-
-    actualPrepareStackTrace = undefined
   } catch (e) {
     iastLog.warn(e)
   }

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
@@ -36,14 +36,21 @@ describe('IAST Rewriter', () => {
         unwrap: sinon.spy()
       }
 
+      const kSymbolPrepareStackTrace = Symbol('kTestSymbolPrepareStackTrace')
+
       rewriter = proxyquire('../../../../src/appsec/iast/taint-tracking/rewriter', {
         '@datadog/native-iast-rewriter': {
           Rewriter,
           getPrepareStackTrace: function (fn) {
-            return function testWrappedPrepareStackTrace (_, callsites) {
+            const testWrap = function testWrappedPrepareStackTrace (_, callsites) {
               return fn(_, callsites)
             }
-          }
+            Object.defineProperty(testWrap, kSymbolPrepareStackTrace, {
+              value: true
+            })
+            return testWrap
+          },
+          kSymbolPrepareStackTrace
         },
         '../../../../../datadog-shimmer': shimmer,
         '../../telemetry': iastTelemetry
@@ -100,6 +107,26 @@ describe('IAST Rewriter', () => {
 
       Error.prepareStackTrace = orig
     })
+
+    it('Should keep original prepareStackTrace fn when calling disable if not marked with the Symbol', () => {
+      const orig = Error.prepareStackTrace
+
+      rewriter.enableRewriter()
+
+      // remove iast property to avoid wrapping the new testPrepareStackTrace fn
+      delete Error.prepareStackTrace
+
+      const testPrepareStackTrace = (_, callsites) => {
+        // do nothing
+      }
+      Error.prepareStackTrace = testPrepareStackTrace
+
+      rewriter.disableRewriter()
+
+      expect(Error.prepareStackTrace).to.be.eq(testPrepareStackTrace)
+
+      Error.prepareStackTrace = orig
+    })
   })
 
   describe('getOriginalPathAndLineFromSourceMap', () => {

diff --git a/yarn.lock b/yarn.lock
@@ -419,10 +419,10 @@
   dependencies:
     node-gyp-build "^3.9.0"
 
-"@datadog/native-iast-rewriter@2.2.3":
-  version "2.2.3"
-  resolved "https://registry.yarnpkg.com/@datadog/native-iast-rewriter/-/native-iast-rewriter-2.2.3.tgz#7d512abdb03dcc238825e8d6c90cebf782686db3"
-  integrity sha512-RCbflf8BJ++h99I7iA4NxTA1lx7YqB+sPQkJNSZKxXyEXtWl9J4XsDV9C/sB9iGbf1PVY77tFvoGm5/WpUV4IA==
+"@datadog/native-iast-rewriter@2.3.0":
+  version "2.3.0"
+  resolved "https://registry.yarnpkg.com/@datadog/native-iast-rewriter/-/native-iast-rewriter-2.3.0.tgz#67abddc698504ad76736c0d49681bcf9e330e1bd"
+  integrity sha512-78ivSaaSXOaHn3VumF9kcSI443nbPfVAWsnDTH9X1ZbqXjHpSlHHTZgK9z/TNbkvuJarS/X1GBioPMcgea1Ejg==
   dependencies:
     lru-cache "^7.14.0"
     node-gyp-build "^4.5.0"