Attempt to fix #21

__tests__/index.test.js

@@ -65,20 +65,35 @@ describe('middleware', () => {
             });
     });
 
-    it.skip('should work with res.end', done => {
+    it('should work with just res.end', done => {
+        app.get('/foo', (req, res) => {
+            res.set('Content-Type', 'text/plain');
+            res.end('foo');
+        });
+
+        request(app)
+            .get('/foo')
+            .end(() => {
+                expect(beforeFn).toHaveBeenCalled();
+                expect(afterFn).toHaveBeenCalledWith('foo');
+                done();
+            });
+    });
+
+    it('should work with res.write + res.end', done => {
         app.get('/foo', (req, res) => {
             res.set('Content-Type', 'text/plain');
             res.write('foo');
             res.write('bar');
             res.write('baz');
-            res.end();
+            res.end('qux');
         });
 
         request(app)
             .get('/foo')
             .end(() => {
                 expect(beforeFn).toHaveBeenCalled();
-                expect(afterFn).toHaveBeenCalledWith('foobarbaz');
+                expect(afterFn).toHaveBeenCalledWith('foobarbazqux');
                 done();
             });
     });

package.json

@@ -43,5 +43,8 @@
     "hooks": {
       "pre-commit": "yarn run lint"
     }
+  },
+  "dependencies": {
+    "get-stream": "^4.1.0"
   }
 }

src/index.js

@@ -1,8 +1,11 @@
 // @flow
+
 import type { $Request, $Response, Middleware, NextFunction } from 'express';
+import getStream from 'get-stream';
+import { PassThrough } from 'stream';
 
-export type RequestDetails = {
-    responseBody: string,
+export type RequestContext = {
+    responseBody?: string,
 };
 
 export type Tween = (
@@ -11,36 +14,98 @@ export type Tween = (
     $Response,
 ) => Promise<void>;
 
-export const getPatchedSendFn = (
-    context: Object,
-    res: $Response,
-): $PropertyType<$Response, 'send'> => {
-    const oldSend = res.send.bind(res);
-
-    return function send(...args: any) {
-        // eslint-disable-next-line prefer-destructuring, no-param-reassign
-        context.responseBody = args[0];
-        oldSend(...args);
-        /* istanbul ignore next */
-        return this;
-    }.bind(res);
-};
-
 export default function tweenz(...tweens: Array<Tween>): Middleware {
     return (req: $Request, res: $Response, next: NextFunction) => {
-        // Save some stateful request context
-        const context = {};
+        // Save some stateful request context that we'll pass back to the tween
+        const context: RequestContext = {};
+
+        // Create a noop pass through pipe and tap into the response body
+        const passThru = new PassThrough();
+        passThru.pipe(res);
+
+        // Set up a promise to wait for the response pipe to complete, and save the body
+        getStream(passThru).then(responseBody => {
+            context.responseBody = responseBody;
+        });
+
+        // We're about to write over res.write and res.end - save the old versions
+        const oldResWrite = res.write.bind(res);
+        const oldResEnd = res.end.bind(res);
+
+        /**
+         * This is a little bit suspect but here's what's going on:
+         *
+         * (1) We create a new pass through stream, 'passThru', and attach it to `res`.*
+         * Anything that flows through `passThru` will get magically passed along to `res`.
+         * e.g. `passThru.write()` eventually calls `res.write()`.
+         *
+         * *res is an http.ServerResponse, which is a Writable Stream. See:
+         * https://nodejs.org/api/stream.html#stream_api_for_stream_consumers
+         *
+         * (2) We create new proxy methods for `res.write` and `res.end`, and overwrite
+         * them back to the res object. This allows us to intercept when an application
+         * calls one of these methods, so we can make them flow through `passThru`.
+         * In other words, when `res.write()` gets hit by a consumer, our proxy will be
+         * called instead, which will in turn call `passThru.write()`
+         *
+         * (3) But wait! If we just had the proxy call the pass through, we'd end up in
+         * and infinite loop, since the pass through would now call the proxy :) To address
+         * this, inside our proxy, we need to work out if we should call the pass through
+         * method (`passThru.write`), or the actual original method (`oldResWrite`) from
+         * inside our proxy.
+         *
+         * (4) We currently do this by "remembering" how many times `res.write` / `res.end`
+         * has been called. If we haven't been called before, then assume it's the application
+         * code calling `res.write`. We should reroute to the pass through method if so. If
+         * we've already been called, then assume we're now being called by the pass through
+         * method. Here, we want to call the original `oldResWrite` method, to avoid the
+         * infinite loop.
+         *
+         * (5) We "remember" by setting a flag (`acceptWrite`, `acceptEnd`). This isn't perfect
+         * since we could have multiple `res.write()`, so we have to remember to flip it back
+         * ASAP. This is a little bit gross, and we should look for a better way to do this.
+         */
+
+        let acceptWrite = true;
+        let acceptEnd = true;
+
+        // $FlowFixMe: Ack that res.write is not writable - we're doing mad science
+        res.write = function PatchedWrite(...argumentsList) {
+            if (acceptWrite) {
+                acceptWrite = false;
+                passThru.write(...argumentsList);
+            } else {
+                acceptWrite = true;
+                oldResWrite(...argumentsList);
+            }
+
+            return res;
+        };
+
+        // $FlowFixMe: Ack that res.end is not writable - we're doing mad science
+        res.end = function PatchedEnd(...argumentsList) {
+            // Avoids the case where application code calls write + end of double
+            // writing to the passThru stream. (end() already calls write() internally.)
+            // @see: https://github.com/sharkcore/tweenz/pull/22#issuecomment-449688350
+            //
+            // $FlowFixMe: Ack that res.write is not writable - we're doing mad science
+            res.write = oldResWrite;
+
+            if (acceptEnd) {
+                acceptEnd = false;
+                passThru.end(...argumentsList);
+            } else {
+                acceptEnd = true;
+                oldResEnd(...argumentsList);
+            }
 
-        // Monkey patch express methods to intercept response body
-        // $FlowFixMe: https://github.com/facebook/flow/issues/3076
-        res.send = getPatchedSendFn(context, res);
+            return res;
+        };
 
         // Construct a Promise to be fulfilled when request has finished
         const requestDetails = new Promise(resolve => {
             function finish() {
-                resolve({
-                    responseBody: context.responseBody,
-                });
+                resolve(context);
                 // eslint-disable-next-line no-use-before-define
                 removeListeners();
             }