From 3d157007f6bec653f524a4df5e4e69d825df9765 Mon Sep 17 00:00:00 2001
From: Steven Bazyl <sqrrrl@gmail.com>
Date: Tue, 30 Oct 2012 13:18:12 -0700
Subject: [PATCH] Auto-refresh OAuth 2 tokens & retry request on 401 response

---
 CHANGELOG.md                                  |  2 ++
 lib/google/api_client.rb                      |  9 +++++
 lib/google/api_client/auth/jwt_asserter.rb    | 17 +++++++++-
 .../google/api_client/service_account_spec.rb | 34 +++++++++++++++++++
 4 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fef21b2b4..31bba593a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@
     * :api_method in request can no longer be a string
     * Deprecated ResumableUpload.send_* methods.
 * Reduce memory utilization when uploading large files
+* Automatic refresh of OAuth 2 credentials & retry of request when 401 errors
+  are returned
 * Simplify internal request processing.
 
 # 0.4.7
diff --git a/lib/google/api_client.rb b/lib/google/api_client.rb
index 937e89ee1..28b027ef7 100644
--- a/lib/google/api_client.rb
+++ b/lib/google/api_client.rb
@@ -540,6 +540,15 @@ module Google
       request.authorization = options[:authorization] || self.authorization unless options[:authenticated] == false
 
       result = request.send(connection)
+      if result.status == 401 && authorization.respond_to?(:refresh_token) 
+        begin
+          authorization.fetch_access_token!
+          result = request.send(connection)
+        rescue Signet::AuthorizationError
+           # Ignore since we want the original error
+        end
+      end
+      
       return result
     end
 
diff --git a/lib/google/api_client/auth/jwt_asserter.rb b/lib/google/api_client/auth/jwt_asserter.rb
index 62cfd8689..a974f8af6 100644
--- a/lib/google/api_client/auth/jwt_asserter.rb
+++ b/lib/google/api_client/auth/jwt_asserter.rb
@@ -14,6 +14,7 @@
 
 require 'jwt'
 require 'signet/oauth_2/client'
+require 'delegate'
 
 module Google
   class APIClient
@@ -117,7 +118,21 @@ module Google
         authorization.grant_type = 'urn:ietf:params:oauth:grant-type:jwt-bearer'
         authorization.extension_parameters = { :assertion => assertion }
         authorization.fetch_access_token!(options)
-        return authorization
+        return JWTAuthorization.new(authorization, self, person)
+      end
+    end
+    
+    class JWTAuthorization < DelegateClass(Signet::OAuth2::Client)
+      def initialize(authorization, asserter, person = nil)
+        @asserter = asserter
+        @person = person
+        super(authorization)
+      end
+      
+      def fetch_access_token!(options={})
+        new_authorization = @asserter.authorize(@person, options)
+        __setobj__(new_authorization)
+        self
       end
     end
   end
diff --git a/spec/google/api_client/service_account_spec.rb b/spec/google/api_client/service_account_spec.rb
index f303d6466..2a2b38cfa 100644
--- a/spec/google/api_client/service_account_spec.rb
+++ b/spec/google/api_client/service_account_spec.rb
@@ -52,5 +52,39 @@ describe Google::APIClient::JWTAsserter do
     auth.access_token.should == "1/abcdef1234567890"
     conn.verify
   end
+  
+  it 'should be refreshable' do
+    conn = stub_connection do |stub|
+      stub.post('/o/oauth2/token') do |env|
+        params = Addressable::URI.form_unencode(env[:body])
+        JWT.decode(params.assoc("assertion").last, @key.public_key)
+        params.assoc("grant_type").should == ['grant_type','urn:ietf:params:oauth:grant-type:jwt-bearer']
+        [200, {}, '{
+          "access_token" : "1/abcdef1234567890",
+          "token_type" : "Bearer",
+          "expires_in" : 3600
+        }']
+      end
+      stub.post('/o/oauth2/token') do |env|
+        params = Addressable::URI.form_unencode(env[:body])
+        JWT.decode(params.assoc("assertion").last, @key.public_key)
+        params.assoc("grant_type").should == ['grant_type','urn:ietf:params:oauth:grant-type:jwt-bearer']
+        [200, {}, '{
+          "access_token" : "1/0987654321fedcba",
+          "token_type" : "Bearer",
+          "expires_in" : 3600
+        }']
+      end
+    end
+    asserter = Google::APIClient::JWTAsserter.new('client1', 'scope1 scope2', @key)
+    auth = asserter.authorize(nil, { :connection => conn })
+    auth.should_not == nil?
+    auth.access_token.should == "1/abcdef1234567890"
+    
+    auth.fetch_access_token!(:connection => conn)
+    auth.access_token.should == "1/0987654321fedcba"
+    
+    conn.verify
+  end    
 end