diff --git a/operations/development-server.js b/operations/development-server.js index 4c56965..0ef2c39 100644 --- a/operations/development-server.js +++ b/operations/development-server.js @@ -5,35 +5,57 @@ const defaultValue = require("default-value"); const assureArray = require("assure-array"); const path = require("path"); const entities = require("entities"); +const util = require("util"); -let reloadClientTag = ""; +let reloadClientTag = Buffer.from(""); -function monkeyPatchWrite(res) { +function monkeyPatchEnd(res) { // NOTE: This is a hack, to auto-inject the Budo livereload client in any outgoing HTML response (but only in development mode). - let write = res.write.bind(res); + let end = res.end.bind(res); - res.write = function monkeyPatchedWrite(... args) { - if (res.getHeader("content-type").startsWith("text/html")) { - let contentLength = res.getHeader("content-length"); - - if (contentLength != null) { - // Compensate for the additional bytes introduced by our injected script tag - res.setHeader("content-length", contentLength + reloadClientTag.length); + res.end = function monkeyPatchedEnd(... args) { + let originalChunk = args[0]; + let prefix; + + if (originalChunk == null || Buffer.isBuffer(originalChunk)) { + if (res.getHeader("content-type").startsWith("text/html")) { + let contentLength = res.getHeader("content-length"); + + if (contentLength != null) { + // Compensate for the additional bytes introduced by our injected script tag + res.setHeader("content-length", parseInt(contentLength) + reloadClientTag.length); + } + + prefix = reloadClientTag; } - write(reloadClientTag); - } - - // Reset the write method back to the original method; we don't need to get in the way anymore - res.write = write; + // Reset the `end` method back to the original method; we don't need to get in the way anymore + res.end = end; + + if (prefix != null) { + let chunk = (originalChunk == null) + ? reloadClientTag + : Buffer.concat([ originalChunk, reloadClientTag ]); - write(... args); + end(chunk, ... args.slice(1)); + } else { + end(... args); + } + } else { + throw new Error(`Expected a buffer in 'end', but got something else: ${util.inspect(args[0])}`); + } }; } function createExpressMiddleware(app) { return function expressMiddleware(req, res, next) { - monkeyPatchWrite(res); + try { + monkeyPatchEnd(res); + } catch (error) { + process.nextTick(() => { + throw error; + }); + } app.handle(req, res, (err) => { if (err != null && err instanceof Error) {