Fix res.end patch to call correct upstream res.write

fixes #72

Author: dougwilson

History.md

@@ -1,3 +1,8 @@
+unreleased
+==========
+
+  * Fix `res.end` patch to call correct upstream `res.write`
+
 1.7.2 / 2014-07-27
 ==================
 

index.js

@@ -179,9 +179,10 @@ function session(options){
     });
 
     // proxy end() to commit the session
-    var end = res.end;
+    var _end = res.end;
+    var _write = res.write;
     var ended = false;
-    res.end = function(chunk, encoding){
+    res.end = function end(chunk, encoding) {
       if (ended) {
         return false;
       }
@@ -206,16 +207,16 @@ function session(options){
           debug('destroyed');
 
           if (sync) {
-            ret = end.call(res, chunk, encoding);
+            ret = _end.call(res, chunk, encoding);
             sync = false;
             return;
           }
 
-          end.call(res);
+          _end.call(res);
         });
 
         if (sync) {
-          ret = res.write(chunk, encoding);
+          ret = _write.call(res, chunk, encoding);
           sync = false;
         }
 
@@ -225,7 +226,7 @@ function session(options){
       // no session to save
       if (!req.session) {
         debug('no session');
-        return end.call(res, chunk, encoding);
+        return _end.call(res, chunk, encoding);
       }
 
       req.session.resetMaxAge();
@@ -240,23 +241,23 @@ function session(options){
           debug('saved');
 
           if (sync) {
-            ret = end.call(res, chunk, encoding);
+            ret = _end.call(res, chunk, encoding);
             sync = false;
             return;
           }
 
-          end.call(res);
+          _end.call(res);
         });
 
         if (sync) {
-          ret = res.write(chunk, encoding);
+          ret = _write.call(res, chunk, encoding);
           sync = false;
         }
 
         return ret;
       }
 
-      return end.call(res, chunk, encoding);
+      return _end.call(res, chunk, encoding);
     };
 
     // generate the session

test/session.js

@@ -937,6 +937,40 @@ describe('session()', function(){
     })
   });
 
+  describe('res.end patch', function () {
+    it('should correctly handle res.end/res.write patched prior', function (done) {
+      var app = express()
+
+      app.use(writePatch())
+      app.use(createSession())
+      app.use(function (req, res) {
+        req.session.hit = true
+        res.write('hello, ')
+        res.end('world')
+      })
+
+      request(app)
+      .get('/')
+      .expect(200, 'hello, world', done)
+    })
+
+    it('should correctly handle res.end/res.write patched after', function (done) {
+      var app = express()
+
+      app.use(createSession())
+      app.use(writePatch())
+      app.use(function (req, res) {
+        req.session.hit = true
+        res.write('hello, ')
+        res.end('world')
+      })
+
+      request(app)
+      .get('/')
+      .expect(200, 'hello, world', done)
+    })
+  })
+
   describe('req.session', function(){
     it('should persist', function(done){
       var store = new session.MemoryStore()
@@ -1584,19 +1618,9 @@ function cookie(res) {
 }
 
 function createServer(opts, fn) {
-  var options = opts || {}
+  var _session = createSession(opts)
   var respond = fn || end
 
-  if (!('cookie' in options)) {
-    options.cookie = { maxAge: 60 * 1000 }
-  }
-
-  if (!('secret' in options)) {
-    options.secret = 'keyboard cat'
-  }
-
-  var _session = session(options)
-
   var server = http.createServer(function (req, res) {
     _session(req, res, function (err) {
       if (err && !res._header) {
@@ -1617,6 +1641,20 @@ function createServer(opts, fn) {
   return server
 }
 
+function createSession(opts) {
+  var options = opts || {}
+
+  if (!('cookie' in options)) {
+    options.cookie = { maxAge: 60 * 1000 }
+  }
+
+  if (!('secret' in options)) {
+    options.secret = 'keyboard cat'
+  }
+
+  return session(options)
+}
+
 function end(req, res) {
   res.end()
 }
@@ -1632,6 +1670,29 @@ function sid(res) {
   return val
 }
 
+function writePatch() {
+  var ended = false
+  return function addWritePatch(req, res, next) {
+    var _end = res.end
+    var _write = res.write
+
+    res.end = function end() {
+      ended = true
+      return _end.apply(this, arguments)
+    }
+
+    res.write = function write() {
+      if (ended) {
+        throw new Error('write after end')
+      }
+
+      return _write.apply(this, arguments)
+    }
+
+    next()
+  }
+}
+
 function SyncStore() {
   this.sessions = Object.create(null);
 }