Merge pull request #9 from alexsanya/master

Add strippedHeaders and userResHeadersDecorator
nsimmons · Oct 20, 2020 · d4c94ae · d4c94ae
2 parents 879af61 + 39896a2
commit d4c94ae
Show file tree

Hide file tree

Showing 5 changed files with 121 additions and 25 deletions.
diff --git a/README.md b/README.md
@@ -50,6 +50,19 @@ app.use(proxy('www.google.com', {
 }));
 ```
 
+
+#### strippedHeaders
+
+Headers to remove from proxy response.
+
+```js
+app.use(proxy('www.google.com', {
+  strippedHeaders: [
+    'set-cookie'
+  ]
+}));
+```
+
 #### preserveReqSession
 
 Pass the session along to the proxied request
@@ -116,6 +129,10 @@ instance, but this is not a reliable interface. I expect to close this
 exploit in a future release, while providing an additional hook for mutating
 the userRes before sending.
 
+#### userResHeadersDecorator (supports Promise)
+
+You can modify the proxy's headers before sending it to the client.
+
 ##### gzip responses
 
 If your proxy response is gzipped, this program will automatically unzip

diff --git a/app/steps/copyProxyResHeadersToUserRes.js b/app/steps/copyProxyResHeadersToUserRes.js
@@ -4,19 +4,23 @@ function copyProxyResHeadersToUserRes(container) {
   return new Promise(function(resolve) {
     var ctx = container.user.ctx;
     var rsp = container.proxy.res;
+    var strippedHeaders = container.options.strippedHeaders || [];
+    var userResHeadersDecorator = container.options.userResHeadersDecorator || function(headers) { return headers; };
 
-    if (!ctx.headerSent && ctx.status !== 504) {
-      ctx.status = rsp.statusCode;
-      Object.keys(rsp.headers)
+    if (ctx.headerSent || ctx.status === 504) {
+      return resolve(container);
+    }
+    ctx.status = rsp.statusCode;
+    Promise.resolve(userResHeadersDecorator(rsp.headers)).then(function(decoratedHeaders) {
+      Object.keys(decoratedHeaders)
+      .filter(function(item) { return strippedHeaders.indexOf(item) < 0; })
       .filter(function(item) { return item !== 'transfer-encoding'; })
       .forEach(function(item) {
-        ctx.set(item, rsp.headers[item]);
+        ctx.set(item, decoratedHeaders[item]);
       });
-    }
-
-    resolve(container);
+      resolve(container);
+    });
   });
 }
 
 module.exports = copyProxyResHeadersToUserRes;
-
diff --git a/lib/resolveOptions.js b/lib/resolveOptions.js
@@ -19,11 +19,13 @@ function resolveOptions(options) {
     proxyReqOptDecorator: options.proxyReqOptDecorator,
     proxyReqBodyDecorator: options.proxyReqBodyDecorator,
     userResDecorator: options.userResDecorator,
+    userResHeadersDecorator: options.userResHeadersDecorator,
     filter: options.filter || defaultFilter,
     // For backwards compatability, we default to legacy behavior for newly added settings.
     parseReqBody: isUnset(options.parseReqBody) ? true : options.parseReqBody,
     reqBodyEncoding: resolveBodyEncoding(options.reqBodyEncoding),
     headers: options.headers,
+    strippedHeaders: options.strippedHeaders,
     preserveReqSession: options.preserveReqSession,
     https: options.https,
     port: options.port,

diff --git a/test/userResDecorator.js b/test/userResDecorator.js
@@ -5,11 +5,39 @@ var Koa = require('koa');
 var agent = require('supertest').agent;
 var proxy = require('../');
 
-describe('userResDecorator', function() {
+function proxyTarget(port) {
+  var other = new Koa();
+  other.use(function(ctx, next) {
+    if (ctx.request.url !== '/json') {
+      return next();
+    }
+    ctx.set('content-type', 'app,lication/json');
+    ctx.body = JSON.stringify({foo: 'bar'});
+  });
+  other.use(function(ctx) {
+    ctx.status = 200;
+    ctx.set('x-wombat-alliance', 'mammels');
+    ctx.set('x-custom-header', 'something');
+    ctx.body = 'Success';
+  });
+  return other.listen(port);
+}
+
+describe.only('userResDecorator', function() {
+  var other;
+
+  beforeEach(function() {
+    other = proxyTarget(8080);
+  });
+
+  afterEach(function() {
+    other.close();
+  });
 
   it('has access to original response', function(done) {
     var app = new Koa();
-    app.use(proxy('httpbin.org', {
+    app.use(proxy('http://localhost', {
+      port: 8080,
       userResDecorator: function(proxyRes, proxyResData) {
         assert(proxyRes.connection);
         assert(proxyRes.socket);
@@ -24,7 +52,8 @@ describe('userResDecorator', function() {
 
   it('works with promises', function(done) {
     var app = new Koa();
-    app.use(proxy('httpbin.org', {
+    app.use(proxy('http://localhost', {
+      port: 8080,
       userResDecorator: function(proxyRes, proxyResData) {
         return new Promise(function(resolve) {
           proxyResData.funkyMessage = 'oi io oo ii';
@@ -36,7 +65,7 @@ describe('userResDecorator', function() {
     }));
 
     agent(app.callback())
-    .get('/ip')
+    .get('/')
     .end(function(err, res) {
       if (err) { return done(err); }
 
@@ -48,7 +77,8 @@ describe('userResDecorator', function() {
 
   it('can modify the response data', function(done) {
     var app = new Koa();
-    app.use(proxy('httpbin.org', {
+    app.use(proxy('http://localhost', {
+      port: 8080,
       userResDecorator: function(proxyRes, proxyResData) {
         proxyResData = JSON.parse(proxyResData.toString('utf8'));
         proxyResData.intercepted = true;
@@ -57,7 +87,7 @@ describe('userResDecorator', function() {
     }));
 
     agent(app.callback())
-    .get('/ip')
+    .get('/json')
     .end(function(err, res) {
       if (err) { return done(err); }
 
@@ -66,39 +96,80 @@ describe('userResDecorator', function() {
     });
   });
 
+  it('can filter response headers', function(done) {
+    var proxiedApp = new Koa();
+    var app = new Koa();
+    var p1Done, p2Done;
+    var p1 = new Promise(function(resolve) { p1Done = resolve; });
+    var p2 = new Promise(function(resolve) { p2Done = resolve; });
+    app.use(proxy('http://localhost', {
+      port: 8080
+    }));
+    proxiedApp.use(proxy('http://localhost', {
+      port: 8080,
+      strippedHeaders: ['x-wombat-alliance', 'x-custom-header']
+    }));
 
-  it('can modify the response headers, [deviant case, supported by pass-by-reference atm]', function(done) {
+    agent(app.callback())
+    .get('/')
+    .end(function(err, res) {
+      if (err) { return done(err); }
+      assert(typeof res.headers['x-custom-header'] === 'string');
+      assert(typeof res.headers['x-wombat-alliance'] === 'string');
+      p1Done();
+    });
+
+    agent(proxiedApp.callback())
+    .get('/')
+    .end(function(err, res) {
+      if (err) { return done(err); }
+      assert(typeof res.headers['x-custom-header'] !== 'string');
+      assert(typeof res.headers['x-wombat-alliance'] !== 'string');
+      p2Done();
+    });
+
+    Promise.all([p1, p2]).then(function() { done(); });
+  });
+
+  it('can modify the response headers', function(done) {
     var app = new Koa();
-    app.use(proxy('httpbin.org', {
-      userResDecorator: function(rsp, data, ctx) {
-        ctx.set('x-wombat-alliance', 'mammels');
-        ctx.set('content-type', 'wiki/wiki');
-        return data;
+    app.use(proxy('http://localhost', {
+      port: 8080,
+      userResHeadersDecorator: function(headers) {
+        var newHeaders = Object.keys(headers)
+          .reduce(function(result, key) {
+            result[key] = headers[key];
+            return result;
+          }, {});
+        newHeaders['x-transaction-id'] = '12345';
+        newHeaders['x-entity-id'] = 'abcdef';
+        return newHeaders;
       }
     }));
 
     agent(app.callback())
     .get('/ip')
     .end(function(err, res) {
       if (err) { return done(err); }
-      assert(res.headers['content-type'] === 'wiki/wiki');
-      assert(res.headers['x-wombat-alliance'] === 'mammels');
+      assert(res.headers['x-transaction-id'] === '12345');
+      assert(res.headers['x-entity-id'] === 'abcdef');
       done();
     });
   });
 
   it('can mutuate an html response', function(done) {
     var app = new Koa();
-    app.use(proxy('httpbin.org', {
+    app.use(proxy('http://localhost', {
+      port: 8080,
       userResDecorator: function(rsp, data) {
-        data = data.toString().replace('Oh', '<strong>Hey</strong>');
+        data = data.toString().replace('Success', '<strong>Hey</strong>');
         assert(data !== '');
         return data;
       }
     }));
 
     agent(app.callback())
-    .get('/html')
+    .get('/')
     .end(function(err, res) {
       if (err) { return done(err); }
       assert(res.text.indexOf('<strong>Hey</strong>') > -1);

diff --git a/types.d.ts b/types.d.ts
@@ -6,6 +6,7 @@ declare function koaHttpProxy(host: string, options: koaHttpProxy.IOptions): koa
 declare namespace koaHttpProxy {
   export interface IOptions {
     headers?: { [key: string]: any },
+    strippedHeaders?: [string],
     https?: boolean,
     limit?: string,
     parseReqBody?: boolean,
@@ -21,6 +22,7 @@ declare namespace koaHttpProxy {
     proxyReqOptDecorator?(proxyReqOpts: IRequestOption, ctx: koa.Context): IRequestOption | Promise<IRequestOption>,
     proxyReqPathResolver?(ctx: koa.Context): string | Promise<string>,
     userResDecorator?(proxyRes: http.IncomingMessage, proxyResData: string | Buffer, ctx: koa.Context): string | Buffer | Promise<string> | Promise<Buffer>,
+    userResHeadersDecorator?(headers: {[key: string]: string}): Promise<{[key: string]: string}> | {[key: string]: string},
   }
 
   export interface IRequestOption {