Merge pull request #241 from stephenplusplus/spp--token-retry

connection: refetch token if invalid.
googleapis · Sep 30, 2014 · a525e86 · a525e86
2 parents d370c7f + 797ca1b
commit a525e86
Show file tree

Hide file tree

Showing 2 changed files with 77 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 refreshing a token. */
+var MAX_TOKEN_REFRESH_ATTEMPTS = 1;
+
 /** @const {object} gcloud-node's package.json file. */
 var PKG = require('../../package.json');
 
@@ -217,7 +221,9 @@ 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_TOKEN_REFRESH_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 +233,25 @@ Connection.prototype.fetchServiceAccountToken_ = function(callback) {
  */
 Connection.prototype.req = function(requestOptions, callback) {
   var that = this;
+  var tokenRefreshAttempts = 0;
   callback = callback || util.noop;
-  this.createAuthorizedReq(requestOptions, function(err, authorizedReq) {
+  function onAuthorizedReq(err, authorizedReq) {
     if (err) {
       callback(err);
       return;
     }
-    that.requester(authorizedReq, callback);
-  });
+    that.requester(authorizedReq, function(err) {
+      if (err && err.code === 401 &&
+          ++tokenRefreshAttempts <= MAX_TOKEN_REFRESH_ATTEMPTS) {
+        // Invalid token. Try to fetch a new one.
+        that.token = null;
+        that.createAuthorizedReq(requestOptions, onAuthorizedReq);
+        return;
+      }
+      callback.apply(null, util.toArray(arguments));
+    });
+  }
+  this.createAuthorizedReq(requestOptions, onAuthorizedReq);
 };
 
 /**
@@ -246,10 +263,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 +322,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
@@ -23,6 +23,7 @@ var async = require('async');
 var path = require('path');
 
 var connection = require('../../lib/common/connection.js');
+var util = require('../../lib/common/util.js');
 
 describe('Connection', function() {
   var conn;
@@ -141,13 +142,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(util.toArray(arguments), args);
         done();
       });
     });