Support schema extraction in fastify response objects (#5894)

Author: IlyasShabi

packages/datadog-instrumentations/src/fastify.js

@@ -8,6 +8,7 @@ const handleChannel = channel('apm:fastify:request:handle')
 const routeAddedChannel = channel('apm:fastify:route:added')
 const bodyParserReadCh = channel('datadog:fastify:body-parser:finish')
 const queryParamsReadCh = channel('datadog:fastify:query-params:finish')
+const responsePayloadReadCh = channel('datadog:fastify:response:finish')
 const pathParamsReadCh = channel('datadog:fastify:path-params:finish')
 
 const parsingResources = new WeakMap()
@@ -167,9 +168,12 @@ function preParsing (request, reply, payload, done) {
 }
 
 function wrapSend (send, req) {
-  return function sendWithTrace (error) {
-    if (error instanceof Error) {
-      errorChannel.publish({ req, error })
+  return function sendWithTrace (payload) {
+    if (payload instanceof Error) {
+      errorChannel.publish({ req, error: payload })
+    } else if (canPublishResponsePayload(payload)) {
+      const res = getRes(this)
+      responsePayloadReadCh.publish({ req, res, body: payload })
     }
 
     return send.apply(this, arguments)
@@ -200,6 +204,18 @@ function onRoute (routeOptions) {
   routeAddedChannel.publish({ routeOptions, onRoute })
 }
 
+// send() payload types: https://fastify.dev/docs/latest/Reference/Reply/#senddata
+function canPublishResponsePayload (payload) {
+  return responsePayloadReadCh.hasSubscribers &&
+    payload &&
+    typeof payload === 'object' &&
+    typeof payload.pipe !== 'function' && // Node streams
+    typeof payload.body?.pipe !== 'function' && // Response with body stream
+    !Buffer.isBuffer(payload) && // Buffer
+    !(payload instanceof ArrayBuffer) && // ArrayBuffer
+    !ArrayBuffer.isView(payload) // TypedArray
+}
+
 addHook({ name: 'fastify', versions: ['>=3'] }, fastify => {
   const wrapped = shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, true))
 

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

@@ -13,6 +13,7 @@ module.exports = {
   expressProcessParams: dc.channel('datadog:express:process_params:start'),
   expressSession: dc.channel('datadog:express-session:middleware:finish'),
   fastifyBodyParser: dc.channel('datadog:fastify:body-parser:finish'),
+  fastifyResponseChannel: dc.channel('datadog:fastify:response:finish'),
   fastifyQueryParams: dc.channel('datadog:fastify:query-params:finish'),
   fastifyPathParams: dc.channel('datadog:fastify:path-params:finish'),
   fsOperationStart: dc.channel('apm:fs:operation:start'),

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

@@ -22,6 +22,7 @@ const {
   responseWriteHead,
   responseSetHeader,
   routerParam,
+  fastifyResponseChannel,
   fastifyPathParams
 } = require('./channels')
 const waf = require('./waf')
@@ -85,6 +86,7 @@ function enable (_config) {
     fastifyPathParams.subscribe(onRequestProcessParams)
     routerParam.subscribe(onRequestProcessParams)
     responseBody.subscribe(onResponseBody)
+    fastifyResponseChannel.subscribe(onResponseBody)
     responseWriteHead.subscribe(onResponseWriteHead)
     responseSetHeader.subscribe(onResponseSetHeader)
 
@@ -378,6 +380,7 @@ function disable () {
   if (fastifyPathParams.hasSubscribers) fastifyPathParams.unsubscribe(onRequestProcessParams)
   if (routerParam.hasSubscribers) routerParam.unsubscribe(onRequestProcessParams)
   if (responseBody.hasSubscribers) responseBody.unsubscribe(onResponseBody)
+  if (fastifyResponseChannel.hasSubscribers) fastifyResponseChannel.unsubscribe(onResponseBody)
   if (responseWriteHead.hasSubscribers) responseWriteHead.unsubscribe(onResponseWriteHead)
   if (responseSetHeader.hasSubscribers) responseSetHeader.unsubscribe(onResponseSetHeader)
 }

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

@@ -5,6 +5,8 @@ const { assert } = require('chai')
 const getPort = require('get-port')
 const path = require('path')
 const semver = require('semver')
+const zlib = require('zlib')
+const fs = require('node:fs')
 const agent = require('../plugins/agent')
 const appsec = require('../../src/appsec')
 const Config = require('../../src/config')
@@ -167,11 +169,12 @@ withVersions('fastify', 'fastify', version => {
         assert.deepEqual(e.response.data, JSON.parse(json))
         sinon.assert.notCalled(requestBody)
 
-        await agent.assertSomeTraces((traces) => {
-          const span = traces[0][0]
-          assert.equal(span.metrics['_dd.appsec.truncated.string_length'], 5000)
-          assert.equal(span.metrics['_dd.appsec.truncated.container_size'], 300)
-          assert.equal(span.metrics['_dd.appsec.truncated.container_depth'], 20)
+        await agent.assertFirstTraceSpan({
+          metrics: {
+            '_dd.appsec.truncated.string_length': 5000,
+            '_dd.appsec.truncated.container_size': 300,
+            '_dd.appsec.truncated.container_depth': 20
+          }
         })
       }
     })
@@ -446,6 +449,159 @@ withVersions('fastify', 'fastify', version => {
   })
 })
 
+describe('Api Security - Fastify', () => {
+  withVersions('fastify', 'fastify', version => {
+    let config, server, axios
+
+    before(() => {
+      return agent.load(['fastify', 'http'], { client: false })
+    })
+
+    before((done) => {
+      const fastify = require(`../../../../versions/fastify@${version}`).get()
+
+      const app = fastify()
+
+      app.post('/send', (request, reply) => {
+        reply.send({ sendResKey: 'sendResValue' })
+      })
+
+      app.post('/return', async (request, reply) => {
+        return { returnResKey: 'returnResValue' }
+      })
+
+      app.get('/', (request, reply) => {
+        reply.send('DONE')
+      })
+
+      app.get('/buffer', (request, reply) => {
+        reply.send(Buffer.from('DONE'))
+      })
+
+      app.get('/stream', (request, reply) => {
+        const stream = fs.createReadStream(__filename)
+        reply.header('Content-Type', 'application/octet-stream')
+        reply.send(stream)
+      })
+
+      app.get('/typedarray', (request, reply) => {
+        reply.send(new Uint16Array(10))
+      })
+
+      getPort().then((port) => {
+        app.listen({ port }, () => {
+          axios = Axios.create({ baseURL: `http://localhost:${port}` })
+          done()
+        })
+        server = app.server
+      })
+    })
+
+    after(() => {
+      server.close()
+      return agent.close({ ritmReset: false })
+    })
+
+    beforeEach(() => {
+      config = new Config({
+        appsec: {
+          enabled: true,
+          rules: path.join(__dirname, 'api_security_rules.json'),
+          apiSecurity: {
+            enabled: true,
+            sampleDelay: 10
+          }
+        }
+      })
+      appsec.enable(config)
+    })
+
+    afterEach(() => {
+      appsec.disable()
+    })
+
+    function formatSchema (body) {
+      return zlib.gzipSync(JSON.stringify(body)).toString('base64')
+    }
+
+    it('should get the response body schema with reply.send', async () => {
+      const expectedResponseBodySchema = formatSchema([{ sendResKey: [8] }])
+      const res = await axios.post('/send', { key: 'value' })
+
+      await agent.assertFirstTraceSpan({
+        meta: {
+          '_dd.appsec.s.res.body': expectedResponseBodySchema
+        }
+      })
+
+      assert.equal(res.status, 200)
+      assert.deepEqual(res.data, { sendResKey: 'sendResValue' })
+    })
+
+    it('should get the response body schema with return', async () => {
+      const expectedResponseBodySchema = formatSchema([{ returnResKey: [8] }])
+      const res = await axios.post('/return', { key: 'value' })
+
+      await agent.assertFirstTraceSpan({
+        meta: {
+          '_dd.appsec.s.res.body': expectedResponseBodySchema
+        }
+      })
+
+      assert.equal(res.status, 200)
+      assert.deepEqual(res.data, { returnResKey: 'returnResValue' })
+    })
+
+    it('should not get the schema for string', async () => {
+      const res = await axios.get('/')
+
+      await agent.assertFirstTraceSpan(span => {
+        assert.notProperty(span.meta, '_dd.appsec.s.res.body')
+      })
+
+      assert.equal(res.status, 200)
+      assert.equal(res.data, 'DONE')
+    })
+
+    it('should not get the schema for Buffer', async () => {
+      const res = await axios.get('/buffer')
+
+      await agent.assertFirstTraceSpan(span => {
+        if (span.meta) {
+          assert.notProperty(span.meta, '_dd.appsec.s.res.body')
+        }
+      })
+
+      assert.equal(res.status, 200)
+      assert.equal(res.data, 'DONE')
+    })
+
+    it('should not get the schema for stream', async () => {
+      const res = await axios.get('/stream', { responseType: 'arraybuffer' })
+
+      await agent.assertFirstTraceSpan(span => {
+        if (span.meta) {
+          assert.notProperty(span.meta, '_dd.appsec.s.res.body')
+        }
+      })
+
+      assert.equal(res.status, 200)
+    })
+
+    it('should not get the schema for TypedArray', async () => {
+      const res = await axios.get('/typedarray', { responseType: 'arraybuffer' })
+
+      await agent.assertFirstTraceSpan(span => {
+        if (span.meta) {
+          assert.notProperty(span.meta, '_dd.appsec.s.res.body')
+        }
+      })
+
+      assert.equal(res.status, 200)
+    })
+  })
+})
+
 const createNestedObject = (n, obj) => {
   if (n > 0) {
     return { a: createNestedObject(n - 1, obj) }