Persist cookies across redirects when follow_set_cookies is true

lib/cookies.js

@@ -48,6 +48,7 @@ function parseSetCookieString(str) {
 // Parses a set-cookie-header and returns a key/value object.
 // Each key represents the name of a cookie.
 function parseSetCookieHeader(header) {
+  if (!header) return {};
   header = (header instanceof Array) ? header : [header];
 
   return header.reduce(function(res, str) {

lib/needle.js

@@ -33,6 +33,9 @@ var tls_options = 'agent pfx key passphrase cert ca ciphers rejectUnauthorized s
 // we'll default new requests to set a Connection: close header.
 var close_by_default = !http.Agent || http.Agent.defaultMaxSockets != Infinity;
 
+// see if we have Object.assign. otherwise fall back to util._extend
+var extend = Object.assign ? Object.assign : require('util')._extend;
+
 //////////////////////////////////////////
 // decompressors for gzip/deflate bodies
 
@@ -414,8 +417,10 @@ Needle.prototype.send_request = function(count, method, uri, config, post_data,
     if (timer) clearTimeout(timer);
     set_timeout(config.read_timeout);
 
-    if (headers['set-cookie'] && config.parse_cookies) {
-      resp.cookies = cookies.read(headers['set-cookie']);
+    // if we got cookies, parse them unless we were instructed not to. make sure to include any
+    // cookies that might have been set on previous redirects.
+    if (config.parse_cookies && (headers['set-cookie'] || config.stored_cookies)) {
+      resp.cookies = extend(config.stored_cookies || {}, cookies.read(headers['set-cookie']));
       debug('Got cookies', resp.cookies);
     }
 
@@ -432,8 +437,11 @@ Needle.prototype.send_request = function(count, method, uri, config, post_data,
           delete config.headers['content-length']; // in case the original was a multipart POST request.
         }
 
-        if (config.follow_set_cookies && resp.cookies)
+        // if follow_set_cookies is true, make sure to put any cookies in the next request's headers.
+        if (config.follow_set_cookies && resp.cookies) {
+          config.stored_cookies    = resp.cookies;
           config.headers['cookie'] = cookies.write(resp.cookies);
+        }
 
         if (config.follow_set_referer)
           config.headers['referer'] = uri; // the original, not the destination URL.

test/cookies_spec.js

@@ -1,8 +1,9 @@
-var needle = require('../'),
-  sinon = require('sinon'),
-  http = require('http'),
-  should = require('should'),
-  assert = require('assert');
+var needle  = require('../'),
+    cookies = require('../lib/cookies'),
+    sinon   = require('sinon'),
+    http    = require('http'),
+    should  = require('should'),
+    assert  = require('assert');
 
 var WEIRD_COOKIE_NAME = 'wc',
   BASE64_COOKIE_NAME = 'bc',
@@ -19,22 +20,23 @@ var WEIRD_COOKIE_NAME = 'wc',
   NUMBER_COOKIE = 'nc=12354342',
   COOKIE_HEADER = WEIRD_COOKIE + '; ' + BASE64_COOKIE + '; ' +
   FORBIDDEN_COOKIE + '; ' + NUMBER_COOKIE,
-  TEST_HOST = 'localhost';
-NO_COOKIES_TEST_PORT = 11112, ALL_COOKIES_TEST_PORT = 11113;
-
-function decode(str) {
-  return decodeURIComponent(str);
-}
-
-function encode(str) {
-  str = str.toString().replace(/[\x00-\x1F\x7F]/g, encodeURIComponent);
-  return str.replace(/[\s\"\,;\\%]/g, encodeURIComponent);
-}
+  TEST_HOST = 'localhost',
+  NO_COOKIES_TEST_PORT = 11112,
+  ALL_COOKIES_TEST_PORT = 11113;
 
 describe('cookies', function() {
 
   var headers, server, opts;
 
+  function decode(str) {
+    return decodeURIComponent(str);
+  }
+
+  function encode(str) {
+    str = str.toString().replace(/[\x00-\x1F\x7F]/g, encodeURIComponent);
+    return str.replace(/[\s\"\,;\\%]/g, encodeURIComponent);
+  }
+
   before(function() {
     setCookieHeader = [
       WEIRD_COOKIE_NAME + '=' + encode(WEIRD_COOKIE_VALUE) + ';',
@@ -120,10 +122,79 @@ describe('cookies', function() {
     });
 
     describe('and response is a redirect', function() {
+
+      var redirectServer, testPort = 22222;
+
+      var responseCookies = [
+        [ // first req
+          WEIRD_COOKIE_NAME + '=' + encode(WEIRD_COOKIE_VALUE) + ';',
+          BASE64_COOKIE_NAME + '=' + encode(BASE64_COOKIE_VALUE) + ';',
+          'FOO=123;'
+        ], [ // second req
+          FORBIDDEN_COOKIE_NAME + '=' + encode(FORBIDDEN_COOKIE_VALUE) + ';',
+          NUMBER_COOKIE_NAME + '=' + encode(NUMBER_COOKIE_VALUE) + ';'
+        ], [ // third red
+          'FOO=BAR;'
+        ]
+      ]
+
+      before(function() {
+        redirectServer = http.createServer(function(req, res) {
+          var number  = parseInt(req.url.replace('/', ''));
+          var nextUrl = 'http://' + TEST_HOST + ':' + testPort + '/' + (number + 1);
+
+          if (responseCookies[number]) { // got cookies
+            res.statusCode = 302;
+            res.setHeader('Set-Cookie', responseCookies[number]);
+            res.setHeader('Location', nextUrl);
+          } else if (number == 3) {
+            res.statusCode = 302; // redirect but without cookies
+            res.setHeader('Location', nextUrl);
+          }
+
+          res.end('OK');
+        }).listen(22222, TEST_HOST);
+      });
+
+      after(function(done) {
+        redirectServer.close(done);
+      })
+
       describe('and follow_set_cookies is false', function() {
-        it('no cookie header set on redirection request', function() {});
+
+        var opts = {
+          follow_set_cookies: false,
+          follow_max: 4
+        };
+
+        it('no cookie header set on redirection request', function(done) {
+          var spy = sinon.spy(cookies, 'write');
+
+          needle.get(TEST_HOST + ':' + testPort + '/0', opts, function(err, resp) {
+            spy.callCount.should.eql(0);
+            done();
+          });
+        });
+      });
+
+      describe('and follow_set_cookies is true', function() {
+        var opts = {
+          follow_set_cookies: true,
+          follow_max: 4
+        };
+
+        it('should have all the cookies', function(done) {
+          needle.get(TEST_HOST + ':' + testPort + '/0', opts, function(err, resp) {
+            resp.cookies.should.have.property(WEIRD_COOKIE_NAME);
+            resp.cookies.should.have.property(BASE64_COOKIE_NAME);
+            resp.cookies.should.have.property(FORBIDDEN_COOKIE_NAME);
+            resp.cookies.should.have.property(NUMBER_COOKIE_NAME);
+            resp.cookies.should.have.property('FOO');
+            resp.cookies.FOO.should.eql('BAR'); // should overwrite previous one
+            done();
+          });
+        });
       });
-      describe('and follow_set_cookies is true', function() {});
     });
 
     describe('with parse_cookies = false', function() {
@@ -211,4 +282,4 @@ describe('cookies', function() {
       });
     });
   });
-});
+});