refactor(responseBody): handle client close event

lib/middleware/response-body.js

@@ -14,6 +14,7 @@ module.exports = function transformResponseBody(middleware, filter) {
 
     var buf = []
     var length = 0
+    var closed = false
     var headArgs = null
 
     var _end = res.end
@@ -26,35 +27,52 @@ module.exports = function transformResponseBody(middleware, filter) {
     }
 
     res.write = function (data) {
+      var cb = getCallback(arguments)
+      if (closed) return cb && cb()
+
       if (data) {
         length += +data.length || 0
         buf.push(data)
       }
 
-      var cb = getCallback(arguments)
       if (cb) cb()
     }
 
     res.end = function (data) {
+      var cb = getCallback(arguments)
+      if (closed) return cb && cb()
+
       if (data && typeof data !== 'function') {
         buf.push(data)
       }
 
       var body = Buffer.concat(buf, length)
 
+      // Expose the body
+      res.body = res._originalBody = body
+
+      // Restore and clean references
+      cleanupAndRestore()
+
+      middleware(req, res, finisher(cb))
+    }
+
+    req.on('close', onClose)
+
+    function onClose() {
+      closed = true
+      cleanupAndRestore()
+    }
+
+    function cleanupAndRestore() {
       // Restore methods
       res.writeHead = _writeHead
       res.write = _write
       res.end = _end
 
-      // Expose the body
-      res.body = res._originalBody = body
-
       // Clean references to prevent leaks
       buf = body = _end = _write = headArgs = null
-
-      var cb = getCallback(arguments)
-      middleware(req, res, finisher(cb))
+      req.removeListener('close', onClose)
     }
 
     function finisher(cb) {

test/middleware/response-body.js

@@ -1,4 +1,5 @@
 const expect = require('chai').expect
+const Emitter = require('events').EventEmitter
 const middleware = require('../../lib/middleware')
 const noop = function () {}
 
@@ -13,6 +14,10 @@ suite('middleware#responseBody', function () {
     }
   })
 
+  beforeEach(function () {
+    req = new Emitter
+  })
+
   function middlewareFn(req, res, next) {
     var body = res.body.toString()
     var newBody = body.split(' ').reverse().join(' ')
@@ -33,7 +38,7 @@ suite('middleware#responseBody', function () {
 
     middleware.responseBody
       (middlewareFn)
-        (null, res, noop)
+        (req, res, noop)
 
     writeData()
   })
@@ -50,7 +55,7 @@ suite('middleware#responseBody', function () {
 
     middleware.responseBody
       (middlewareFn, filter)
-        (null, res, noop)
+        (req, res, noop)
 
     writeData()
   })