From 9eadf69141679fa9903380dc6a2a308bd5f2591b Mon Sep 17 00:00:00 2001
From: Jakub Frej <frej@appliedgenerics.com>
Date: Tue, 16 Feb 2016 11:10:59 +0000
Subject: [PATCH] Added options for transformReq and transformResp functions

---
 README.md    |  2 ++
 index.js     | 18 ++++++++++++++++
 test/test.js | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/README.md b/README.md
index 6bcb842..7e14ba5 100644
--- a/README.md
+++ b/README.md
@@ -26,6 +26,8 @@ Other options:
 - `via`: by default no [via header](http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.45) is added. If you pass `true` for this option the local hostname will be used for the via header. You can also pass a string for this option in which case that will be used for the via header.
 - `cookieRewrite`: this option can be used to support cookies via the proxy by rewriting the cookie domain to that of the proxy server. By default cookie domains are not rewritten. The `cookieRewrite` option works as the `via` option - if you pass `true` the local hostname will be used, and if you pass a string that will be used as the rewritten cookie domain.
 - `preserveHost`: When enabled, this option will pass the Host: line from the incoming request to the proxied host. Default: `false`.
+- `transformReq`: transform function called on requests enables dynamic rewriting of headers sent to the server.
+- `transformResp`: transform function called on responses enables dynamic rewriting of headers coming back from the server.
 
 ### Usage with route:
 
diff --git a/index.js b/index.js
index b8c573e..627c01a 100644
--- a/index.js
+++ b/index.js
@@ -47,6 +47,7 @@ module.exports = function proxyMiddleware(options) {
     opts.headers = options.headers ? merge(req.headers, options.headers) : req.headers;
 
     applyViaHeader(req.headers, opts, opts.headers);
+    transformReq(req.headers, opts, opts.headers);
 
     if (!options.preserveHost) {
       // Forwarding the host breaks dotcloud
@@ -64,6 +65,7 @@ module.exports = function proxyMiddleware(options) {
       }
       applyViaHeader(myRes.headers, opts, myRes.headers);
       rewriteCookieHosts(myRes.headers, opts, myRes.headers, req);
+      transformResp(myRes.headers, opts, myRes.headers);
       resp.writeHead(myRes.statusCode, myRes.headers);
       myRes.on('error', function (err) {
         next(err);
@@ -118,6 +120,22 @@ function rewriteCookieHosts(existingHeaders, opts, applyTo, req) {
   applyTo['set-cookie'] = rewrittenCookies;
 }
 
+function transformReq (existingHeaders, opts, applyTo) {
+  if (typeof opts.transformReq !== 'function') {
+    return;
+  }
+
+  opts.transformReq(existingHeaders, opts, applyTo);
+}
+
+function transformResp (existingHeaders, opts, applyTo) {
+  if (typeof opts.transformResp !== 'function') {
+    return;
+  }
+
+  opts.transformResp(existingHeaders, opts, applyTo);
+}
+
 function slashJoin(p1, p2) {
   var trailing_slash = false;
 
diff --git a/test/test.js b/test/test.js
index 45e4f9b..6a98ef8 100644
--- a/test/test.js
+++ b/test/test.js
@@ -547,6 +547,65 @@ describe("proxy", function() {
     })
   });
 
+  it("calls the transformReq function and transforms the request headers", function(done) {
+    var destServer = createServerWithLibName('http', function(req, resp) {
+      assert.strictEqual(req.headers.host, 'localhost:8086');
+      assert.strictEqual(req.headers['x-custom-header'], 'hello localhost:8087');
+      resp.statusCode = 200;
+      resp.write(req.url);
+      resp.end();
+    });
+
+    var proxyOptions = url.parse('http://localhost:8086/');
+    proxyOptions.transformReq = function (existingHeaders, opts, applyTo) {
+      applyTo['x-custom-header'] = 'hello ' + existingHeaders.host;
+    };
+
+    var app = connect();
+    app.use(proxy(proxyOptions));
+
+    destServer.listen(8086, 'localhost', function() {
+      app.listen(8087);
+
+      var options = url.parse('http://localhost:8087/foo/test/');
+      http.get(options, function () {
+        // ok...
+        done();
+      }).on('error', function () {
+        assert.fail('Request proxy failed');
+      });
+    });
+  });
+
+  it("calls the transformResp function and transforms the response headers", function(done) {
+    var destServer = createServerWithLibName('http', function(req, resp) {
+      resp.statusCode = 200;
+      resp.setHeader('x-custom-header', 'hello');
+      resp.write(req.url);
+      resp.end();
+    });
+
+    var proxyOptions = url.parse('http://localhost:8088/');
+    proxyOptions.transformResp = function (existingHeaders, opts, applyTo) {
+      applyTo['x-custom-reply'] = 'hi there ' + existingHeaders['x-custom-header'];
+    };
+
+    var app = connect();
+    app.use(proxy(proxyOptions));
+
+    destServer.listen(8088, 'localhost', function() {
+      app.listen(8089);
+
+      var options = url.parse('http://localhost:8089/foo/test/');
+      http.get(options, function (resp) {
+        assert.strictEqual(resp.headers['x-custom-reply'], 'hi there hello');
+        done();
+      }).on('error', function () {
+        assert.fail('Request proxy failed');
+      });
+    });
+  });
+
 });
 
 function createServerWithLibName(libName, requestListener) {