Merge pull request #675 from PlasmaPower/flush-headers

Add support for flushing headers

docs/api/response.md

@@ -289,3 +289,7 @@ this.response.etag = crypto.createHash('md5').update(this.body).digest('hex');
 ### response.vary(field)
 
   Vary on `field`.
+
+### response.flushHeaders()
+
+  Flush any set headers, and begin the body.

lib/application.js

@@ -181,7 +181,7 @@ function respond(ctx) {
   if (false === ctx.respond) return;
 
   const res = ctx.res;
-  if (res.headersSent || !ctx.writable) return;
+  if (!ctx.writable) return;
 
   let body = ctx.body;
   const code = ctx.status;
@@ -194,15 +194,19 @@ function respond(ctx) {
   }
 
   if ('HEAD' == ctx.method) {
-    if (isJSON(body)) ctx.length = Buffer.byteLength(JSON.stringify(body));
+    if (!res.headersSent && isJSON(body)) {
+      ctx.length = Buffer.byteLength(JSON.stringify(body));
+    }
     return res.end();
   }
 
   // status body
   if (null == body) {
-    ctx.type = 'text';
     body = ctx.message || String(code);
-    ctx.length = Buffer.byteLength(body);
+    if (!res.headersSent) {
+      ctx.type = 'text';
+      ctx.length = Buffer.byteLength(body);
+    }
     return res.end(body);
   }
 
@@ -213,6 +217,8 @@ function respond(ctx) {
 
   // body: json
   body = JSON.stringify(body);
-  ctx.length = Buffer.byteLength(body);
+  if (!res.headersSent) {
+    ctx.length = Buffer.byteLength(body);
+  }
   res.end(body);
 }

lib/context.js

@@ -149,6 +149,7 @@ delegate(proto, 'response')
   .method('vary')
   .method('set')
   .method('append')
+  .method('flushHeaders')
   .access('status')
   .access('message')
   .access('body')

lib/response.js

@@ -80,6 +80,7 @@ module.exports = {
   set status(code) {
     assert('number' == typeof code, 'status code must be a number');
     assert(statuses[code], `invalid status code: ${code}`);
+    assert(!this.res.headersSent, 'headers have already been sent');
     this._explicitStatus = true;
     this.res.statusCode = code;
     this.res.statusMessage = statuses[code];
@@ -130,6 +131,8 @@ module.exports = {
     const original = this._body;
     this._body = val;
 
+    if (this.res.headersSent) return;
+
     // no content
     if (null == val) {
       if (!statuses.empty[this.status]) this.status = 204;
@@ -521,5 +524,12 @@ module.exports = {
       'message',
       'header'
     ]);
+  },
+
+  /**
+   * Flush any set headers, and begin the body
+   */
+  flushHeaders() {
+    this.res.writeHead(this.res.statusCode);
   }
 };

test/application/respond.js

@@ -234,7 +234,12 @@ describe('app.respond', () => {
         ctx.status = 200;
         res.setHeader('Content-Type', 'text/html');
         res.write('Hello');
-        setTimeout(() => res.end('Goodbye'), 0);
+        return new Promise(resolve => {
+          setTimeout(() => {
+            res.end('Goodbye');
+            resolve();
+          }, 0);
+        });
       });
 
       const server = app.listen();

test/response/flushHeaders.js

@@ -0,0 +1,98 @@
+
+'use strict';
+
+const request = require('supertest');
+const assert = require('assert');
+const Koa = require('../..');
+
+describe('ctx.flushHeaders()', () => {
+  it('should set headersSent', done => {
+    const app = new Koa();
+
+    app.use((ctx, next) => {
+      ctx.body = 'Body';
+      ctx.status = 200;
+      ctx.flushHeaders();
+      assert(ctx.res.headersSent);
+    });
+
+    const server = app.listen();
+
+    request(server)
+      .get('/')
+      .expect(200)
+      .expect('Body', done);
+  });
+
+  it('should allow a response afterwards', done => {
+    const app = new Koa();
+
+    app.use((ctx, next) => {
+      ctx.status = 200;
+      ctx.res.setHeader('Content-Type', 'text/plain');
+      ctx.flushHeaders();
+      ctx.body = 'Body';
+    });
+
+    const server = app.listen();
+    request(server)
+      .get('/')
+      .expect(200)
+      .expect('Content-Type', 'text/plain')
+      .expect('Body', done);
+  });
+
+  it('should send the correct status code', done => {
+    const app = new Koa();
+
+    app.use((ctx, next) => {
+      ctx.status = 401;
+      ctx.res.setHeader('Content-Type', 'text/plain');
+      ctx.flushHeaders();
+      ctx.body = 'Body';
+    });
+
+    const server = app.listen();
+    request(server)
+      .get('/')
+      .expect(401)
+      .expect('Content-Type', 'text/plain')
+      .expect('Body', done);
+  });
+
+  it('should fail to set the headers after flushHeaders', done => {
+    const app = new Koa();
+
+    app.use((ctx, next) => {
+      ctx.status = 401;
+      ctx.res.setHeader('Content-Type', 'text/plain');
+      ctx.flushHeaders();
+      ctx.body = '';
+      try {
+        ctx.set('X-Shouldnt-Work', 'Value');
+      } catch (err) {
+        ctx.body += 'ctx.set fail ';
+      }
+      try {
+        ctx.status = 200;
+      } catch (err) {
+        ctx.body += 'ctx.status fail ';
+      }
+      try {
+        ctx.length = 10;
+      } catch (err) {
+        ctx.body += 'ctx.length fail';
+      }
+    });
+
+    const server = app.listen();
+    request(server)
+      .get('/')
+      .expect(401)
+      .expect('Content-Type', 'text/plain')
+      .expect('ctx.set fail ctx.status fail ctx.length fail', (err, res) => {
+        assert(res.headers['x-shouldnt-work'] === undefined, 'header set after flushHeaders');
+        done(err);
+      });
+  });
+});