From 57b92a932bb8a0333b39e4ce68ad24825706e13c Mon Sep 17 00:00:00 2001
From: Igor Unanua <igor.unanua@datadoghq.com>
Date: Wed, 28 Aug 2024 11:01:55 +0200
Subject: [PATCH] Enable AppsecFsPlugin for iast

---
 packages/dd-trace/src/appsec/fs-plugin.js     | 31 ++++++++++++++-----
 packages/dd-trace/src/appsec/iast/index.js    |  3 ++
 packages/dd-trace/src/appsec/rasp.js          |  6 ++--
 .../analyzers/path-traversal-analyzer.spec.js |  4 +--
 .../test/appsec/response_blocking.spec.js     |  2 +-
 5 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/packages/dd-trace/src/appsec/fs-plugin.js b/packages/dd-trace/src/appsec/fs-plugin.js
index ba5bfe5529c..450cfdadd28 100644
--- a/packages/dd-trace/src/appsec/fs-plugin.js
+++ b/packages/dd-trace/src/appsec/fs-plugin.js
@@ -2,6 +2,12 @@
 
 const Plugin = require('../plugins/plugin')
 const { storage } = require('../../../datadog-core')
+const log = require('../log')
+
+const enabledFor = {
+  rasp: false,
+  iast: false
+}
 
 let fsPlugin
 
@@ -51,20 +57,29 @@ class AppsecFsPlugin extends Plugin {
   }
 }
 
-function enable () {
-  if (fsPlugin) return
+function enable (mod) {
+  if (!mod || enabledFor[mod]) return
 
-  fsPlugin = new AppsecFsPlugin()
-  fsPlugin.enable()
+  enabledFor[mod] = true
+
+  if (!fsPlugin) {
+    fsPlugin = new AppsecFsPlugin()
+    fsPlugin.enable()
+  }
+
+  log.info(`Enabled AppsecFsPlugin for ${mod}`)
 }
 
-function disable () {
-  if (!fsPlugin) return
+function disable (mod) {
+  if (!mod || !enabledFor[mod]) return
+
+  enabledFor[mod] = false
 
-  // FIXME: AppsecFsPlugin could be used by appsec and iast
-  fsPlugin.disable()
+  fsPlugin?.disable()
 
   fsPlugin = undefined
+
+  log.info(`Disabled AppsecFsPlugin for ${mod}`)
 }
 
 module.exports = {
diff --git a/packages/dd-trace/src/appsec/iast/index.js b/packages/dd-trace/src/appsec/iast/index.js
index 0facaa39a2a..41a88bb3c54 100644
--- a/packages/dd-trace/src/appsec/iast/index.js
+++ b/packages/dd-trace/src/appsec/iast/index.js
@@ -14,6 +14,7 @@ const {
 } = require('./taint-tracking')
 const { IAST_ENABLED_TAG_KEY } = require('./tags')
 const iastTelemetry = require('./telemetry')
+const { enable: enableFsPlugin, disable: disableFsPlugin } = require('../fs-plugin')
 
 // TODO Change to `apm:http:server:request:[start|close]` when the subscription
 //  order of the callbacks can be enforce
@@ -27,6 +28,7 @@ function enable (config, _tracer) {
   if (isEnabled) return
 
   iastTelemetry.configure(config, config.iast?.telemetryVerbosity)
+  enableFsPlugin('iast')
   enableAllAnalyzers(config)
   enableTaintTracking(config.iast, iastTelemetry.verbosity)
   requestStart.subscribe(onIncomingHttpRequestStart)
@@ -44,6 +46,7 @@ function disable () {
   isEnabled = false
 
   iastTelemetry.stop()
+  disableFsPlugin('iast')
   disableAllAnalyzers()
   disableTaintTracking()
   overheadController.finishGlobalContext()
diff --git a/packages/dd-trace/src/appsec/rasp.js b/packages/dd-trace/src/appsec/rasp.js
index 63b3bf841c3..d8283a7f72d 100644
--- a/packages/dd-trace/src/appsec/rasp.js
+++ b/packages/dd-trace/src/appsec/rasp.js
@@ -8,7 +8,7 @@ const { reportStackTrace } = require('./stack_trace')
 const waf = require('./waf')
 const { getBlockingAction, block } = require('./blocking')
 const log = require('../log')
-const { enable: fsPluginEnable, disable: fsPluginDisable } = require('./fs-plugin')
+const { enable: enableFsPlugin, disable: disableFsPlugin } = require('./fs-plugin')
 
 const RULE_TYPES = {
   SSRF: 'ssrf',
@@ -104,7 +104,7 @@ function handleUncaughtExceptionMonitor (err) {
 function enable (_config) {
   config = _config
 
-  fsPluginEnable()
+  enableFsPlugin('rasp')
 
   httpClientRequestStart.subscribe(analyzeSsrf)
   fsOperationStart.subscribe(analyzeLfi)
@@ -121,7 +121,7 @@ function disable () {
   if (httpClientRequestStart.hasSubscribers) httpClientRequestStart.unsubscribe(analyzeSsrf)
   if (fsOperationStart.hasSubscribers) fsOperationStart.unsubscribe(analyzeLfi)
 
-  fsPluginDisable()
+  disableFsPlugin('rasp')
 
   process.off('uncaughtExceptionMonitor', handleUncaughtExceptionMonitor)
 }
diff --git a/packages/dd-trace/test/appsec/iast/analyzers/path-traversal-analyzer.spec.js b/packages/dd-trace/test/appsec/iast/analyzers/path-traversal-analyzer.spec.js
index 42b01301d85..387d33dc633 100644
--- a/packages/dd-trace/test/appsec/iast/analyzers/path-traversal-analyzer.spec.js
+++ b/packages/dd-trace/test/appsec/iast/analyzers/path-traversal-analyzer.spec.js
@@ -163,9 +163,9 @@ describe('path-traversal-analyzer', () => {
 prepareTestServerForIast('integration test', (testThatRequestHasVulnerability, testThatRequestHasNoVulnerability) => {
   function runFsMethodTest (description, vulnerableIndex, fn, ...args) {
     describe(description, () => {
-      before(() => enableFsPlugin())
+      before(() => enableFsPlugin('iast'))
 
-      after(() => disableFsPlugin())
+      after(() => disableFsPlugin('iast'))
 
       describe('vulnerable', () => {
         testThatRequestHasVulnerability(function () {
diff --git a/packages/dd-trace/test/appsec/response_blocking.spec.js b/packages/dd-trace/test/appsec/response_blocking.spec.js
index dd641e8227c..699c051dd8a 100644
--- a/packages/dd-trace/test/appsec/response_blocking.spec.js
+++ b/packages/dd-trace/test/appsec/response_blocking.spec.js
@@ -57,7 +57,7 @@ describe('HTTP Response Blocking', () => {
       }
     }))
 
-    disableFsPlugin()
+    disableFsPlugin('rasp')
   })
 
   beforeEach(() => {