Merge pull request #126 from proxyco/issue-109

#109 Implement a whitelist for headers as an option to LambdaAPI

Author: jeremydaly

index.d.ts

@@ -94,6 +94,7 @@ export declare interface Options {
   };
   serializer?: SerializerFunction;
   version?: string;
+  errorHeaderWhitelist?: string[];
 }
 
 export declare class Request {

index.js

@@ -26,6 +26,7 @@ class API {
     this._callbackName = props && props.callback ? props.callback.trim() : 'callback'
     this._mimeTypes = props && props.mimeTypes && typeof props.mimeTypes === 'object' ? props.mimeTypes : {}
     this._serializer = props && props.serializer && typeof props.serializer === 'function' ? props.serializer : JSON.stringify
+    this._errorHeaderWhitelist = props && (props.errorHeaderWhitelist || []).map(header => header.toLowerCase())
 
     // Set sampling info
     this._sampleCounts = {}
@@ -253,8 +254,17 @@ class API {
     // Error messages should never be base64 encoded
     response._isBase64 = false
 
-    // Strip the headers (TODO: find a better way to handle this)
-    response._headers = {}
+    // Strip the headers, keep whitelist
+    const strippedHeaders = Object.entries(response._headers).reduce((acc, [headerName, value]) => {
+      if (!this._errorHeaderWhitelist.includes(headerName.toLowerCase())) return acc
+    
+      return {
+        ...acc,
+        [headerName]: value
+      }
+    }, {})
+
+    response._headers = strippedHeaders
 
     let message
 

test/headers.js

@@ -3,7 +3,14 @@
 const expect = require('chai').expect // Assertion library
 
 // Init API instance
-const api = require('../index')({ version: 'v1.0' })
+
+const api = require('../index')({
+  version: 'v1.0',
+  errorHeaderWhitelist: [
+    'Access-Control-Allow-Origin',
+    'Access-Control-Allow-Methods',
+  ]
+})
 
 let event = {
   httpMethod: 'get',
@@ -93,6 +100,13 @@ api.get('/removeHeader', function(req,res) {
   })
 })
 
+api.get('/whitelistHeaders', function(req,res) {
+  res.status(200).header('TestStrippedHeader', 'RemoveMe')
+  res.status(200).header('access-control-allow-methods', ['GET, OPTIONS'])
+  res.status(200).header('access-control-allow-origin', ['example.com'])
+  throw new Error('TestError')
+})
+
 api.get('/cors', function(req,res) {
   res.cors().json({})
 })
@@ -244,6 +258,20 @@ describe('Header Tests:', function() {
       })
     }) // end it
 
+    it('Pass whitelisted headers on error', async function() {
+      let _event = Object.assign({},event,{ path: '/whitelistHeaders'})
+      let result = await new Promise(r => api.run(_event,{},(e,res) => { r(res) }))
+      expect(result).to.deep.equal({
+        multiValueHeaders: {
+          'content-type': ['application/json'],
+          'access-control-allow-methods': ['GET, OPTIONS'],
+          'access-control-allow-origin': ['example.com'],
+        }, statusCode: 500,
+        body: '{"error":"TestError"}',
+        isBase64Encoded: false
+      })
+    }) // end it
+
   }) // end Standard tests
 
   describe('CORS Tests:', function() {