Proxy res.end(callback) calls correctly

index.js

@@ -247,7 +247,7 @@ function session(options) {
     var _end = res.end;
     var _write = res.write;
     var ended = false;
-    res.end = function end(chunk, encoding) {
+    res.end = function end(chunk, encoding, callback) {
       if (ended) {
         return false;
       }
@@ -257,14 +257,42 @@ function session(options) {
       var ret;
       var sync = true;
 
+      // Keep a reference to the original arguments so we can preserve their
+      // order when passing them on to the original _end()
+      var endArgs = arguments;
+
+      // The callback argument is optional. If provided, it may be the
+      // first, second, or third argument, and must be the last argument.
+      var callbackArgumentIndex = typeof chunk === 'function'
+        ? 0
+        : typeof encoding === 'function'
+        ? 1
+        : 2;
+
+      if (callbackArgumentIndex === 0) {
+        callback = chunk;
+        chunk = null;
+        encoding = null;
+      } else if (callbackArgumentIndex === 1) {
+        callback = encoding;
+        encoding = null;
+      }
+
       function writeend() {
         if (sync) {
-          ret = _end.call(res, chunk, encoding);
+          ret = _end.apply(res, endArgs);
           sync = false;
           return;
         }
 
-        _end.call(res);
+        // This differs very slightly from (undocumented) Node behaviour in versions
+        // prior to 0.11.6 (which introduced callback support), because we explicitly
+        // call write() followed by end().
+        // The behaviour seems unintentional, as the callback is invoked immediately
+        // after write(), as opposed to after end().
+        var argumentsWithoutChunkOrEncoding = [null, null, null];
+        argumentsWithoutChunkOrEncoding[callbackArgumentIndex] = callback;
+        _end.apply(res, argumentsWithoutChunkOrEncoding);
       }
 
       function writetop() {
@@ -284,12 +312,17 @@ function session(options) {
           chunk = !Buffer.isBuffer(chunk)
             ? Buffer.from(chunk, encoding)
             : chunk;
+          endArgs[0] = chunk;
           encoding = undefined;
+          if (callbackArgumentIndex !== 1) {
+            endArgs[1] = encoding;
+          }
 
           if (chunk.length !== 0) {
             debug('split response');
             ret = _write.call(res, chunk.slice(0, chunk.length - 1));
             chunk = chunk.slice(chunk.length - 1, chunk.length);
+            endArgs[0] = chunk;
             return ret;
           }
         }
@@ -309,7 +342,11 @@ function session(options) {
           }
 
           debug('destroyed');
-          writeend();
+          try {
+            writeend();
+          } catch (err) {
+            next(err);
+          }
         });
 
         return writetop();
@@ -318,7 +355,7 @@ function session(options) {
       // no session to save
       if (!req.session) {
         debug('no session');
-        return _end.call(res, chunk, encoding);
+        return _end.apply(res, endArgs);
       }
 
       if (!touched) {
@@ -333,7 +370,11 @@ function session(options) {
             defer(next, err);
           }
 
-          writeend();
+          try {
+            writeend();
+          } catch (err) {
+            next(err);
+          }
         });
 
         return writetop();
@@ -346,13 +387,17 @@ function session(options) {
           }
 
           debug('touched');
-          writeend();
+          try {
+            writeend();
+          } catch (err) {
+            next(err);
+          }
         });
 
         return writetop();
       }
 
-      return _end.call(res, chunk, encoding);
+      return _end.apply(res, endArgs);
     };
 
     // generate the session

test/session.js

@@ -189,6 +189,94 @@ describe('session()', function(){
     .expect(200, 'Hello, world!', done);
   })
 
+  describe('res.end() proxy', function () {
+    var nodeVersion = utils.getNodeVersion();
+
+    // Support for the callback argument was added in Node 0.11.6.
+    var nodeVersionSupportsResEndCallback = nodeVersion.major > 0 || nodeVersion.minor > 11 || (nodeVersion.minor === 11 && nodeVersion.patch >= 6);
+
+    it('should correctly handle callback as only argument', function (done) {
+      var callbackHasBeenCalled = false
+
+      var server = createServer(null, function (req, res) {
+        function callback() {
+          callbackHasBeenCalled = true
+        }
+
+        server.on('error', function onerror (err) {
+          assert.ok(err)
+          assert.strictEqual(err.message, 'first argument must be a string or Buffer')
+          done()
+        })
+
+        res.end(callback)
+      });
+
+      request(server)
+      .get('/')
+      .expect(200, '', function () {
+        if (nodeVersionSupportsResEndCallback) {
+          // If the Node version does not support the callback,
+          // the error listener will make the necessary assertions.
+          assert.ok(callbackHasBeenCalled)
+          done()
+        }
+      })
+    });
+
+    it('should correctly handle callback as second argument', function (done) {
+      var callbackHasBeenCalled = false
+
+      var server = createServer(null, function (req, res) {
+        function callback() {
+          callbackHasBeenCalled = true
+        }
+
+        res.end('hello', callback)
+      });
+
+      request(server).get('/').expect(200, 'hello', function () {
+        if (nodeVersionSupportsResEndCallback) {
+          assert.ok(callbackHasBeenCalled)
+        } else {
+          // This is (very slightly) different from the default Node behaviour,
+          // where the callback is invoked if `chunk` is a non-empty string or Buffer.
+          // This behaviour is undocumented, and seemingly happens unintentionally
+          // when the chunk and encoding are passed to write(), and in turn to
+          // afterWrite(), which considers the encoding argument to possibly be
+          // a function, and invokes it.
+          assert.strictEqual(callbackHasBeenCalled, false)
+        }
+
+        done()
+      })
+    })
+
+    it('should correctly handle callback as third argument', function (done) {
+      var callbackHasBeenCalled = false
+
+      var server = createServer(null, function (req, res) {
+        function callback() {
+          callbackHasBeenCalled = true
+        }
+
+        res.end('hello', 'utf8', callback)
+      });
+
+      request(server).get('/').expect(200, 'hello', function () {
+        // The third argument is completely ignored on Node versions
+        // that do not support the callback.
+        if (nodeVersionSupportsResEndCallback) {
+          assert.ok(callbackHasBeenCalled)
+        } else {
+          assert.strictEqual(callbackHasBeenCalled, false)
+        }
+
+        done()
+      })
+    })
+  })
+
   it('should handle res.end(null) calls', function (done) {
     var server = createServer(null, function (req, res) {
       res.end(null)

test/support/utils.js

@@ -2,6 +2,7 @@
 
 module.exports.parseSetCookie = parseSetCookie
 module.exports.writePatch = writePatch
+module.exports.getNodeVersion = getNodeVersion
 
 function parseSetCookie (header) {
   var match
@@ -40,3 +41,13 @@ function writePatch (res) {
     return _write.apply(this, arguments)
   }
 }
+
+function getNodeVersion () {
+  var nodeVersionNumbers = process.versions.node.split('.').map(Number);
+
+  return {
+    major: nodeVersionNumbers[0],
+    minor: nodeVersionNumbers[1],
+    patch: nodeVersionNumbers[2]
+  }
+}