connection: refetch token if invalid.

googleapis · Sep 30, 2014 · a8f31ff · a8f31ff
1 parent 4e6f384
commit a8f31ff
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 16 deletions.
diff --git a/lib/common/connection.js b/lib/common/connection.js
@@ -22,6 +22,7 @@
 'use strict';
 
 var events = require('events');
+var extend = require('extend');
 var fs = require('fs');
 var GAPIToken = require('gapitoken');
 var nodeutil = require('util');
@@ -35,6 +36,9 @@ var METADATA_TOKEN_URL =
     'http://metadata/computeMetadata/v1/instance/service-accounts/default/' +
     'token';
 
+/** @const {number} Maximum amount of times to attempt an API request. */
+var MAX_ATTEMPTS = 2;
+
 /** @const {object} gcloud-node's package.json file. */
 var PKG = require('../../package.json');
 
@@ -217,7 +221,8 @@ Connection.prototype.fetchServiceAccountToken_ = function(callback) {
 
 /**
  * Make an authorized request if the current connection token is still valid. If
- * it's not, try to reconnect.
+ * it's not, try to reconnect to the limit specified by MAX_ATTEMPTS. If a valid
+ * connection still cannot be made, execute the callback with the API error.
  *
  * @param {object} requestOptions - Request options.
  * @param {function=} callback - The callback function.
@@ -227,14 +232,24 @@ Connection.prototype.fetchServiceAccountToken_ = function(callback) {
  */
 Connection.prototype.req = function(requestOptions, callback) {
   var that = this;
+  var attempts = 1;
   callback = callback || util.noop;
-  this.createAuthorizedReq(requestOptions, function(err, authorizedReq) {
+  function onAuthorized(err, authorizedReq) {
     if (err) {
       callback(err);
       return;
     }
-    that.requester(authorizedReq, callback);
-  });
+    that.requester(authorizedReq, function(err) {
+      if (err && err.code === 401 && ++attempts <= MAX_ATTEMPTS) {
+        // Invalid token. Try to fetch a new one.
+        that.token = null;
+        that.createAuthorizedReq(requestOptions, onAuthorized);
+        return;
+      }
+      callback.apply(null, util.toArray(arguments));
+    });
+  }
+  this.createAuthorizedReq(requestOptions, onAuthorized);
 };
 
 /**
@@ -246,10 +261,10 @@ Connection.prototype.req = function(requestOptions, callback) {
  * @example
  * conn.createAuthorizedReq({}, function(err) {});
  */
-Connection.prototype.createAuthorizedReq = function(reqOpts, callback) {
+Connection.prototype.createAuthorizedReq = function(requestOptions, callback) {
   var that = this;
-  // Add user agent.
-  reqOpts.headers = reqOpts.headers || {};
+
+  var reqOpts = extend(true, {}, requestOptions, { headers: {} });
 
   if (reqOpts.headers['User-Agent']) {
     reqOpts.headers['User-Agent'] += '; ' + USER_AGENT;
@@ -305,9 +320,11 @@ Connection.prototype.isConnected = function() {
  * @return {object} Authorized request options.
  */
 Connection.prototype.authorizeReq = function(requestOptions) {
-  requestOptions.headers = requestOptions.headers || {};
-  requestOptions.headers.Authorization = 'Bearer ' + this.token.accessToken;
-  return requestOptions;
+  return extend(true, {}, requestOptions, {
+    headers: {
+      Authorization: 'Bearer ' + this.token.accessToken
+    }
+  });
 };
 
 /**

diff --git a/test/common/connection.js b/test/common/connection.js
@@ -141,13 +141,54 @@ describe('Connection', function() {
       conn.req({ uri: 'https://someuri' }, function() {});
     });
 
-    it('should pass error to callback', function(done) {
-      var error = new Error('Something terrible happened.');
-      conn.fetchToken = function(cb) {
-        cb(error);
+    it('should fetch a new token if API returns a 401', function() {
+      var fetchTokenCount = 0;
+      conn.fetchToken = function(callback) {
+        fetchTokenCount++;
+        callback(null, tokenNeverExpires);
+      };
+      conn.requester = function(req, callback) {
+        if (fetchTokenCount === 1) {
+          callback({ code: 401 });
+        } else {
+          callback(null);
+        }
+      };
+      conn.req({ uri: 'https://someuri' }, function() {});
+      assert.equal(fetchTokenCount, 2);
+    });
+
+    it('should try API request 2 times', function(done) {
+      // Fail 1: invalid token.
+      // -- try to get token --
+      // Fail 2: invalid token.
+      // -- execute callback with error.
+      var error = { code: 401 };
+      var requesterCount = 0;
+      conn.fetchToken = function(callback) {
+        callback(null, tokenNeverExpires);
+      };
+      conn.requester = function(req, callback) {
+        requesterCount++;
+        callback(error);
+      };
+      conn.req({ uri: 'https://someuri' }, function(err) {
+        assert.equal(requesterCount, 2);
+        assert.deepEqual(err, error);
+        done();
+      });
+    });
+
+    it('should pass all arguments from requester to callback', function(done) {
+      var args = [null, 1, 2, 3];
+      conn.fetchToken = function(callback) {
+        callback(null, tokenNeverExpires);
+      };
+      conn.requester = function(req, callback) {
+        callback.apply(null, args);
       };
-      conn.req({}, function(err) {
-        assert.equal(error, err);
+      conn.req({ uri: 'https://someuri' }, function() {
+        assert.deepEqual([].slice.call(arguments), args);
         done();
       });
     });