From ae53038f1cbd4c2478cdddc01596b8c85f6c7ffa Mon Sep 17 00:00:00 2001
From: Paul Covell <paul@done.com>
Date: Thu, 10 May 2012 15:41:12 -0400
Subject: [PATCH] assertion types can provide callbacks

---
 lib/rack/oauth2/server.rb       | 25 +++++++++++++++++++++--
 test/oauth/access_grant_test.rb | 35 +++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/lib/rack/oauth2/server.rb b/lib/rack/oauth2/server.rb
index 7cec605..063bb15 100644
--- a/lib/rack/oauth2/server.rb
+++ b/lib/rack/oauth2/server.rb
@@ -188,7 +188,21 @@ def get_issuer(identifier)
       #     user = User.find_by_username(username)
       #     user if user && user.authenticated?(password)
       #   end
-      Options = Struct.new(:access_token_path, :authenticator, :authorization_types,
+      #
+      # Assertion handler is a hash of blocks keyed by assertion_type.  Blocks receive
+      # three parameters: the client, the assertion, and the scope.  If authenticated, 
+      # it returns an identity.  Otherwise it can return nil or false.  For example:
+      #   oauth.assertion_handler['facebook.com'] = lambda do |client, assertion, scope|
+      #     facebook = URI.parse('https://graph.facebook.com/me?access_token=' + assertion)
+      #     response = Net::HTTP.get_response(facebook)
+      #
+      #     user_data = JSON.parse(response.body)
+      #     user   = User.from_facebook_data(user_data)
+      #   end
+      # Assertion handlers are optional; if one is not present for a given assertion
+      # type, no error will result.
+      #
+      Options = Struct.new(:access_token_path, :authenticator, :assertion_handler, :authorization_types,
         :authorize_path, :database, :host, :param_authentication, :path, :realm, 
         :expires_in,:logger, :collection_prefix)
 
@@ -204,6 +218,7 @@ def initialize(app, options = nil, &authenticator)
         @app = app
         @options = options || Server.options
         @options.authenticator ||= authenticator
+        @options.assertion_handler ||= {}
         @options.access_token_path ||= "/oauth/access_token"
         @options.authorize_path ||= "/oauth/authorize"
         @options.authorization_types ||=  %w{code token}
@@ -422,10 +437,16 @@ def respond_with_access_token(request, logger)
             assertion_type, assertion = request.POST.values_at("assertion_type", "assertion")
             raise InvalidGrantError, "Missing assertion_type/assertion" unless assertion_type && assertion
             # TODO: Add other supported assertion types (i.e. SAML) here
-            raise InvalidGrantError, "Unsupported assertion_type" if assertion_type != "urn:ietf:params:oauth:grant-type:jwt-bearer"
             if assertion_type == "urn:ietf:params:oauth:grant-type:jwt-bearer"
               identity = process_jwt_assertion(assertion)
               access_token = AccessToken.get_token_for(identity, client, requested_scope, options.expires_in)
+            elsif options.assertion_handler[assertion_type]
+              args = [client, assertion, requested_scope]
+              identity = options.assertion_handler[assertion_type].call(*args)
+              raise InvalidGrantError, "Unknown assertion for #{assertion_type}" unless identity
+              access_token = AccessToken.get_token_for(identity, client, requested_scope, options.expires_in)
+            else
+              raise InvalidGrantError, "Unsupported assertion_type" if assertion_type != "urn:ietf:params:oauth:grant-type:jwt-bearer"
             end
           else
             raise UnsupportedGrantType
diff --git a/test/oauth/access_grant_test.rb b/test/oauth/access_grant_test.rb
index c8c8eb8..a06b24d 100644
--- a/test/oauth/access_grant_test.rb
+++ b/test/oauth/access_grant_test.rb
@@ -277,6 +277,41 @@ def request_with_assertion(assertion_type, assertion)
       setup { request_with_assertion "urn:some:assertion:type", nil }
       should_return_error :invalid_grant
     end
+    
+    context "assertion_type with callback" do
+      setup do
+        config.assertion_handler['special_assertion_type'] = lambda do |client, assertion, scope|
+          @client = client
+          @assertion = assertion
+          @scope = scope
+          if assertion == 'myassertion'
+            "Spiderman"
+          else
+            nil
+          end
+        end
+        request_with_assertion 'special_assertion_type', 'myassertion'
+      end
+      
+      context "valid credentials" do
+        setup { request_with_assertion 'special_assertion_type', 'myassertion' }
+        
+        should_respond_with_access_token "read write"
+        should "receive client" do
+          assert_equal client, @client
+        end
+        should "receieve assertion" do
+          assert_equal 'myassertion', @assertion
+        end
+      end
+      
+      context "invalid credentials" do
+        setup { request_with_assertion 'special_assertion_type', 'dunno' }
+        should_return_error :invalid_grant
+      end
+
+      teardown { config.assertion_handler['special_assertion_type'] = nil }
+    end
 
     context "unsupported assertion_type" do
       setup { request_with_assertion "urn:some:assertion:type", "myassertion" }