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) {