feat(aap): In App WAF support for lambda (#7783)

Adds AppSec support for AWS lambda to dd-trace-js by introducing DC handlers that allow the datadog-lambda-js layer to delegate WAF execution to the tracer.

Author: CarlesDD

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

@@ -27,6 +27,8 @@ module.exports = {
   httpClientResponseFinish: dc.channel('apm:http:client:response:finish'),
   incomingHttpRequestEnd: dc.channel('dd-trace:incomingHttpRequestEnd'),
   incomingHttpRequestStart: dc.channel('dd-trace:incomingHttpRequestStart'),
+  lambdaStartInvocation: dc.channel('datadog:lambda:start-invocation'),
+  lambdaEndInvocation: dc.channel('datadog:lambda:end-invocation'),
   multerParser: dc.channel('datadog:multer:read:finish'),
   mysql2OuterQueryStart: dc.channel('datadog:mysql2:outerquery:start'),
   nextBodyParsed: dc.channel('apm:next:body-parsed'),

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

@@ -16,6 +16,8 @@ const {
   fastifyCookieParser,
   incomingHttpRequestStart,
   incomingHttpRequestEnd,
+  lambdaStartInvocation,
+  lambdaEndInvocation,
   passportVerify,
   passportUser,
   expressSession,
@@ -43,6 +45,7 @@ const { isBlocked, block, callBlockDelegation, setTemplates, getBlockingAction }
 const { getActiveRequest } = require('./store')
 const UserTracking = require('./user_tracking')
 const graphql = require('./graphql')
+const lambda = require('./lambda')
 const rasp = require('./rasp')
 
 const responseAnalyzedSet = new WeakSet()
@@ -99,11 +102,15 @@ function enable (_config) {
     stripeCheckoutSessionCreate.subscribe(onStripeCheckoutSessionCreate)
     stripePaymentIntentCreate.subscribe(onStripePaymentIntentCreate)
     stripeConstructEvent.subscribe(onStripeConstructEvent)
+    lambdaStartInvocation.subscribe(lambda.onLambdaStartInvocation)
+    lambdaEndInvocation.subscribe(lambda.onLambdaEndInvocation)
 
     isEnabled = true
     config = _config
   } catch (err) {
-    if (!IS_SERVERLESS) {
+    if (IS_SERVERLESS) {
+      log.debug('[ASM] Serverless mode: suppressing error log, calling disable()')
+    } else {
       log.error('[ASM] Unable to start AppSec', err)
     }
 
@@ -544,6 +551,8 @@ function disable () {
   if (stripeCheckoutSessionCreate.hasSubscribers) stripeCheckoutSessionCreate.unsubscribe(onStripeCheckoutSessionCreate)
   if (stripePaymentIntentCreate.hasSubscribers) stripePaymentIntentCreate.unsubscribe(onStripePaymentIntentCreate)
   if (stripeConstructEvent.hasSubscribers) stripeConstructEvent.unsubscribe(onStripeConstructEvent)
+  if (lambdaStartInvocation.hasSubscribers) lambdaStartInvocation.unsubscribe(lambda.onLambdaStartInvocation)
+  if (lambdaEndInvocation.hasSubscribers) lambdaEndInvocation.unsubscribe(lambda.onLambdaEndInvocation)
 }
 
 /**

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

@@ -0,0 +1,132 @@
+'use strict'
+
+const { HTTP_CLIENT_IP } = require('../../../../ext/tags')
+
+const log = require('../log')
+const addresses = require('./addresses')
+const Reporter = require('./reporter')
+const waf = require('./waf')
+
+// Tracks spans for which start-invocation has been processed, so that
+// end-invocation can gate correctly
+const activeInvocations = new WeakSet()
+
+/**
+ * Maps pre-extracted HTTP data from the Lambda event to WAF addresses,
+ * runs the WAF, and reports results on the span.
+ *
+ * @param {{ span: object, headers: Record<string, string>, method: string, path: string,
+ *           query: Record<string, string | string[]> | undefined, body: string | object | undefined,
+ *           isBase64Encoded: boolean, clientIp: string | undefined,
+ *           pathParams: Record<string, string> | undefined,
+ *           cookies: Record<string, string> | undefined,
+ *           route: string | undefined }} data
+ */
+function onLambdaStartInvocation (data) {
+  try {
+    const { span, headers, method, path, query, body, clientIp, pathParams, cookies } = data
+
+    if (!span) {
+      log.warn('[ASM] No span provided in Lambda start invocation')
+      return
+    }
+
+    activeInvocations.add(span)
+
+    span.setTag('_dd.appsec.enabled', 1)
+
+    const persistent = {}
+
+    if (path) {
+      persistent[addresses.HTTP_INCOMING_URL] = path
+    }
+
+    if (method) {
+      persistent[addresses.HTTP_INCOMING_METHOD] = method
+    }
+
+    if (headers) {
+      // Cookie header is already stripped by the Lambda layer's event-data-extractor
+      persistent[addresses.HTTP_INCOMING_HEADERS] = headers
+    }
+
+    if (clientIp) {
+      span.setTag(HTTP_CLIENT_IP, clientIp)
+      persistent[addresses.HTTP_CLIENT_IP] = clientIp
+    }
+
+    if (query) {
+      persistent[addresses.HTTP_INCOMING_QUERY] = query
+    }
+
+    if (body !== undefined && body !== null) {
+      persistent[addresses.HTTP_INCOMING_BODY] = body
+    }
+
+    if (pathParams) {
+      persistent[addresses.HTTP_INCOMING_PARAMS] = pathParams
+    }
+
+    if (cookies) {
+      persistent[addresses.HTTP_INCOMING_COOKIES] = cookies
+    }
+
+    waf.run({ persistent }, span, undefined, span)
+  } catch (err) {
+    log.error('[ASM] Error in Lambda start-invocation handler', err)
+  }
+}
+
+/**
+ * Maps response data to WAF addresses, runs a final WAF pass,
+ * disposes the WAF context, and finishes the request report.
+ *
+ * @param {{ span: object, statusCode: string | undefined,
+ *           responseHeaders: Record<string, string> | undefined }} data
+ */
+function onLambdaEndInvocation (data) {
+  try {
+    const { span, statusCode, responseHeaders } = data
+
+    if (!span) {
+      log.warn('[ASM] No span provided in Lambda end invocation')
+      return
+    }
+
+    if (!activeInvocations.has(span)) {
+      return
+    }
+
+    activeInvocations.delete(span)
+
+    let hasPersistentData = false
+    const persistent = {}
+
+    if (statusCode) {
+      persistent[addresses.HTTP_INCOMING_RESPONSE_CODE] = String(statusCode)
+      hasPersistentData = true
+    }
+
+    if (responseHeaders) {
+      const filteredHeaders = { ...responseHeaders }
+      delete filteredHeaders['set-cookie']
+      persistent[addresses.HTTP_INCOMING_RESPONSE_HEADERS] = filteredHeaders
+      hasPersistentData = true
+    }
+
+    if (hasPersistentData) {
+      waf.run({ persistent }, span, undefined, span)
+    }
+
+    waf.disposeContext(span)
+
+    Reporter.finishRequest(span, null, {}, undefined, span)
+  } catch (err) {
+    log.error('[ASM] Error in Lambda end-invocation handler', err)
+  }
+}
+
+module.exports = {
+  onLambdaStartInvocation,
+  onLambdaEndInvocation,
+}

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

@@ -318,13 +318,21 @@ function reportWafConfigUpdate (product, rcConfigId, diagnostics, wafVersion) {
   }
 }
 
-function reportMetrics (metrics, raspRule, req) {
+/**
+ * @param {object} metrics - WAF run metrics
+ * @param {string} [raspRule] - RASP rule identifier
+ * @param {object} [req] - Request key (plain object for lambda)
+ * @param {object} [rootSpan] - Root span (required for lambda)
+ */
+function reportMetrics (metrics, raspRule, req, rootSpan) {
   if (!req) {
     req = getActiveRequest()
   }
-  const rootSpan = req && web.root(req)
+  if (!rootSpan) {
+    rootSpan = req && web.root(req)
+  }
 
-  if (!rootSpan) return
+  if (!req || !rootSpan) return
 
   if (metrics.rulesVersion) {
     rootSpan.setTag('_dd.appsec.event_rules.version', metrics.rulesVersion)
@@ -353,13 +361,26 @@ function reportTruncationMetrics (rootSpan, metrics) {
   }
 }
 
-function reportAttack ({ events: attackData, actions }, req) {
+// NOTE: `req` in the WAF execution path may be any object, not necessarily
+// an HTTP IncomingMessage. In Lambda, it is a plain context key ({}) with no
+// HTTP properties. Always guard HTTP-specific property access
+// See tests in lambda.spec.js for enforcement.
+
+/**
+ * @param {{ events: Array, actions: object }} result - WAF result with attack data
+ * @param {object} [req] - Request key. May be an HTTP IncomingMessage or a plain object (Lambda)
+ * @param {object} [rootSpan] - Root span (required for lambda)
+ */
+function reportAttack ({ events: attackData, actions }, req, rootSpan) {
   if (!req) {
     req = getActiveRequest()
   }
 
-  const rootSpan = web.root(req)
-  if (!rootSpan) return
+  if (!rootSpan) {
+    rootSpan = web.root(req)
+  }
+
+  if (!req || !rootSpan) return
 
   const spanContext = rootSpan.context()
 
@@ -494,14 +515,21 @@ function isSchemaAttribute (attribute) {
   return attribute.startsWith('_dd.appsec.s.')
 }
 
-function reportAttributes (attributes, req) {
+/**
+ * @param {object} [attributes] - WAF result attributes
+ * @param {object} [req] - Request key (plain object for lambda)
+ * @param {object} [rootSpan] - Root span (required for lambda)
+ */
+function reportAttributes (attributes, req, rootSpan) {
   if (!attributes) return
 
   if (!req) {
     req = getActiveRequest()
   }
 
-  const rootSpan = web.root(req)
+  if (!rootSpan) {
+    rootSpan = web.root(req)
+  }
 
   if (!rootSpan) return
 
@@ -517,8 +545,17 @@ function reportAttributes (attributes, req) {
   rootSpan.addTags(tags)
 }
 
-function finishRequest (req, res, storedResponseHeaders, requestBody) {
-  const rootSpan = web.root(req)
+/**
+ * @param {object} [req] - Request key (null for Lambda)
+ * @param {object} [res] - Response object (null for lambda)
+ * @param {object} storedResponseHeaders
+ * @param {object} [requestBody]
+ * @param {object} [rootSpan] - Root span (required for lambda)
+ */
+function finishRequest (req, res, storedResponseHeaders, requestBody, rootSpan) {
+  if (!rootSpan) {
+    rootSpan = web.root(req)
+  }
   if (!rootSpan) return
 
   if (metricsQueue.size) {
@@ -569,6 +606,8 @@ function finishRequest (req, res, storedResponseHeaders, requestBody) {
     rootSpan.setTag('_dd.appsec.rasp.rule.eval', metrics.raspEvalCount)
   }
 
+  if (!req) return
+
   incrementWafRequestsMetric(req)
 
   const tags = rootSpan.context().getTags()

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

@@ -110,7 +110,19 @@ function removeConfig (configPath) {
   }
 }
 
-function run (data, req, raspRule) {
+// `req` in the WAF execution path may be any object, not necessarily
+// an HTTP IncomingMessage (it is a plain object for lambda).
+// Always guard HTTP-specific property access.
+// @see packages/dd-trace/test/appsec/lambda.spec.js
+
+/**
+ * @param {object} data - WAF address data ({ persistent, ephemeral })
+ * @param {object} req - Request key for WAF context lookup. May be an HTTP
+ *   IncomingMessage or a plain object (Lambda invocation key).
+ * @param {string} [raspRule] - RASP rule identifier
+ * @param {object} [rootSpan] - Root span to tag (required for Lambda)
+ */
+function run (data, req, raspRule, rootSpan) {
   if (!req) {
     req = getActiveRequest()
     if (!req) {
@@ -120,12 +132,12 @@ function run (data, req, raspRule) {
   }
 
   const wafContext = waf.wafManager.getWAFContext(req)
-  const result = wafContext.run(data, raspRule, req)
+  const result = wafContext.run(data, raspRule, req, rootSpan)
 
   if (result?.keep) {
     if (limiter.isAllowed()) {
-      const rootSpan = web.root(req)
-      keepTrace(rootSpan, ASM)
+      const span = rootSpan || web.root(req)
+      keepTrace(span, ASM)
     } else {
       updateRateLimitedMetric(req)
     }

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

@@ -36,7 +36,7 @@ class WAFContextWrapper {
     }
   }
 
-  run ({ persistent, ephemeral }, raspRule, req) {
+  run ({ persistent, ephemeral }, raspRule, req, rootSpan) {
     if (this.ddwafContext.disposed) {
       log.warn('[ASM] Calling run on a disposed context')
       if (raspRule) {
@@ -156,10 +156,10 @@ class WAFContextWrapper {
       metrics.wafTimeout = result.timeout
 
       if (ruleTriggered) {
-        Reporter.reportAttack(result, req)
+        Reporter.reportAttack(result, req, rootSpan)
       }
 
-      Reporter.reportAttributes(result.attributes, req)
+      Reporter.reportAttributes(result.attributes, req, rootSpan)
 
       return result
     } catch (err) {
@@ -171,7 +171,7 @@ class WAFContextWrapper {
         wafRunFinished.publish({ payload })
       }
 
-      Reporter.reportMetrics(metrics, raspRule, req)
+      Reporter.reportMetrics(metrics, raspRule, req, rootSpan)
     }
   }