Get response headers on Fastify by the end of the request (#5831)

* set headers manually before calling writeHead

* use writeHead dc instead

* linter

* get right fastify span

* merge reply and request headers

* remove get-port

* fix tests

* avoid unecessary object creation

* avoid nullish storedResponseHeaders

* linter

Author: IlyasShabi

integration-tests/appsec/data-collection/fastify.js

@@ -0,0 +1,23 @@
+'use strict'
+
+const tracer = require('dd-trace')
+tracer.init({
+  flushInterval: 0
+})
+
+const fastify = require('fastify')
+const app = fastify()
+
+app.get('/', (request, reply) => {
+  reply.headers({
+    'content-type': 'text/plain',
+    'content-language': 'en-US',
+    'content-length': '3'
+  })
+  return 'end'
+})
+
+const port = process.env.APP_PORT ? Number(process.env.APP_PORT) : 0
+app.listen({ port }, (err) => {
+  process.send?.({ port: app.server.address().port })
+})

integration-tests/appsec/response-headers.spec.js

@@ -0,0 +1,88 @@
+'use strict'
+
+const { assert } = require('chai')
+const path = require('path')
+const Axios = require('axios')
+
+const {
+  createSandbox,
+  FakeAgent,
+  spawnProc
+} = require('../helpers')
+
+describe('Headers collection - Fastify', () => {
+  let axios, sandbox, cwd, appFile, agent, proc
+
+  before(async () => {
+    sandbox = await createSandbox(['fastify'])
+    cwd = sandbox.folder
+    appFile = path.join(cwd, 'appsec/data-collection/fastify.js')
+  })
+
+  after(async () => {
+    await sandbox.remove()
+  })
+
+  beforeEach(async () => {
+    agent = await new FakeAgent().start()
+
+    const env = {
+      DD_TRACE_AGENT_PORT: agent.port,
+      DD_APPSEC_ENABLED: true,
+      DD_APPSEC_RULES: path.join(cwd, 'appsec/data-collection/data-collection-rules.json')
+    }
+    proc = await spawnProc(appFile, { cwd, env, execArgv: [] })
+    axios = Axios.create({ baseURL: proc.url })
+  })
+
+  afterEach(async () => {
+    proc.kill()
+    await agent.stop()
+  })
+
+  it('should collect response headers with Fastify', async () => {
+    let fastifyRequestReceived = false
+    const requestHeaders = [
+      'user-agent',
+      'accept',
+      'host',
+      'accept-encoding'
+    ]
+
+    const responseHeaders = [
+      'content-type',
+      'content-language',
+      'content-length'
+    ]
+
+    await axios.get('/', {
+      headers: { 'User-Agent': 'Arachni/v1' }
+    })
+
+    await agent.assertMessageReceived(({ headers, payload }) => {
+      if (payload[0][0].name !== 'fastify.request') {
+        throw new Error('Not the span we are looking for')
+      }
+
+      fastifyRequestReceived = true
+      assert.equal(
+        Object.keys(payload[0][0].meta).filter(tagName => tagName.startsWith('http.request.headers.')).length,
+        requestHeaders.length
+      )
+      requestHeaders.forEach((headerName) => {
+        assert.property(payload[0][0].meta, `http.request.headers.${headerName}`)
+      })
+
+      // Response headers
+      assert.equal(
+        Object.keys(payload[0][0].meta).filter(tagName => tagName.startsWith('http.response.headers.')).length,
+        responseHeaders.length
+      )
+      responseHeaders.forEach((headerName) => {
+        assert.property(payload[0][0].meta, `http.response.headers.${headerName}`)
+      })
+    }, 30000, 10, true)
+
+    assert.isTrue(fastifyRequestReceived)
+  })
+})

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

@@ -37,6 +37,7 @@ const rasp = require('./rasp')
 const { isInServerlessEnvironment } = require('../serverless')
 
 const responseAnalyzedSet = new WeakSet()
+const storedResponseHeaders = new WeakMap()
 
 let isEnabled = false
 let config
@@ -187,7 +188,13 @@ function incomingHttpEndTranslator ({ req, res }) {
 
   waf.disposeContext(req)
 
-  Reporter.finishRequest(req, res)
+  const storedHeaders = storedResponseHeaders.get(req) || {}
+
+  Reporter.finishRequest(req, res, storedHeaders)
+
+  if (storedHeaders) {
+    storedResponseHeaders.delete(req)
+  }
 }
 
 function onPassportVerify ({ framework, login, user, success, abortController }) {
@@ -285,6 +292,10 @@ function onResponseBody ({ req, res, body }) {
 }
 
 function onResponseWriteHead ({ req, res, abortController, statusCode, responseHeaders }) {
+  if (Object.keys(responseHeaders).length) {
+    storedResponseHeaders.set(req, responseHeaders)
+  }
+
   // avoid "write after end" error
   if (isBlocked(res)) {
     abortController?.abort()

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

@@ -140,14 +140,16 @@ function filterExtendedHeaders (headers, excludedHeaderNames, tagPrefix, limit =
   return result
 }
 
-function getCollectedHeaders (req, res, shouldCollectEventHeaders) {
+function getCollectedHeaders (req, res, shouldCollectEventHeaders, storedResponseHeaders = {}) {
   // Mandatory
   const mandatoryCollectedHeaders = filterHeaders(req.headers, REQUEST_HEADERS_MAP)
 
   // Basic collection
   if (!shouldCollectEventHeaders) return mandatoryCollectedHeaders
 
-  const responseHeaders = res.getHeaders()
+  const responseHeaders = Object.keys(storedResponseHeaders).length === 0
+    ? res.getHeaders()
+    : { ...storedResponseHeaders, ...res.getHeaders() }
 
   const requestEventCollectedHeaders = filterHeaders(req.headers, EVENT_HEADERS_MAP)
   const responseEventCollectedHeaders = filterHeaders(responseHeaders, RESPONSE_HEADERS_MAP)
@@ -399,7 +401,7 @@ function reportDerivatives (derivatives) {
   rootSpan.addTags(tags)
 }
 
-function finishRequest (req, res) {
+function finishRequest (req, res, storedResponseHeaders) {
   const rootSpan = web.root(req)
   if (!rootSpan) return
 
@@ -453,7 +455,7 @@ function finishRequest (req, res) {
 
   const tags = rootSpan.context()._tags
 
-  const newTags = getCollectedHeaders(req, res, shouldCollectEventHeaders(tags))
+  const newTags = getCollectedHeaders(req, res, shouldCollectEventHeaders(tags), storedResponseHeaders)
 
   if (tags['appsec.event'] === 'true' && typeof req.route?.path === 'string') {
     newTags['http.endpoint'] = req.route.path

packages/dd-trace/test/appsec/index.spec.js

@@ -424,7 +424,52 @@ describe('AppSec Index', function () {
 
       expect(waf.run).to.have.not.been.called
 
-      expect(Reporter.finishRequest).to.have.been.calledOnceWithExactly(req, res)
+      expect(Reporter.finishRequest).to.have.been.calledOnceWithExactly(req, res, {})
+    })
+
+    it('should pass stored response headers to Reporter.finishRequest', () => {
+      const req = {
+        url: '/path',
+        headers: {
+          'user-agent': 'Arachni',
+          host: 'localhost'
+        },
+        method: 'POST',
+        socket: {
+          remoteAddress: '127.0.0.1',
+          remotePort: 8080
+        }
+      }
+      const res = {
+        getHeaders: () => ({
+          'content-type': 'application/json',
+          'content-length': 42
+        }),
+        statusCode: 200
+      }
+
+      const storedHeaders = {
+        'content-type': 'text/plain',
+        'content-language': 'en-US',
+        'content-length': '15'
+      }
+
+      web.patch(req)
+
+      sinon.stub(Reporter, 'finishRequest')
+      sinon.stub(waf, 'disposeContext')
+
+      responseWriteHead.publish({
+        req,
+        res,
+        abortController: { abort: sinon.stub() },
+        statusCode: 200,
+        responseHeaders: storedHeaders
+      })
+
+      AppSec.incomingHttpEndTranslator({ req, res })
+
+      expect(Reporter.finishRequest).to.have.been.calledOnceWithExactly(req, res, storedHeaders)
     })
 
     it('should not propagate incoming http end data with invalid framework properties', () => {
@@ -462,7 +507,7 @@ describe('AppSec Index', function () {
 
       expect(waf.run).to.have.not.been.called
 
-      expect(Reporter.finishRequest).to.have.been.calledOnceWithExactly(req, res)
+      expect(Reporter.finishRequest).to.have.been.calledOnceWithExactly(req, res, {})
     })
 
     it('should propagate incoming http end data with express', () => {
@@ -512,7 +557,7 @@ describe('AppSec Index', function () {
           'server.request.query': { b: '2' }
         }
       }, req)
-      expect(Reporter.finishRequest).to.have.been.calledOnceWithExactly(req, res)
+      expect(Reporter.finishRequest).to.have.been.calledOnceWithExactly(req, res, {})
     })
   })