From 48ac0669c5e3c92560b5f3c38e0aaadcc65905fb Mon Sep 17 00:00:00 2001
From: dead_horse <dead_horse@qq.com>
Date: Mon, 28 Apr 2014 20:34:26 -0700
Subject: [PATCH] fix content-length when body is re-assigned. Closes #267

---
 lib/response.js         | 11 ++++++-
 test/application.js     | 72 +++++++++++++++++++++++++++++++++++++++++
 test/response/length.js | 27 ++++++++++++++--
 3 files changed, 107 insertions(+), 3 deletions(-)

diff --git a/lib/response.js b/lib/response.js
index e4f8aa7..ef15f95 100644
--- a/lib/response.js
+++ b/lib/response.js
@@ -4,6 +4,7 @@
  */
 
 var ensureErrorHandler = require('error-inject');
+var isJSON = require('koa-is-json');
 var escape = require('escape-html');
 var onfinish = require('finished');
 var getType = require('set-type');
@@ -98,6 +99,7 @@ module.exports = {
    */
 
   set body(val) {
+    var original = this._body;
     this._body = val;
 
     // no content
@@ -133,11 +135,16 @@ module.exports = {
     if ('function' == typeof val.pipe) {
       onfinish(this, destroy.bind(null, val));
       ensureErrorHandler(val, this.ctx.onerror);
+
+      // overwriting
+      if (null != original && original != val) this.remove('Content-Length');
+
       if (setType) this.type = 'bin';
       return;
     }
 
     // json
+    this.remove('Content-Length');
     this.type = 'json';
   },
 
@@ -166,7 +173,9 @@ module.exports = {
     if (null == len) {
       if (!body) return;
       if ('string' == typeof body) return Buffer.byteLength(body);
-      return body.length;
+      if (Buffer.isBuffer(body)) return body.length;
+      if (isJSON(body)) return Buffer.byteLength(JSON.stringify(body));
+      return;
     }
 
     return ~~len;
diff --git a/test/application.js b/test/application.js
index dc2a4b1..d3f8e1d 100644
--- a/test/application.js
+++ b/test/application.js
@@ -565,11 +565,83 @@ describe('app.respond', function(){
       .end(function(err, res){
         if (err) return done(err);
         var pkg = require('../package');
+        res.should.not.have.header('Content-Length');
         res.body.should.eql(pkg);
         done();
       });
     })
 
+    it('should strip content-length when overwriting', function(done){
+      var app = koa();
+
+      app.use(function *(){
+        this.body = 'hello';
+        this.body = fs.createReadStream('package.json');
+        this.set('Content-Type', 'application/json');
+      });
+
+      var server = app.listen();
+
+      request(server)
+      .get('/')
+      .expect('Content-Type', 'application/json')
+      .end(function(err, res){
+        if (err) return done(err);
+        var pkg = require('../package');
+        res.should.not.have.header('Content-Length');
+        res.body.should.eql(pkg);
+        done();
+      })
+    })
+
+    it('should keep content-length if not overwritten', function(done){
+      var app = koa();
+
+      app.use(function *(){
+        this.length = fs.readFileSync('package.json').length;
+        this.body = fs.createReadStream('package.json');
+        this.set('Content-Type', 'application/json');
+      });
+
+      var server = app.listen();
+
+      request(server)
+      .get('/')
+      .expect('Content-Type', 'application/json')
+      .end(function(err, res){
+        if (err) return done(err);
+        var pkg = require('../package');
+        res.should.have.header('Content-Length');
+        res.body.should.eql(pkg);
+        done();
+      })
+    })
+
+    it('should keep content-length if overwritten with the same stream', function(done){
+      var app = koa();
+
+      app.use(function *(){
+        this.length = fs.readFileSync('package.json').length;
+        var stream = fs.createReadStream('package.json');
+        this.body = stream;
+        this.body = stream;
+        this.set('Content-Type', 'application/json');
+      });
+
+      var server = app.listen();
+
+      request(server)
+      .get('/')
+      .expect('Content-Type', 'application/json')
+      .end(function(err, res){
+        if (err) return done(err);
+        var pkg = require('../package');
+        res.should.have.header('Content-Length');
+        res.body.should.eql(pkg);
+        done();
+      })
+    })
+
     it('should handle errors', function(done){
       var app = koa();
 
diff --git a/test/response/length.js b/test/response/length.js
index 53adba1..8fd8456 100644
--- a/test/response/length.js
+++ b/test/response/length.js
@@ -1,4 +1,6 @@
 
+var fs = require('fs');
+var should = require('should');
 var response = require('../context').response;
 var assert = require('assert');
 
@@ -27,10 +29,31 @@ describe('res.length', function(){
         var res = response();
 
         res.body = 'foo';
+        res.remove('Content-Length');
         res.length.should.equal(3);
 
-        res.body = new Buffer('foo');
+        res.body = 'foo';
         res.length.should.equal(3);
+
+        res.body = new Buffer('foo bar');
+        res.remove('Content-Length');
+        res.length.should.equal(7);
+
+        res.body = new Buffer('foo bar');
+        res.length.should.equal(7);
+
+        res.body = { hello: 'world' };
+        res.remove('Content-Length');
+        res.length.should.equal(17);
+
+        res.body = { hello: 'world' };
+        res.length.should.equal(17);
+
+        res.body = fs.createReadStream('package.json');
+        should.not.exist(res.length);
+
+        res.body = null;
+        should.not.exist(res.length);
       })
     })
 
@@ -41,4 +64,4 @@ describe('res.length', function(){
       })
     })
   })
-})
\ No newline at end of file
+})