Add support for headers in errors

Author: PlasmaPower

lib/context.js

@@ -117,8 +117,26 @@ var proto = module.exports = {
       return;
     }
 
-    // unset all headers
+    // unset all headers not explicitly kept
+    // while it's inconvenient to use setHeader and getHeader,
+    // res._headers isn't documentented and can't be relied on much
+    // for example: the keys are always stored in lowercase
+    var errHeaders = err.headers ? Object.keys(err.headers) : [];
+    var newHeaders = {};
+    for (var i = 0; i < errHeaders.length; i++) {
+      var header = errHeaders[i];
+      if (err.headers[header] === true) { // keep the header
+        newHeaders[header] = this.res.getHeader(header);
+      } else {
+        newHeaders[header] = err.headers[header];
+      }
+    }
     this.res._headers = {};
+    var newHeaderKeys = Object.keys(newHeaders);
+    for (var i = 0; i < newHeaderKeys.length; i++) {
+      var key = newHeaderKeys[i];
+      this.res.setHeader(key, newHeaders[key]);
+    }
 
     // force text/plain
     this.type = 'text';

test/context/onerror.js

@@ -52,6 +52,43 @@ describe('ctx.onerror(err)', function(){
     })
   })
 
+  it('should keep headers specified in the error', function(done){
+    var app = koa();
+
+    app.use(function *(next){
+      this.set('Vary', 'Accept-Encoding');
+      this.set('X-CSRF-Token', 'asdf');
+      this.set('X-Example-Kept', 'should be kept');
+      this.body = 'response';
+
+      throw Object.assign(new Error('boom'), {
+        status: 418,
+        expose: true,
+        headers: {
+          'X-Example-Kept': true,
+          'X-New-Header': 'Value'
+        }
+      })
+    })
+
+    var server = app.listen();
+
+    request(server)
+    .get('/')
+    .expect(418)
+    .expect('Content-Type', 'text/plain; charset=utf-8')
+    .expect('X-Example-Kept', 'should be kept')
+    .expect('X-New-Header', 'Value')
+    .end(function(err, res){
+      if (err) return done(err);
+
+      res.headers.should.not.have.property('vary');
+      res.headers.should.not.have.property('x-csrf-token');
+
+      done();
+    })
+  })
+
   describe('when invalid err.status', function(){
     describe('not number', function(){
       it('should respond 500', function(done){