Add support for setting the maximum number of concurrent transactions (#2926)

* Add support for setting the maximum number of concurrent transactions in iast operations.

* Update taint tracking module to version 1.3.1.

Author: hoolioh

package.json

@@ -64,7 +64,7 @@
   "dependencies": {
     "@datadog/native-appsec": "2.0.0",
     "@datadog/native-iast-rewriter": "2.0.1",
-    "@datadog/native-iast-taint-tracking": "1.1.1",
+    "@datadog/native-iast-taint-tracking": "1.3.1",
     "@datadog/native-metrics": "^1.5.0",
     "@datadog/pprof": "^2.1.0",
     "@datadog/sketches-js": "^2.1.0",

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

@@ -6,6 +6,7 @@ const overheadController = require('./overhead-controller')
 const dc = require('diagnostics_channel')
 const iastContextFunctions = require('./iast-context')
 const { enableTaintTracking, disableTaintTracking, createTransaction, removeTransaction } = require('./taint-tracking')
+
 const telemetryLogs = require('./telemetry/logs')
 const IAST_ENABLED_TAG_KEY = '_dd.iast.enabled'
 
@@ -16,7 +17,7 @@ const requestClose = dc.channel('dd-trace:incomingHttpRequestEnd')
 
 function enable (config, _tracer) {
   enableAllAnalyzers()
-  enableTaintTracking()
+  enableTaintTracking(config.iast)
   requestStart.subscribe(onIncomingHttpRequestStart)
   requestClose.subscribe(onIncomingHttpRequestEnd)
   overheadController.configure(config.iast)

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

@@ -1,20 +1,27 @@
 'use strict'
 
 const { enableRewriter, disableRewriter } = require('./rewriter')
-const { createTransaction, removeTransaction, enableTaintOperations, disableTaintOperations } = require('./operations')
+const { createTransaction,
+  removeTransaction,
+  setMaxTransactions,
+  enableTaintOperations,
+  disableTaintOperations } = require('./operations')
+
 const taintTrackingPlugin = require('./plugin')
 
 module.exports = {
-  enableTaintTracking () {
+  enableTaintTracking (config) {
     enableRewriter()
     enableTaintOperations()
     taintTrackingPlugin.enable()
+    setMaxTransactions(config.maxConcurrentRequests)
   },
   disableTaintTracking () {
     disableRewriter()
     disableTaintOperations()
     taintTrackingPlugin.disable()
   },
+  setMaxTransactions: setMaxTransactions,
   createTransaction: createTransaction,
   removeTransaction: removeTransaction
 }

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

@@ -87,6 +87,14 @@ function disableTaintOperations () {
   global._ddiast = TaintTrackingDummy
 }
 
+function setMaxTransactions (transactions) {
+  if (!transactions) {
+    return
+  }
+
+  TaintedUtils.setMaxTransactions(transactions)
+}
+
 module.exports = {
   createTransaction,
   removeTransaction,
@@ -96,5 +104,6 @@ module.exports = {
   getRanges,
   enableTaintOperations,
   disableTaintOperations,
+  setMaxTransactions,
   IAST_TRANSACTION_ID
 }

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

@@ -5,14 +5,21 @@ const proxyquire = require('proxyquire')
 
 describe('IAST TaintTracking', () => {
   let taintTracking
+  const config = {
+    iast: {
+      maxConcurrentRequests: 2
+    }
+  }
+
   const rewriter = {
     enableRewriter: sinon.spy(),
     disableRewriter: sinon.spy()
   }
 
   const taintTrackingOperations = {
     enableTaintOperations: sinon.spy(),
-    disableTaintOperations: sinon.spy()
+    disableTaintOperations: sinon.spy(),
+    setMaxTransactions: sinon.spy()
   }
 
   const taintTrackingPlugin = {
@@ -31,10 +38,12 @@ describe('IAST TaintTracking', () => {
   afterEach(sinon.restore)
 
   it('Should enable rewriter, taint tracking operations and plugin', () => {
-    taintTracking.enableTaintTracking()
+    taintTracking.enableTaintTracking(config.iast)
     expect(rewriter.enableRewriter).to.be.calledOnce
     expect(taintTrackingOperations.enableTaintOperations).to.be.calledOnce
     expect(taintTrackingPlugin.enable).to.be.calledOnce
+    expect(taintTrackingOperations.setMaxTransactions)
+      .to.have.been.calledOnceWithExactly(config.iast.maxConcurrentRequests)
   })
 
   it('Should disable both rewriter, taint tracking operations, plugin', () => {

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

@@ -11,6 +11,7 @@ describe('IAST TaintTracking Operations', () => {
   const taintedUtils = {
     createTransaction: id => id,
     removeTransaction: id => id,
+    setMaxTransactions: () => {},
     newTaintedString: id => id,
     isTainted: id => id,
     getRanges: id => id,
@@ -101,6 +102,20 @@ describe('IAST TaintTracking Operations', () => {
     })
   })
 
+  describe('SetMaxTransactions', () => {
+    it('Given a number of concurrent transactions should call setMaxTransactions', () => {
+      const transactions = 3
+
+      taintTrackingOperations.setMaxTransactions(transactions)
+      expect(taintedUtils.setMaxTransactions).to.have.been.calledOnceWithExactly(transactions)
+    })
+
+    it('Given undefined as a number of concurrent transactions should not call setMaxTransactions', () => {
+      taintTrackingOperations.setMaxTransactions()
+      expect(taintedUtils.setMaxTransactions).not.to.have.been.called
+    })
+  })
+
   describe('enableTaintTracking', () => {
     beforeEach(() => {
       iastContextFunctions.saveIastContext(

yarn.lock

@@ -267,10 +267,10 @@
   dependencies:
     node-gyp-build "^4.5.0"
 
-"@datadog/native-iast-taint-tracking@1.1.1":
-  version "1.1.1"
-  resolved "https://registry.yarnpkg.com/@datadog/native-iast-taint-tracking/-/native-iast-taint-tracking-1.1.1.tgz#cbeace022b6c1f3a0a40dc0000cc40079c6d4895"
-  integrity sha512-VkESVYpVlLHqw38UHqqEYsJaJTp3+JpKIJhfB9nlQO13dYBc3Sgq/QJZNdPViU73SVsCJtuw4D0SXRyjTXP1IA==
+"@datadog/native-iast-taint-tracking@1.3.1":
+  version "1.3.1"
+  resolved "https://registry.yarnpkg.com/@datadog/native-iast-taint-tracking/-/native-iast-taint-tracking-1.3.1.tgz#49b3befc3049370f4034babcf57c3d67e9f4d56b"
+  integrity sha512-KWKmK4/GANisxqVZ1TtGlBIOw2RIXdUO0r7361QJHiBVUxwNKmKNVDVuCTKGpRRH/0GZcxY0yVgl38ee/6HM3A==
   dependencies:
     node-gyp-build "^3.9.0"