Extract rewriter telemetry to another module

Author: iunanua

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

@@ -6,7 +6,7 @@ const taintTrackingPlugin = require('./plugin')
 
 module.exports = {
   enableTaintTracking (telemetryVerbosity) {
-    enableRewriter()
+    enableRewriter(telemetryVerbosity)
     enableTaintOperations(telemetryVerbosity)
     taintTrackingPlugin.enable()
   },

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

@@ -0,0 +1,50 @@
+'use strict'
+
+const telemetry = require('../../telemetry')
+const { Verbosity } = require('../../telemetry/verbosity')
+const { IastMetric, PropagationType } = require('../iast-metric')
+
+const telemetryRewriter = {
+  off (content, filename, rewriter) {
+    return rewriter.rewrite(content, filename)
+  },
+
+  information (content, filename, rewriter) {
+    const response = this.off(content, filename, rewriter)
+
+    const metrics = response.metrics
+    if (metrics && metrics.instrumentedPropagation) {
+      telemetry.add(IastMetric.INSTRUMENTED_PROPAGATION, metrics.instrumentedPropagation, PropagationType.STRING)
+    }
+
+    return response
+  },
+
+  debug (content, filename, rewriter) {
+    const start = process.hrtime.bigint()
+    const response = this.information(content, filename, rewriter)
+
+    // TODO: propagationDebug!
+    const metrics = response.metrics
+    if (metrics && metrics.propagationDebug) {
+      // debug metrics are using logs telemetry API instead metrics telemetry API
+    }
+
+    const rewriteTime = parseInt(process.hrtime.bigint() - start) * 1e-6
+    telemetry.add(IastMetric.INSTRUMENTATION_TIME, rewriteTime)
+    return response
+  }
+}
+
+function getRewriteFunction (rewriter) {
+  switch (telemetry.verbosity) {
+    case Verbosity.OFF:
+      return (content, filename) => telemetryRewriter.off(content, filename, rewriter)
+    case Verbosity.DEBUG:
+      return (content, filename) => telemetryRewriter.debug(content, filename, rewriter)
+    default:
+      return (content, filename) => telemetryRewriter.information(content, filename, rewriter)
+  }
+}
+
+module.exports = { getRewriteFunction }

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

@@ -5,19 +5,18 @@ const shimmer = require('../../../../../datadog-shimmer')
 const iastLog = require('../iast-log')
 const { isPrivateModule, isNotLibraryFile } = require('./filter')
 const { csiMethods } = require('./csi-methods')
-const telemetry = require('../../telemetry')
-const { Verbosity } = require('../../telemetry/verbosity')
-const { IastMetric, PropagationType } = require('../iast-metric')
+const { getName } = require('../../telemetry/verbosity')
+const { getRewriteFunction } = require('./rewriter-telemetry')
 
 let rewriter
 let getPrepareStackTrace
-function getRewriter () {
+function getRewriter (telemetryVerbosity) {
   if (!rewriter) {
     try {
       const iastRewriter = require('@datadog/native-iast-rewriter')
       const Rewriter = iastRewriter.Rewriter
       getPrepareStackTrace = iastRewriter.getPrepareStackTrace
-      rewriter = new Rewriter({ csiMethods, telemetryVerbosity: telemetry.getVerbosityName() })
+      rewriter = new Rewriter({ csiMethods, telemetryVerbosity: getName(telemetryVerbosity) })
     } catch (e) {
       iastLog.error('Unable to initialize TaintTracking Rewriter')
         .errorAndPublish(e)
@@ -41,51 +40,8 @@ function getPrepareStackTraceAccessor () {
   }
 }
 
-const telemetryOffRewrite = function (content, filename) {
-  if (isPrivateModule(filename) && isNotLibraryFile(filename)) {
-    return rewriter.rewrite(content, filename)
-  }
-}
-
-const telemetryInformationRewrite = function (content, filename) {
-  const response = telemetryOffRewrite(content, filename)
-
-  const metrics = response.metrics
-  if (metrics && metrics.instrumentedPropagation) {
-    telemetry.add(IastMetric.INSTRUMENTED_PROPAGATION, metrics.instrumentedPropagation, PropagationType.STRING)
-  }
-
-  return response
-}
-
-const telemetryDebugRewrite = function (content, filename) {
-  const start = process.hrtime.bigint()
-  const response = telemetryInformationRewrite(content, filename)
-
-  // TODO: propagationDebug!
-  const metrics = response.metrics
-  if (metrics && metrics.propagationDebug) {
-    // debug metrics are using logs telemetry API instead metrics telemetry API
-  }
-
-  const rewriteTime = parseInt(process.hrtime.bigint() - start) * 1e-6
-  telemetry.add(IastMetric.INSTRUMENTATION_TIME, rewriteTime)
-  return response
-}
-
-function getRewriteFunction () {
-  switch (telemetry.verbosity) {
-    case Verbosity.OFF:
-      return telemetryOffRewrite
-    case Verbosity.DEBUG:
-      return telemetryDebugRewrite
-    default:
-      return telemetryInformationRewrite
-  }
-}
-
 function getCompileMethodFn (compileMethod) {
-  const rewriteFn = getRewriteFunction()
+  const rewriteFn = getRewriteFunction(rewriter)
   return function (content, filename) {
     try {
       if (isPrivateModule(filename) && isNotLibraryFile(filename)) {
@@ -102,9 +58,9 @@ function getCompileMethodFn (compileMethod) {
   }
 }
 
-function enableRewriter () {
+function enableRewriter (telemetryVerbosity) {
   try {
-    const rewriter = getRewriter()
+    const rewriter = getRewriter(telemetryVerbosity)
     if (rewriter) {
       const pstDescriptor = Object.getOwnPropertyDescriptor(global.Error, 'prepareStackTrace')
       if (!pstDescriptor || pstDescriptor.configurable) {
@@ -124,5 +80,5 @@ function disableRewriter () {
 }
 
 module.exports = {
-  enableRewriter, disableRewriter, getRewriteFunction
+  enableRewriter, disableRewriter
 }

packages/dd-trace/src/appsec/telemetry/index.js

@@ -2,7 +2,7 @@
 
 const logs = require('./api/logs-plugin')
 const metrics = require('./api/metrics-plugin')
-const { Verbosity, isDebugAllowed, isInfoAllowed, parseVerbosity, getName } = require('./verbosity')
+const { Verbosity, isDebugAllowed, isInfoAllowed, parseVerbosity } = require('./verbosity')
 const {
   add,
   drain,
@@ -46,10 +46,6 @@ class Telemetry {
     return this.isEnabled() && isInfoAllowed(this.verbosity)
   }
 
-  getVerbosityName () {
-    return getName(this.verbosity)
-  }
-
   increase (metric, tag, context) {
     add(metric, 1, tag, context)
   }

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

@@ -0,0 +1,77 @@
+'use strict'
+const { expect } = require('chai')
+const proxyquire = require('proxyquire')
+
+const { IastMetric, PropagationType } = require('../../../../src/appsec/iast/iast-metric')
+const { Verbosity } = require('../../../../src/appsec/telemetry/verbosity')
+
+describe('rewriter telemetry', () => {
+  let telemetry, rewriter, getRewriteFunction
+
+  beforeEach(() => {
+    telemetry = {
+      add: sinon.spy()
+    }
+    const rewriterTelemetry = proxyquire('../../../../src/appsec/iast/taint-tracking/rewriter-telemetry', {
+      '../../telemetry': telemetry
+    })
+    getRewriteFunction = rewriterTelemetry.getRewriteFunction
+    rewriter = {
+      rewrite: (content) => {
+        return {
+          content: content + 'rewritten',
+          metrics: {
+            instrumentedPropagation: 2
+          }
+        }
+      }
+    }
+  })
+
+  it('should not increase any metrics with OFF verbosity', () => {
+    telemetry.verbosity = Verbosity.OFF
+
+    const rewriteFn = getRewriteFunction(rewriter)
+    rewriteFn('const a = b + c', 'test.js')
+
+    expect(telemetry.add).to.not.be.called
+  })
+
+  it('should increase information metrics with MANDATORY verbosity', () => {
+    telemetry.verbosity = Verbosity.MANDATORY
+
+    const rewriteFn = getRewriteFunction(rewriter)
+    const result = rewriteFn('const a = b + c', 'test.js')
+
+    expect(telemetry.add).to.be.calledOnceWith(IastMetric.INSTRUMENTED_PROPAGATION,
+      result.metrics.instrumentedPropagation,
+      PropagationType.STRING)
+  })
+
+  it('should increase information metrics with INFORMATION verbosity', () => {
+    telemetry.verbosity = Verbosity.INFORMATION
+
+    const rewriteFn = getRewriteFunction(rewriter)
+    const result = rewriteFn('const a = b + c', 'test.js')
+
+    expect(telemetry.add).to.be.calledOnceWith(IastMetric.INSTRUMENTED_PROPAGATION,
+      result.metrics.instrumentedPropagation,
+      PropagationType.STRING)
+  })
+
+  it('should increase debug metrics with DEBUG verbosity', () => {
+    telemetry.verbosity = Verbosity.DEBUG
+
+    const rewriteFn = getRewriteFunction(rewriter)
+    const result = rewriteFn('const a = b + c', 'test.js')
+
+    expect(telemetry.add).to.be.calledTwice
+    const instrumentedPropagation = telemetry.add.getCall(0).args
+    expect(instrumentedPropagation).to.be.deep.eq([IastMetric.INSTRUMENTED_PROPAGATION,
+      result.metrics.instrumentedPropagation,
+      PropagationType.STRING])
+
+    const instrumentationTime = telemetry.add.getCall(1).args
+    expect(instrumentationTime[0]).to.be.eq(IastMetric.INSTRUMENTATION_TIME)
+  })
+})

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

@@ -2,8 +2,6 @@
 
 const { expect } = require('chai')
 const proxyquire = require('proxyquire')
-const { IastMetric, PropagationType } = require('../../../../src/appsec/iast/iast-metric')
-const { Verbosity } = require('../../../../src/appsec/telemetry/verbosity')
 
 describe('IAST Rewriter', () => {
   it('Addon should return a rewritter instance', () => {
@@ -35,8 +33,7 @@ describe('IAST Rewriter', () => {
 
     beforeEach(() => {
       telemetry = {
-        add: sinon.spy(),
-        getVerbosityName: () => 'OFF'
+        add: sinon.spy()
       }
       rewriter = proxyquire('../../../../src/appsec/iast/taint-tracking/rewriter', {
         '@datadog/native-iast-rewriter': { Rewriter, getPrepareStackTrace: function () {} },
@@ -60,58 +57,5 @@ describe('IAST Rewriter', () => {
       expect(shimmer.unwrap).to.be.calledOnce
       expect(shimmer.unwrap.getCall(0).args[1]).eq('_compile')
     })
-
-    describe('telemetry', () => {
-      it('should not increase any metrics with OFF verbosity', () => {
-        telemetry.verbosity = Verbosity.OFF
-
-        rewriter.enableRewriter()
-        const rewriteFn = rewriter.getRewriteFunction()
-        rewriteFn('const a = b + c', 'test.js')
-
-        expect(telemetry.add).to.not.be.called
-      })
-
-      it('should increase information metrics with MANDATORY verbosity', () => {
-        telemetry.verbosity = Verbosity.MANDATORY
-
-        rewriter.enableRewriter()
-        const rewriteFn = rewriter.getRewriteFunction()
-        const result = rewriteFn('const a = b + c', 'test.js')
-
-        expect(telemetry.add).to.be.calledOnceWith(IastMetric.INSTRUMENTED_PROPAGATION,
-          result.metrics.instrumentedPropagation,
-          PropagationType.STRING)
-      })
-
-      it('should increase information metrics with INFORMATION verbosity', () => {
-        telemetry.verbosity = Verbosity.INFORMATION
-
-        rewriter.enableRewriter()
-        const rewriteFn = rewriter.getRewriteFunction()
-        const result = rewriteFn('const a = b + c', 'test.js')
-
-        expect(telemetry.add).to.be.calledOnceWith(IastMetric.INSTRUMENTED_PROPAGATION,
-          result.metrics.instrumentedPropagation,
-          PropagationType.STRING)
-      })
-
-      it('should increase debug metrics with DEBUG verbosity', () => {
-        telemetry.verbosity = Verbosity.DEBUG
-
-        rewriter.enableRewriter()
-        const rewriteFn = rewriter.getRewriteFunction()
-        const result = rewriteFn('const a = b + c', 'test.js')
-
-        expect(telemetry.add).to.be.calledTwice
-        const instrumentedPropagation = telemetry.add.getCall(0).args
-        expect(instrumentedPropagation).to.be.deep.eq([IastMetric.INSTRUMENTED_PROPAGATION,
-          result.metrics.instrumentedPropagation,
-          PropagationType.STRING])
-
-        const instrumentationTime = telemetry.add.getCall(1).args
-        expect(instrumentationTime[0]).to.be.eq(IastMetric.INSTRUMENTATION_TIME)
-      })
-    })
   })
 })