refactor(remote-config): replace kPreUpdate with batch handler API (DEBUG-4402) (#7121)

- Remove the leaky kPreUpdate hook and introduce a first-class batch API:
  setBatchHandler()/removeBatchHandler()
- Introduce explicit product subscription via subscribeProducts()/
  unsubscribeProducts(), and have setProductHandler()/
  removeProductHandler() subscribe/unsubscribe for clarity
- Centralize ASM/WAF RC product names in appsec/rc-products.js
- Update AppSec WAF RC integration to use batched transaction (ack/error)
  instead of mutating configs
- Improve JSDoc/TS inference for batch transaction/descriptors and
  WAF manager typing

Author: watson

packages/dd-trace/src/appsec/rc-products.js

@@ -0,0 +1,10 @@
+'use strict'
+
+// Remote Config product names used by ASM/WAF.
+const ASM_WAF_PRODUCTS = ['ASM', 'ASM_DD', 'ASM_DATA']
+const ASM_WAF_PRODUCTS_SET = new Set(ASM_WAF_PRODUCTS)
+
+module.exports = {
+  ASM_WAF_PRODUCTS,
+  ASM_WAF_PRODUCTS_SET
+}

packages/dd-trace/src/appsec/rule_manager.js

@@ -1,75 +1,92 @@
 'use strict'
 
-const fs = require('fs')
+const { readFileSync } = require('node:fs')
+
 const waf = require('./waf')
 const { DIAGNOSTIC_KEYS } = require('./waf/diagnostics')
-const { ACKNOWLEDGED, ERROR } = require('../remote_config/apply_states')
-const Reporter = require('./reporter')
-
 const blocking = require('./blocking')
-
-const ASM_PRODUCTS = new Set(['ASM', 'ASM_DD', 'ASM_DATA'])
+const Reporter = require('./reporter')
+const { ASM_WAF_PRODUCTS_SET } = require('./rc-products')
 
 /*
   ASM Actions must be tracked in order to update the defaultBlockingActions in blocking. These actions are used
   by blockRequest method exposed in the user blocking SDK (see packages/dd-trace/src/appsec/sdk/user_blocking.js)
  */
 let appliedActions = new Map()
 
+/**
+ * @typedef {object} AsmConfigFile
+ * @property {Array<Record<string, unknown>>} [actions]
+ */
+
+/**
+ * @typedef {import('./waf').WAFConfig & { rules?: string }} AppSecConfig
+ */
+
+/**
+ * @param {AppSecConfig} config
+ */
 function loadRules (config) {
   const defaultRules = config.rules
-    ? JSON.parse(fs.readFileSync(config.rules))
+    ? JSON.parse(readFileSync(config.rules, 'utf8'))
     : require('./recommended.json')
 
   waf.init(defaultRules, config)
 
   blocking.setDefaultBlockingActionParameters(defaultRules?.actions)
 }
 
-function updateWafFromRC ({ toUnapply, toApply, toModify }) {
+/**
+ * Apply ASM remote-config updates to the WAF in a single batch.
+ *
+ * @param {import('../remote_config/manager').RcBatchUpdateTransaction} transaction
+ */
+function updateWafFromRC (transaction) {
+  const { toUnapply, toApply, toModify } = transaction
+
   const newActions = new SpyMap(appliedActions)
 
   let wafUpdated = false
   let wafUpdatedFailed = false
 
   for (const item of toUnapply) {
-    if (!ASM_PRODUCTS.has(item.product)) continue
+    if (!ASM_WAF_PRODUCTS_SET.has(item.product)) continue
 
     try {
       waf.removeConfig(item.path)
 
-      item.apply_state = ACKNOWLEDGED
+      transaction.ack(item.path)
       wafUpdated = true
 
       // ASM actions
       if (item.product === 'ASM') {
         newActions.delete(item.id)
       }
     } catch (e) {
-      item.apply_state = ERROR
-      item.apply_error = e.toString()
+      transaction.error(item.path, e)
       wafUpdatedFailed = true
     }
   }
 
   for (const item of [...toApply, ...toModify]) {
-    if (!ASM_PRODUCTS.has(item.product)) continue
+    if (!ASM_WAF_PRODUCTS_SET.has(item.product)) continue
 
     try {
       waf.updateConfig(item.product, item.id, item.path, item.file)
 
-      item.apply_state = ACKNOWLEDGED
+      transaction.ack(item.path)
       wafUpdated = true
 
       // ASM actions
-      if (item.product === 'ASM' && item.file?.actions?.length) {
-        newActions.set(item.id, item.file.actions)
+      if (item.product === 'ASM') {
+        const asmFile = /** @type {AsmConfigFile} */ (item.file)
+        if (asmFile?.actions?.length) {
+          newActions.set(item.id, asmFile.actions)
+        }
       }
     } catch (e) {
-      item.apply_state = ERROR
-      item.apply_error = e instanceof waf.WafUpdateError
-        ? JSON.stringify(extractErrors(e.diagnosticErrors))
-        : e.toString()
+      const error = e instanceof waf.WafUpdateError ? JSON.stringify(extractErrors(e.diagnosticErrors)) : e
+      transaction.error(item.path, error)
       wafUpdatedFailed = true
     }
   }

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

@@ -19,7 +19,17 @@ class WafUpdateError extends Error {
 
 let limiter = new Limiter(100)
 
+/** @typedef {import('./waf_manager')} WAFManager */
+
+/** @type {typeof import('./waf_manager') | null} */
+let WAFManager = null
+
+/**
+ * @typedef {import('./waf_manager').WAFManagerConfig & { rateLimit: number }} WAFConfig
+ */
+
 const waf = {
+  /** @type {WAFManager | null} */
   wafManager: null,
   init,
   destroy,
@@ -31,13 +41,17 @@ const waf = {
   WafUpdateError
 }
 
+/**
+ * @param {object} rules
+ * @param {WAFConfig} config
+ */
 function init (rules, config) {
   destroy()
 
   limiter = new Limiter(config.rateLimit)
 
-  // dirty require to make startup faster for serverless
-  const WAFManager = require('./waf_manager')
+  // Lazy loading improves the startup time
+  WAFManager = require('./waf_manager')
 
   waf.wafManager = new WAFManager(rules, config)
 
@@ -70,7 +84,7 @@ function updateConfig (product, configId, configPath, config) {
 
   try {
     if (product === 'ASM_DD') {
-      waf.wafManager.removeConfig(waf.wafManager.constructor.defaultWafConfigPath)
+      waf.wafManager.removeConfig((/** @type {NonNullable<typeof WAFManager>} */ (WAFManager)).defaultWafConfigPath)
     }
 
     const updateSucceeded = waf.wafManager.updateConfig(configPath, config)

packages/dd-trace/src/appsec/waf/waf_manager.js

@@ -6,9 +6,20 @@ const WAFContextWrapper = require('./waf_context_wrapper')
 
 const contexts = new WeakMap()
 
+/**
+ * @typedef {object} WAFManagerConfig
+ * @property {number} wafTimeout - Maximum time in microseconds for WAF execution
+ * @property {string} obfuscatorKeyRegex - Regex to redact sensitive data by key
+ * @property {string} obfuscatorValueRegex - Regex to redact sensitive data by value
+ */
+
 class WAFManager {
   static defaultWafConfigPath = 'datadog/00/ASM_DD/default/config'
 
+  /**
+   * @param {object} rules
+   * @param {WAFManagerConfig} config
+   */
   constructor (rules, config) {
     this.config = config
     this.wafTimeout = config.wafTimeout

packages/dd-trace/src/remote_config/index.js

@@ -88,6 +88,7 @@ function enableOrDisableAppsec (action, rcConfig, config, appsec) {
 function enableWafUpdate (appsecConfig) {
   if (rc && appsecConfig && !appsecConfig.rules) {
     // dirty require to make startup faster for serverless
+    const { ASM_WAF_PRODUCTS } = require('../appsec/rc-products')
     const RuleManager = require('../appsec/rule_manager')
 
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_IP_BLOCKING, true)
@@ -119,17 +120,14 @@ function enableWafUpdate (appsecConfig) {
       rc.updateCapabilities(RemoteConfigCapabilities.ASM_RASP_CMDI, true)
     }
 
-    // TODO: delete noop handlers and kPreUpdate and replace with batched handlers
-    rc.setProductHandler('ASM_DATA', noop)
-    rc.setProductHandler('ASM_DD', noop)
-    rc.setProductHandler('ASM', noop)
-
-    rc.on(RemoteConfigManager.kPreUpdate, RuleManager.updateWafFromRC)
+    rc.subscribeProducts(...ASM_WAF_PRODUCTS)
+    rc.setBatchHandler(ASM_WAF_PRODUCTS, RuleManager.updateWafFromRC)
   }
 }
 
 function disableWafUpdate () {
   if (rc) {
+    const { ASM_WAF_PRODUCTS } = require('../appsec/rc-products')
     const RuleManager = require('../appsec/rule_manager')
 
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_IP_BLOCKING, false)
@@ -158,16 +156,11 @@ function disableWafUpdate () {
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_RASP_SHI, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_RASP_CMDI, false)
 
-    rc.removeProductHandler('ASM_DATA')
-    rc.removeProductHandler('ASM_DD')
-    rc.removeProductHandler('ASM')
-
-    rc.off(RemoteConfigManager.kPreUpdate, RuleManager.updateWafFromRC)
+    rc.unsubscribeProducts(...ASM_WAF_PRODUCTS)
+    rc.removeBatchHandler(RuleManager.updateWafFromRC)
   }
 }
 
-function noop () {}
-
 module.exports = {
   enable,
   enableWafUpdate,