fix: move via header to proxyRes event

tyler-johnson · tyler-johnson · commit 843388c6053e · 2016-09-14T19:05:45.000-06:00
diff --git a/src/index.js b/src/index.js
@@ -35,6 +35,12 @@ export default function(fn, opts={}) {
 
 	const proxy = httpProxy.createProxyServer({ target });
 
+	proxy.on("proxyRes", function(proxyRes, req, res) {
+		const existing = res.getHeader("Via");
+		const viaheader = `${existing ? existing + ", " : ""}${req.httpVersion} ${via} (${name}/${version})`;
+		res.setHeader("Via", viaheader);
+	});
+
 	return async function(req, res, next) {
 		try {
 			// hijack the root response and inject proxy information
@@ -51,20 +57,6 @@ export default function(fn, opts={}) {
 				if (secret) req.headers[headerFields.token] = signRequest(name, secret);
 			}
 
-			// attach Via header on response
-			// do this last in case there was an error
-			if (via) {
-				const writeHead = res.writeHead;
-				res.writeHead = function(code, headers) {
-					const existing = res.getHeader("Via");
-					const viaheader = `${existing ? existing + ", " : ""}${req.httpVersion} ${via} (${name}/${version})`;
-					res.setHeader("Via", viaheader);
-					if (headers) headers["Via"] = viaheader;
-
-					return writeHead.apply(res, arguments);
-				};
-			}
-
 			proxy.web(req, res);
 		} catch(e) {
 			if (next) next(e);
