Add support for custom primary keys

lib/rack/oauth2/models/access_token.rb

@@ -55,7 +55,7 @@ def from_identity(identity)
           # Returns all access tokens for a given client, Use limit and offset
           # to return a subset of tokens, sorted by creation date.
           def for_client(client_id, offset = 0, limit = 100)
-            client_id = BSON::ObjectId(client_id.to_s)
+            client_id = Server.options[:pk_generator].from_string(client_id.to_s)
             collection.find({ :client_id=>client_id }, { :sort=>[[:created_at, Mongo::ASCENDING]], :skip=>offset, :limit=>limit }).
               map { |token| Server.new_instance self, token }
           end
@@ -75,7 +75,7 @@ def count(filter = {})
             elsif filter.has_key?(:revoked)
               select[:revoked] = filter[:revoked] ? { :$ne=>nil } : { :$eq=>nil }
             end
-            select[:client_id] = BSON::ObjectId(filter[:client_id].to_s) if filter[:client_id]
+            select[:client_id] = Server.options[:pk_generator].from_string(filter[:client_id].to_s) if filter[:client_id]
             collection.find(select).count
           end
 
@@ -84,7 +84,7 @@ def historical(filter = {})
             select = { :$gt=> { :created_at=>Time.now - 86400 * days } }
             select = {}
             if filter[:client_id]
-              select[:client_id] = BSON::ObjectId(filter[:client_id].to_s)
+              select[:client_id] = Server.options[:pk_generator].from_string(filter[:client_id].to_s)
             end
             raw = Server::AccessToken.collection.group("function (token) { return { ts: Math.floor(token.created_at / 86400) } }",
               select, { :granted=>0 }, "function (token, state) { state.granted++ }")

lib/rack/oauth2/models/auth_request.rb

@@ -9,17 +9,18 @@ class AuthRequest
         class << self
           # Find AuthRequest from identifier.
           def find(request_id)
-            id = BSON::ObjectId(request_id.to_s)
-            Server.new_instance self, collection.find_one(id)
-          rescue BSON::InvalidObjectId
+            id = Server.options[:pk_generator].from_string(request_id.to_s)
+            Server.new_instance self, collection.find_one({ :_id=>id })
+          rescue BSON::InvalidObjectId, InvalidUUID
           end
 
           # Create a new authorization request. This holds state, so in addition
           # to client ID and scope, we need to know the URL to redirect back to
           # and any state value to pass back in that redirect.
           def create(client, scope, redirect_uri, response_type, state)
             scope = Utils.normalize_scope(scope) & client.scope # Only allowed scope
-            fields = { :client_id=>client.id, :scope=>scope, :redirect_uri=>client.redirect_uri || redirect_uri,
+            fields = { :_id=>Server.options[:pk_generator].generate, :client_id=>client.id,
+                       :scope=>scope, :redirect_uri=>client.redirect_uri || redirect_uri,
                        :response_type=>response_type, :state=>state,
                        :grant_code=>nil, :authorized_at=>nil,
                        :created_at=>Time.now.to_i, :revoked=>nil }

lib/rack/oauth2/models/client.rb

@@ -8,9 +8,9 @@ class << self
           # Authenticate a client request. This method takes three arguments,
           # Find Client from client identifier.
           def find(client_id)
-            id = BSON::ObjectId(client_id.to_s)
-            Server.new_instance self, collection.find_one(id)
-          rescue BSON::InvalidObjectId
+            id = Server.options[:pk_generator].from_string(client_id.to_s)
+            Server.new_instance self, collection.find_one({ :_id=>id })
+          rescue BSON::InvalidObjectId, InvalidUUID
           end
 
           # Create a new client. Client provides the following properties:
@@ -33,20 +33,21 @@ def create(args)
                         :notes=>args[:notes].to_s, :scope=>scope,
                         :created_at=>Time.now.to_i, :revoked=>nil }
             if args[:id] && args[:secret]
-              fields[:_id], fields[:secret] = BSON::ObjectId(args[:id].to_s), args[:secret]
+              fields[:_id], fields[:secret] = Server.options[:pk_generator].from_string(args[:id].to_s), args[:secret]
               collection.insert(fields, :safe=>true)
             else
               fields[:secret] = Server.secure_random
-              fields[:_id] = collection.insert(fields)
+              fields[:_id] = Server.options[:pk_generator].generate
+              collection.insert(fields)
             end
             Server.new_instance self, fields
           end
 
           # Lookup client by ID, display name or URL.
           def lookup(field)
-            id = BSON::ObjectId(field.to_s)
+            id = Server.options[:pk_generator].from_string(field.to_s)
             Server.new_instance self, collection.find_one(id)
-          rescue BSON::InvalidObjectId
+          rescue BSON::InvalidObjectId, InvalidUUID
             Server.new_instance self, collection.find_one({ :display_name=>field }) || collection.find_one({ :link=>field })
           end
 
@@ -58,7 +59,7 @@ def all
 
           # Deletes client with given identifier (also, all related records).
           def delete(client_id)
-            id = BSON::ObjectId(client_id.to_s)
+            id = Server.options[:pk_generator].from_string(client_id.to_s)
             Client.collection.remove({ :_id=>id })
             AuthRequest.collection.remove({ :client_id=>id })
             AccessGrant.collection.remove({ :client_id=>id })

lib/rack/oauth2/pk_generators.rb

@@ -0,0 +1,2 @@
+require 'rack/oauth2/pk_generators/bson_generator'
+require 'rack/oauth2/pk_generators/uuid_generator'

lib/rack/oauth2/pk_generators/bson_generator.rb

@@ -0,0 +1,15 @@
+module Rack
+  module OAuth2
+    class BSONGenerator
+      class << self
+        def from_string(value)
+          BSON::ObjectId(value)
+        end
+
+        def generate
+          BSON::ObjectId.new
+        end
+      end
+    end
+  end
+end

lib/rack/oauth2/pk_generators/uuid_generator.rb

@@ -0,0 +1,35 @@
+module Rack
+  module OAuth2
+    # The uuid string is not valid
+    class UnknownUUIDLibrary < Rack::OAuth2::Server::OAuthError
+      def initialize
+        super :unknown_uuid_library, "You have requested UUID Primary Key but the UUID library is not present."
+      end
+    end
+
+    # The uuid string is not valid
+    class InvalidUUID < Rack::OAuth2::Server::OAuthError
+      def initialize
+        super :invalid_uuid, "The uuid string you provided is invalid."
+      end
+    end
+
+    class UUIDGenerator
+
+      class << self
+        def generate
+          fail UnknownUUIDLibrary unless Object.const_defined?('UUID')
+
+          UUID.generate
+        end
+
+        def from_string(value)
+          fail InvalidUUID unless UUID.validate(value)
+          fail UnknownUUIDLibrary unless Object.const_defined?('UUID')
+
+          value
+        end
+      end
+    end
+  end
+end

lib/rack/oauth2/server.rb

@@ -3,6 +3,7 @@
 require "rack/oauth2/server/errors"
 require "rack/oauth2/server/utils"
 require "rack/oauth2/server/helper"
+require "rack/oauth2/pk_generators"
 require "iconv"
 require "json"
 
@@ -204,7 +205,7 @@ def get_issuer(identifier)
       #
       Options = Struct.new(:access_token_path, :authenticator, :assertion_handler, :authorization_types,
         :authorize_path, :database, :host, :param_authentication, :path, :realm, 
-        :expires_in,:logger, :collection_prefix)
+        :expires_in,:logger, :collection_prefix, :pk_generator)
 
       # Global options. This is what we set during configuration (e.g. Rails'
       # config/application), and options all handlers inherit by default.
@@ -214,6 +215,9 @@ def self.options
 
       @options = Options.new
 
+      # default to BSON::ObjectId
+      @options.pk_generator = BSONGenerator
+
       def initialize(app, options = nil, &authenticator)
         @app = app
         @options = options || Server.options
@@ -224,6 +228,7 @@ def initialize(app, options = nil, &authenticator)
         @options.authorization_types ||=  %w{code token}
         @options.param_authentication ||= false
         @options.collection_prefix ||= "oauth2"
+        @options.pk_generator ||= BSONGenerator
       end
 
       # Options specific for this handle. @see Options

test/oauth/server_methods_test.rb

@@ -19,6 +19,10 @@ def setup
     should "set oauth.collection_prefix" do
       assert_equal "oauth2_prefix", Server.options.collection_prefix
     end
+
+    should "set oauth.pk_generator" do
+      assert_equal BSONGenerator, Server.options.pk_generator
+    end
   end
 
   context "get_auth_request" do

test/oauth/server_test.rb

@@ -9,6 +9,17 @@ class ServerTest < Test::Unit::TestCase
       assert_equal "UberClient", @client.display_name
       assert_same_elements %w(read write oauth-admin), @client.scope
     end
-  end
 
+    context "uuid primary key" do
+      setup do
+        options = Server::Options.new
+        options.pk_generator = UUIDGenerator
+        @server = Server.new({}, options)
+      end
+
+      should "set oauth.pk_generator" do
+        assert_equal UUIDGenerator, @server.options.pk_generator
+      end
+    end
+  end
 end