From 3c23aa5b74d19352e0a78d135df5ed6979d4e242 Mon Sep 17 00:00:00 2001
From: Yiyu He <dead_horse@qq.com>
Date: Sun, 11 Feb 2018 16:25:24 +0800
Subject: [PATCH] feat: ignore set header/status when header sent (#1137)

---
 lib/response.js               | 11 ++++++--
 test/application/respond.js   | 45 ++++++++++++++++++++++++++++++
 test/response/flushHeaders.js | 52 ++++++++++++++++++++++-------------
 3 files changed, 86 insertions(+), 22 deletions(-)

diff --git a/lib/response.js b/lib/response.js
index 6c663a7..fb6b897 100644
--- a/lib/response.js
+++ b/lib/response.js
@@ -80,9 +80,10 @@ module.exports = {
    */
 
   set status(code) {
+    if (this.headerSent) return;
+
     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;
     if (this.req.httpVersionMajor < 2) this.res.statusMessage = statuses[code];
@@ -133,8 +134,6 @@ 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;
@@ -233,6 +232,8 @@ module.exports = {
    */
 
   vary(field) {
+    if (this.headerSent) return;
+
     vary(this.res, field);
   },
 
@@ -434,6 +435,8 @@ module.exports = {
    */
 
   set(field, val) {
+    if (this.headerSent) return;
+
     if (2 == arguments.length) {
       if (Array.isArray(val)) val = val.map(String);
       else val = String(val);
@@ -481,6 +484,8 @@ module.exports = {
    */
 
   remove(field) {
+    if (this.headerSent) return;
+
     this.res.removeHeader(field);
   },
 
diff --git a/test/application/respond.js b/test/application/respond.js
index 65b3751..3bc6a74 100644
--- a/test/application/respond.js
+++ b/test/application/respond.js
@@ -39,6 +39,51 @@ describe('app.respond', () => {
         .expect(200)
         .expect('lol');
     });
+
+    it('should ignore set header after header sent', () => {
+      const app = new Koa();
+      app.use(ctx => {
+        ctx.body = 'Hello';
+        ctx.respond = false;
+
+        const res = ctx.res;
+        res.statusCode = 200;
+        res.setHeader('Content-Type', 'text/plain');
+        res.end('lol');
+        ctx.set('foo', 'bar');
+      });
+
+      const server = app.listen();
+
+      return request(server)
+        .get('/')
+        .expect(200)
+        .expect('lol')
+        .expect(res => {
+          assert(!res.headers.foo);
+        });
+    });
+
+    it('should ignore set status after header sent', () => {
+      const app = new Koa();
+      app.use(ctx => {
+        ctx.body = 'Hello';
+        ctx.respond = false;
+
+        const res = ctx.res;
+        res.statusCode = 200;
+        res.setHeader('Content-Type', 'text/plain');
+        res.end('lol');
+        ctx.status = 201;
+      });
+
+      const server = app.listen();
+
+      return request(server)
+        .get('/')
+        .expect(200)
+        .expect('lol');
+    });
   });
 
   describe('when this.type === null', () => {
diff --git a/test/response/flushHeaders.js b/test/response/flushHeaders.js
index 176db2f..b1c5de6 100644
--- a/test/response/flushHeaders.js
+++ b/test/response/flushHeaders.js
@@ -61,39 +61,27 @@ describe('ctx.flushHeaders()', () => {
       .expect('Body');
   });
 
-  it('should fail to set the headers after flushHeaders', async () => {
+  it('should ignore set header after flushHeaders', async () => {
     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';
-      }
+      ctx.body = 'foo';
+      ctx.set('X-Shouldnt-Work', 'Value');
+      ctx.remove('Content-Type');
+      ctx.vary('Content-Type');
     });
 
     const server = app.listen();
     const res = await request(server)
       .get('/')
       .expect(401)
-      .expect('Content-Type', 'text/plain')
-      .expect('ctx.set fail ctx.status fail ctx.length fail');
+      .expect('Content-Type', 'text/plain');
 
     assert.equal(res.headers['x-shouldnt-work'], undefined, 'header set after flushHeaders');
+    assert.equal(res.headers.vary, undefined, 'header set after flushHeaders');
   });
 
   it('should flush headers first and delay to send data', done => {
@@ -134,4 +122,30 @@ describe('ctx.flushHeaders()', () => {
       .end();
     });
   });
+
+  it('should catch stream error', done => {
+    const PassThrough = require('stream').PassThrough;
+    const app = new Koa();
+    app.once('error', err => {
+      assert(err.message === 'mock error');
+      done();
+    });
+
+    app.use(ctx => {
+      ctx.type = 'json';
+      ctx.status = 200;
+      ctx.headers['Link'] = '</css/mycss.css>; as=style; rel=preload, <https://img.craftflair.com>; rel=preconnect; crossorigin';
+      ctx.length = 20;
+      ctx.flushHeaders();
+      const stream = ctx.body = new PassThrough();
+
+      setTimeout(() => {
+        stream.emit('error', new Error('mock error'));
+      }, 100);
+    });
+
+    const server = app.listen();
+
+    request(server).get('/').end();
+  });
 });