Suspicious Response Blocking

packages/datadog-instrumentations/src/http/server.js

@@ -10,6 +10,7 @@ const startServerCh = channel('apm:http:server:request:start')
 const exitServerCh = channel('apm:http:server:request:exit')
 const errorServerCh = channel('apm:http:server:request:error')
 const finishServerCh = channel('apm:http:server:request:finish')
+const startWriteHeadCh = channel('apm:http:server:response:writeHead:start')
 const finishSetHeaderCh = channel('datadog:http:server:response:set-header:finish')
 
 const requestFinishedSet = new WeakSet()
@@ -20,6 +21,9 @@ const httpsNames = ['https', 'node:https']
 addHook({ name: httpNames }, http => {
   shimmer.wrap(http.ServerResponse.prototype, 'emit', wrapResponseEmit)
   shimmer.wrap(http.Server.prototype, 'emit', wrapEmit)
+  shimmer.wrap(http.ServerResponse.prototype, 'writeHead', wrapWriteHead)
+  shimmer.wrap(http.ServerResponse.prototype, 'write', wrapWrite)
+  shimmer.wrap(http.ServerResponse.prototype, 'end', wrapEnd)
   return http
 })
 
@@ -86,3 +90,97 @@ function wrapSetHeader (res) {
     }
   })
 }
+
+function wrapWriteHead (writeHead) {
+  return function wrappedWriteHead (statusCode, reason, obj) {
+    if (!startWriteHeadCh.hasSubscribers) {
+      return writeHead.apply(this, arguments)
+    }
+
+    const abortController = new AbortController()
+
+    if (typeof reason !== 'string') {
+      obj ??= reason
+    }
+
+    // support writeHead(200, ['key1', 'val1', 'key2', 'val2'])
+    if (Array.isArray(obj)) {
+      const headers = {}
+
+      for (let i = 0; i < obj.length; i += 2) {
+        headers[obj[i]] = obj[i + 1]
+      }
+
+      obj = headers
+    }
+
+    // this doesn't support explicit duplicate headers, but it's an edge case
+    const responseHeaders = Object.assign(this.getHeaders(), obj)
+
+    startWriteHeadCh.publish({
+      req: this.req,
+      res: this,
+      abortController,
+      statusCode,
+      responseHeaders
+    })
+
+    if (abortController.signal.aborted) {
+      return this
+    }
+
+    return writeHead.apply(this, arguments)
+  }
+}
+
+function wrapWrite (write) {
+  return function wrappedWrite () {
+    if (!startWriteHeadCh.hasSubscribers) {
+      return write.apply(this, arguments)
+    }
+
+    const abortController = new AbortController()
+
+    const responseHeaders = this.getHeaders()
+
+    startWriteHeadCh.publish({
+      req: this.req,
+      res: this,
+      abortController,
+      statusCode: this.statusCode,
+      responseHeaders
+    })
+
+    if (abortController.signal.aborted) {
+      return true
+    }
+
+    return write.apply(this, arguments)
+  }
+}
+
+function wrapEnd (end) {
+  return function wrappedEnd () {
+    if (!startWriteHeadCh.hasSubscribers) {
+      return end.apply(this, arguments)
+    }
+
+    const abortController = new AbortController()
+
+    const responseHeaders = this.getHeaders()
+
+    startWriteHeadCh.publish({
+      req: this.req,
+      res: this,
+      abortController,
+      statusCode: this.statusCode,
+      responseHeaders
+    })
+
+    if (abortController.signal.aborted) {
+      return this
+    }
+
+    return end.apply(this, arguments)
+  }
+}

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

@@ -111,6 +111,10 @@ function block (req, res, rootSpan, abortController, actionParameters) {
 
   const { body, headers, statusCode } = getBlockingData(req, null, rootSpan, actionParameters)
 
+  for (const headerName of res.getHeaderNames()) {
+    res.removeHeader(headerName)
+  }
+
   res.writeHead(statusCode, headers).end(body)
 
   abortController?.abort()

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

@@ -18,5 +18,6 @@ module.exports = {
   nextBodyParsed: dc.channel('apm:next:body-parsed'),
   nextQueryParsed: dc.channel('apm:next:query-parsed'),
   responseBody: dc.channel('datadog:express:response:json:start'),
+  responseWriteHead: dc.channel('apm:http:server:response:writeHead:start'),
   httpClientRequestStart: dc.channel('apm:http:client:request:start')
 }

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

@@ -12,7 +12,8 @@ const {
   queryParser,
   nextBodyParsed,
   nextQueryParsed,
-  responseBody
+  responseBody,
+  responseWriteHead
 } = require('./channels')
 const waf = require('./waf')
 const addresses = require('./addresses')
@@ -60,6 +61,7 @@ function enable (_config) {
     queryParser.subscribe(onRequestQueryParsed)
     cookieParser.subscribe(onRequestCookieParser)
     responseBody.subscribe(onResponseBody)
+    responseWriteHead.subscribe(onResponseWriteHead)
 
     if (_config.appsec.eventTracking.enabled) {
       passportVerify.subscribe(onPassportVerify)
@@ -110,14 +112,7 @@ function incomingHttpStartTranslator ({ req, res, abortController }) {
 }
 
 function incomingHttpEndTranslator ({ req, res }) {
-  // TODO: this doesn't support headers sent with res.writeHead()
-  const responseHeaders = Object.assign({}, res.getHeaders())
-  delete responseHeaders['set-cookie']
-
-  const persistent = {
-    [addresses.HTTP_INCOMING_RESPONSE_CODE]: '' + res.statusCode,
-    [addresses.HTTP_INCOMING_RESPONSE_HEADERS]: responseHeaders
-  }
+  const persistent = {}
 
   // we need to keep this to support other body parsers
   // TODO: no need to analyze it if it was already done by the body-parser hook
@@ -139,7 +134,9 @@ function incomingHttpEndTranslator ({ req, res }) {
     persistent[addresses.HTTP_INCOMING_QUERY] = req.query
   }
 
-  waf.run({ persistent }, req)
+  if (Object.keys(persistent).length) {
+    waf.run({ persistent }, req)
+  }
 
   waf.disposeContext(req)
 
@@ -225,12 +222,48 @@ function onPassportVerify ({ credentials, user }) {
   passportTrackEvent(credentials, user, rootSpan, config.appsec.eventTracking.mode)
 }
 
+const responseAnalyzedSet = new WeakSet()
+const responseBlockedSet = new WeakSet()
+
+function onResponseWriteHead ({ req, res, abortController, statusCode, responseHeaders }) {
+  // avoid "write after end" error
+  if (responseBlockedSet.has(res)) {
+    abortController?.abort()
+    return
+  }
+
+  // avoid double waf call
+  if (responseAnalyzedSet.has(res)) {
+    return
+  }
+
+  const rootSpan = web.root(req)
+  if (!rootSpan) return
+
+  responseHeaders = Object.assign({}, responseHeaders)
+  delete responseHeaders['set-cookie']
+
+  const results = waf.run({
+    persistent: {
+      [addresses.HTTP_INCOMING_RESPONSE_CODE]: '' + statusCode,
+      [addresses.HTTP_INCOMING_RESPONSE_HEADERS]: responseHeaders
+    }
+  }, req)
+
+  responseAnalyzedSet.add(res)
+
+  handleResults(results, req, res, rootSpan, abortController)
+}
+
 function handleResults (actions, req, res, rootSpan, abortController) {
   if (!actions || !req || !res || !rootSpan || !abortController) return
 
   const blockingAction = getBlockingAction(actions)
   if (blockingAction) {
     block(req, res, rootSpan, abortController, blockingAction)
+    if (!abortController.signal || abortController.signal.aborted) {
+      responseBlockedSet.add(res)
+    }
   }
 }
 
@@ -256,6 +289,7 @@ function disable () {
   if (cookieParser.hasSubscribers) cookieParser.unsubscribe(onRequestCookieParser)
   if (responseBody.hasSubscribers) responseBody.unsubscribe(onResponseBody)
   if (passportVerify.hasSubscribers) passportVerify.unsubscribe(onPassportVerify)
+  if (responseWriteHead.hasSubscribers) responseWriteHead.unsubscribe(onResponseWriteHead)
 }
 
 module.exports = {

packages/dd-trace/src/appsec/remote_config/capabilities.js

@@ -6,6 +6,7 @@ module.exports = {
   ASM_DD_RULES: 1n << 3n,
   ASM_EXCLUSIONS: 1n << 4n,
   ASM_REQUEST_BLOCKING: 1n << 5n,
+  ASM_RESPONSE_BLOCKING: 1n << 6n,
   ASM_USER_BLOCKING: 1n << 7n,
   ASM_CUSTOM_RULES: 1n << 8n,
   ASM_CUSTOM_BLOCKING_RESPONSE: 1n << 9n,

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

@@ -71,6 +71,7 @@ function enableWafUpdate (appsecConfig) {
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_DD_RULES, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_EXCLUSIONS, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_REQUEST_BLOCKING, true)
+    rc.updateCapabilities(RemoteConfigCapabilities.ASM_RESPONSE_BLOCKING, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_RULES, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_BLOCKING_RESPONSE, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_TRUSTED_IPS, true)
@@ -92,6 +93,7 @@ function disableWafUpdate () {
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_DD_RULES, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_EXCLUSIONS, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_REQUEST_BLOCKING, false)
+    rc.updateCapabilities(RemoteConfigCapabilities.ASM_RESPONSE_BLOCKING, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_RULES, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_BLOCKING_RESPONSE, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_TRUSTED_IPS, false)

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

@@ -37,7 +37,9 @@ describe('blocking', () => {
     res = {
       setHeader: sinon.stub(),
       writeHead: sinon.stub(),
-      end: sinon.stub()
+      end: sinon.stub(),
+      getHeaderNames: sinon.stub().returns([]),
+      removeHeader: sinon.stub()
     }
     res.writeHead.returns(res)
 
@@ -109,6 +111,22 @@ describe('blocking', () => {
       expect(res.end).to.have.been.calledOnceWithExactly('jsonBody')
       expect(abortController.signal.aborted).to.be.true
     })
+
+    it('should remove all headers before sending blocking response', () => {
+      res.getHeaderNames.returns(['header1', 'header2'])
+
+      block(req, res, rootSpan)
+
+      expect(rootSpan.addTags).to.have.been.calledOnceWithExactly({ 'appsec.blocked': 'true' })
+      expect(res.removeHeader).to.have.been.calledTwice
+      expect(res.removeHeader.firstCall).to.have.been.calledWithExactly('header1')
+      expect(res.removeHeader.secondCall).to.have.been.calledWithExactly('header2')
+      expect(res.writeHead).to.have.been.calledOnceWithExactly(403, {
+        'Content-Type': 'application/json',
+        'Content-Length': 8
+      })
+      expect(res.end).to.have.been.calledOnceWithExactly('jsonBody')
+    })
   })
 
   describe('block with default templates', () => {