diff --git a/__tests__/index.test.js b/__tests__/index.test.js index 17d920a..f26b246 100644 --- a/__tests__/index.test.js +++ b/__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(); }); }); diff --git a/package.json b/package.json index 4e51cf4..0c81a17 100644 --- a/package.json +++ b/package.json @@ -43,5 +43,8 @@ "hooks": { "pre-commit": "yarn run lint" } + }, + "dependencies": { + "get-stream": "^4.1.0" } } diff --git a/src/index.js b/src/index.js index 28d83db..b69510c 100644 --- a/src/index.js +++ b/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; -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): 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(); }