Additional checks and tests for extension handling.

2 files changed, 53 insertions, 6 deletions

diff --git a/lib/rack/auth/openid.rb b/lib/rack/auth/openid.rb
index b364a207..6225e28b 100644
--- a/lib/rack/auth/openid.rb
+++ b/lib/rack/auth/openid.rb
@@ -5,6 +5,7 @@ require 'rack/auth/abstract/handler' #rack
 require 'uri' #std
 require 'pp' #std
 require 'openid' #gem
+require 'openid/extension' #gem
 require 'openid/store/memory' #gem
 
 module Rack
@@ -340,12 +341,16 @@ module Rack
       #
       # This method will return false if the extension will not be included. Otherwise it will return the key at which the Response will be found in the final session data, which is the namespace uri by default.
       def add_extension ext, *args
-        unless ext.is_a? Module \
-            and ::OpenID::Extension > ext::Request \
-            and ::OpenID::Extension > ext::Response \
-            and ext.constants.include? 'NS_URI'
-          warn 'Incompatible extension provided.'
-          return false
+        if not ext.is_a? Module
+          raise TypeError, "Extension #{ext.inspect} is not a module"
+        elsif not %w'Request Response NS_URI'.all?{|c| ext.constants.include?(c) }
+          raise ArgumentError, "Extension #{ext.inspect} does not containt required constants"
+        elsif not %w'Request Response'.all?{|c| (r=ext.const_get(c)).is_a? Class and ::OpenID::Extension > r }
+          raise TypeError, "Extension #{ext.inspect}'s Request or Response not a decendant of OpenID::Extension"
+        elsif not ext::NS_URI.is_a? String
+          raise TypeError, "Extension #{ext.inspect}'s NS_URI is not a string"
+        elsif not URI(ext::NS_URI).absolute?
+          raise ArgumentError, "Extension #{ext.inspect}'s NS_URI is not a string of an absolute uri"
         end
         @extensions[ext] = args
         return ext::NS_URI
diff --git a/test/spec_rack_auth_openid.rb b/test/spec_rack_auth_openid.rb
index 37d97e19..c39c985e 100644
--- a/test/spec_rack_auth_openid.rb
+++ b/test/spec_rack_auth_openid.rb
@@ -50,4 +50,46 @@ context "Rack::Auth::OpenID" do
     lambda{OID.new(realm, {:return_to => 'http://path/arf/blargh'})}.
       should.not.raise
   end
+
+  specify 'extensions should be a module' do
+    ext = Object.new
+    lambda{OID.new(realm).add_extension(ext)}.should.raise TypeError
+    ext2 = Module.new
+    lambda{OID.new(realm).add_extension(ext2)}.should.raise ArgumentError
+  end
+
+  specify 'extensions should have required constants defined' do
+    ext = Module.new
+    lambda{OID.new(realm).add_extension(ext)}.should.raise ArgumentError
+    ext::Request = nil
+    lambda{OID.new(realm).add_extension(ext)}.should.raise ArgumentError
+    ext::Response = nil
+    lambda{OID.new(realm).add_extension(ext)}.should.raise ArgumentError
+    ext::NS_URI = nil
+    lambda{OID.new(realm).add_extension(ext)}.should.raise TypeError
+  end
+
+  specify 'extensions should have Request and Response defined and inherit from OpenID::Extension' do
+    ext = Module.new
+    ext::Request = nil
+    ext::Response = nil
+    ext::NS_URI = nil
+    lambda{OID.new(realm).add_extension(ext)}.should.raise TypeError
+    ext::Request = Class.new(::OpenID::Extension)
+    lambda{OID.new(realm).add_extension(ext)}.should.raise TypeError
+    ext::Response = Class.new(::OpenID::Extension)
+    lambda{OID.new(realm).add_extension(ext)}.should.raise TypeError
+  end
+
+  specify 'extensions should have NS_URI defined and be a string' do
+    ext = Module.new
+    ext::Request = Class.new(::OpenID::Extension)
+    ext::Response = Class.new(::OpenID::Extension)
+    ext::NS_URI = nil
+    lambda{OID.new(realm).add_extension(ext)}.should.raise TypeError
+    ext::NS_URI = 'openid.net'
+    lambda{OID.new(realm).add_extension(ext)}.should.raise ArgumentError
+    ext::NS_URI = 'http://openid.net'
+    lambda{OID.new(realm).add_extension(ext)}.should.not.raise
+  end
 end