support iast with fastify

packages/datadog-instrumentations/src/fastify.js

@@ -11,6 +11,7 @@ const queryParamsReadCh = channel('datadog:fastify:query-params:finish')
 const cookieParserReadCh = channel('datadog:fastify-cookie:read:finish')
 const responsePayloadReadCh = channel('datadog:fastify:response:finish')
 const pathParamsReadCh = channel('datadog:fastify:path-params:finish')
+const finishSetHeaderCh = channel('datadog:fastify:set-header:finish')
 
 const parsingResources = new WeakMap()
 const cookiesPublished = new WeakSet()
@@ -275,3 +276,23 @@ addHook({ name: 'fastify', versions: ['2'] }, fastify => {
 addHook({ name: 'fastify', versions: ['1'] }, fastify => {
   return shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, false))
 })
+
+function wrapReplyHeader (Reply) {
+  shimmer.wrap(Reply.prototype, 'header', header => function (key, value) {
+    const result = header.apply(this, arguments)
+
+    if (finishSetHeaderCh.hasSubscribers && key && value) {
+      finishSetHeaderCh.publish({ name: key, value, res: getRes(this) })
+    }
+
+    return result
+  })
+
+  return Reply
+}
+
+addHook({ name: 'fastify', file: 'lib/reply.js', versions: ['1'] }, wrapReplyHeader)
+
+addHook({ name: 'fastify', file: 'lib/reply.js', versions: ['2'] }, wrapReplyHeader)
+
+addHook({ name: 'fastify', file: 'lib/reply.js', versions: ['>=3'] }, wrapReplyHeader)

packages/dd-trace/src/appsec/iast/analyzers/cookie-analyzer.js

@@ -3,7 +3,14 @@
 const Analyzer = require('./vulnerability-analyzer')
 const { getNodeModulesPaths } = require('../path-line')
 
-const EXCLUDED_PATHS = getNodeModulesPaths('express/lib/response.js')
+const EXCLUDED_PATHS = [
+  // Express
+  getNodeModulesPaths('express/lib/response.js'),
+  // Fastify
+  getNodeModulesPaths('fastify/lib/reply.js'),
+  getNodeModulesPaths('fastify/lib/hooks.js'),
+  getNodeModulesPaths('@fastify/cookie/plugin.js')
+]
 
 class CookieAnalyzer extends Analyzer {
   constructor (type, propertyToBeSafe) {

packages/dd-trace/src/appsec/iast/analyzers/hsts-header-missing-analyzer.js

@@ -10,8 +10,8 @@ class HstsHeaderMissingAnalyzer extends MissingHeaderAnalyzer {
     super(HSTS_HEADER_MISSING, HSTS_HEADER_NAME)
   }
 
-  _isVulnerableFromRequestAndResponse (req, res) {
-    const headerValues = this._getHeaderValues(res, HSTS_HEADER_NAME)
+  _isVulnerableFromRequestAndResponse (req, res, storedHeaders) {
+    const headerValues = this._getHeaderValues(res, storedHeaders, HSTS_HEADER_NAME)
     return this._isHttpsProtocol(req) && (
       headerValues.length === 0 ||
       headerValues.some(headerValue => !this._isHeaderValid(headerValue))

packages/dd-trace/src/appsec/iast/analyzers/missing-header-analyzer.js

@@ -28,8 +28,9 @@ class MissingHeaderAnalyzer extends Analyzer {
     }, (data) => this.analyze(data))
   }
 
-  _getHeaderValues (res, headerName) {
-    const headerValue = res.getHeader(headerName)
+  _getHeaderValues (res, storedHeaders, headerName) {
+    headerName = headerName.toLowerCase()
+    const headerValue = res.getHeader(headerName) || storedHeaders[headerName]
     if (Array.isArray(headerValue)) {
       return headerValue
     }
@@ -46,8 +47,8 @@ class MissingHeaderAnalyzer extends Analyzer {
     return `${type}:${this.config.tracerConfig.service}`
   }
 
-  _getEvidence ({ res }) {
-    const headerValues = this._getHeaderValues(res, this.headerName)
+  _getEvidence ({ res, storedHeaders }) {
+    const headerValues = this._getHeaderValues(res, storedHeaders, this.headerName)
     let value
     if (headerValues.length === 1) {
       value = headerValues[0]
@@ -57,19 +58,19 @@ class MissingHeaderAnalyzer extends Analyzer {
     return { value }
   }
 
-  _isVulnerable ({ req, res }, context) {
-    if (!IGNORED_RESPONSE_STATUS_LIST.has(res.statusCode) && this._isResponseHtml(res)) {
-      return this._isVulnerableFromRequestAndResponse(req, res)
+  _isVulnerable ({ req, res, storedHeaders }, context) {
+    if (!IGNORED_RESPONSE_STATUS_LIST.has(res.statusCode) && this._isResponseHtml(res, storedHeaders)) {
+      return this._isVulnerableFromRequestAndResponse(req, res, storedHeaders)
     }
     return false
   }
 
-  _isVulnerableFromRequestAndResponse (req, res) {
+  _isVulnerableFromRequestAndResponse (req, res, storedHeaders) {
     return false
   }
 
-  _isResponseHtml (res) {
-    const contentTypes = this._getHeaderValues(res, 'content-type')
+  _isResponseHtml (res, storedHeaders) {
+    const contentTypes = this._getHeaderValues(res, storedHeaders, 'content-type')
     return contentTypes.some(contentType => {
       return contentType && HTML_CONTENT_TYPES.some(htmlContentType => {
         return htmlContentType === contentType || contentType.startsWith(htmlContentType + ';')

packages/dd-trace/src/appsec/iast/analyzers/set-cookies-header-interceptor.js

@@ -7,25 +7,32 @@ class SetCookiesHeaderInterceptor extends Plugin {
   constructor () {
     super()
     this.cookiesInRequest = new WeakMap()
-    this.addSub('datadog:http:server:response:set-header:finish', ({ name, value, res }) => {
-      if (name.toLowerCase() === 'set-cookie') {
-        let allCookies = value
-        if (typeof value === 'string') {
-          allCookies = [value]
-        }
-        const alreadyCheckedCookies = this._getAlreadyCheckedCookiesInResponse(res)
-
-        let location
-        allCookies.forEach(cookieString => {
-          if (!alreadyCheckedCookies.includes(cookieString)) {
-            alreadyCheckedCookies.push(cookieString)
-            const parsedCookie = this._parseCookie(cookieString, location)
-            setCookieChannel.publish(parsedCookie)
-            location = parsedCookie.location
-          }
-        })
+
+    this.addSub('datadog:http:server:response:set-header:finish',
+      ({ name, value, res }) => this._handleCookies(name, value, res))
+
+    this.addSub('datadog:fastify:set-header:finish',
+      ({ name, value, res }) => this._handleCookies(name, value, res))
+  }
+
+  _handleCookies (name, value, res) {
+    if (name.toLowerCase() === 'set-cookie') {
+      let allCookies = value
+      if (typeof value === 'string') {
+        allCookies = [value]
       }
-    })
+      const alreadyCheckedCookies = this._getAlreadyCheckedCookiesInResponse(res)
+
+      let location
+      allCookies.forEach(cookieString => {
+        if (!alreadyCheckedCookies.includes(cookieString)) {
+          alreadyCheckedCookies.push(cookieString)
+          const parsedCookie = this._parseCookie(cookieString, location)
+          setCookieChannel.publish(parsedCookie)
+          location = parsedCookie.location
+        }
+      })
+    }
   }
 
   _parseCookie (cookieString, location) {

packages/dd-trace/src/appsec/iast/analyzers/unvalidated-redirect-analyzer.js

@@ -9,7 +9,10 @@ const {
   HTTP_REQUEST_PARAMETER
 } = require('../taint-tracking/source-types')
 
-const EXCLUDED_PATHS = getNodeModulesPaths('express/lib/response.js')
+const EXCLUDED_PATHS = [
+  getNodeModulesPaths('express/lib/response.js'),
+  getNodeModulesPaths('fastify/lib/reply.js'),
+]
 
 const VULNERABLE_SOURCE_TYPES = new Set([
   HTTP_REQUEST_BODY,
@@ -23,6 +26,7 @@ class UnvalidatedRedirectAnalyzer extends InjectionAnalyzer {
 
   onConfigure () {
     this.addSub('datadog:http:server:response:set-header:finish', ({ name, value }) => this.analyze(name, value))
+    this.addSub('datadog:fastify:set-header:finish', ({ name, value }) => this.analyze(name, value))
   }
 
   analyze (name, value) {

packages/dd-trace/src/appsec/iast/analyzers/xcontenttype-header-missing-analyzer.js

@@ -10,8 +10,8 @@ class XcontenttypeHeaderMissingAnalyzer extends MissingHeaderAnalyzer {
     super(XCONTENTTYPE_HEADER_MISSING, XCONTENTTYPEOPTIONS_HEADER_NAME)
   }
 
-  _isVulnerableFromRequestAndResponse (req, res) {
-    const headerValues = this._getHeaderValues(res, XCONTENTTYPEOPTIONS_HEADER_NAME)
+  _isVulnerableFromRequestAndResponse (req, res, storedHeaders) {
+    const headerValues = this._getHeaderValues(res, storedHeaders, XCONTENTTYPEOPTIONS_HEADER_NAME)
     return headerValues.length === 0 || headerValues.some(headerValue => headerValue.trim().toLowerCase() !== 'nosniff')
   }
 }

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

@@ -47,6 +47,10 @@ class IastPluginSubscription {
   }
 
   matchesModuleInstrumented (name) {
+    // Remove node: prefix if present
+    if (name.startsWith('node:')) {
+      name = name.slice(5)
+    }
     // https module is a special case because it's events are published as http
     name = name === 'https' ? 'http' : name
     return this.moduleName === name

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

@@ -18,11 +18,12 @@ const { IAST_ENABLED_TAG_KEY } = require('./tags')
 const iastTelemetry = require('./telemetry')
 const { enable: enableFsPlugin, disable: disableFsPlugin, IAST_MODULE } = require('../rasp/fs-plugin')
 const securityControls = require('./security-controls')
+const { incomingHttpRequestStart, incomingHttpRequestEnd, responseWriteHead } = require('../channels')
+
+const collectedResponseHeaders = new WeakMap()
 
 // TODO Change to `apm:http:server:request:[start|close]` when the subscription
 //  order of the callbacks can be enforce
-const requestStart = dc.channel('dd-trace:incomingHttpRequestStart')
-const requestClose = dc.channel('dd-trace:incomingHttpRequestEnd')
 const iastResponseEnd = dc.channel('datadog:iast:response-end')
 let isEnabled = false
 
@@ -33,8 +34,9 @@ function enable (config, _tracer) {
   enableFsPlugin(IAST_MODULE)
   enableAllAnalyzers(config)
   enableTaintTracking(config.iast, iastTelemetry.verbosity)
-  requestStart.subscribe(onIncomingHttpRequestStart)
-  requestClose.subscribe(onIncomingHttpRequestEnd)
+  incomingHttpRequestStart.subscribe(onIncomingHttpRequestStart)
+  incomingHttpRequestEnd.subscribe(onIncomingHttpRequestEnd)
+  responseWriteHead.subscribe(onResponseWriteHeadCollect)
   overheadController.configure(config.iast)
   overheadController.startGlobalContext()
   securityControls.configure(config.iast)
@@ -53,8 +55,9 @@ function disable () {
   disableAllAnalyzers()
   disableTaintTracking()
   overheadController.finishGlobalContext()
-  if (requestStart.hasSubscribers) requestStart.unsubscribe(onIncomingHttpRequestStart)
-  if (requestClose.hasSubscribers) requestClose.unsubscribe(onIncomingHttpRequestEnd)
+  if (incomingHttpRequestStart.hasSubscribers) incomingHttpRequestStart.unsubscribe(onIncomingHttpRequestStart)
+  if (incomingHttpRequestEnd.hasSubscribers) incomingHttpRequestEnd.unsubscribe(onIncomingHttpRequestEnd)
+  if (responseWriteHead.hasSubscribers) responseWriteHead.unsubscribe(onResponseWriteHeadCollect)
   vulnerabilityReporter.stop()
 }
 
@@ -89,7 +92,13 @@ function onIncomingHttpRequestEnd (data) {
     const topContext = web.getContext(data.req)
     const iastContext = iastContextFunctions.getIastContext(store, topContext)
     if (iastContext?.rootSpan) {
-      iastResponseEnd.publish(data)
+      const storedHeaders = collectedResponseHeaders.get(data.res) || {}
+
+      iastResponseEnd.publish({ ...data, storedHeaders })
+
+      if (storedHeaders) {
+        collectedResponseHeaders.delete(data.res)
+      }
 
       const vulnerabilities = iastContext.vulnerabilities
       const rootSpan = iastContext.rootSpan
@@ -105,4 +114,13 @@ function onIncomingHttpRequestEnd (data) {
   }
 }
 
+// Response headers are collected here because they are not available in the onIncomingHttpRequestEnd when using Fastify
+function onResponseWriteHeadCollect ({ res, responseHeaders = {} }) {
+  if (!res) return
+
+  if (Object.keys(responseHeaders).length) {
+    collectedResponseHeaders.set(res, responseHeaders)
+  }
+}
+
 module.exports = { enable, disable, onIncomingHttpRequestEnd, onIncomingHttpRequestStart }