Add a JWKS implementation

lib/jwt/jwk.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative 'jwk/key_finder'
+require_relative 'jwk/set'
 
 module JWT
   module JWK

lib/jwt/jwk/ec.rb

@@ -22,6 +22,8 @@ def initialize(key, params = nil, options = {})
         params = { kid: params } if params.is_a?(String)
 
         key_params = case key
+                     when JWT::JWK::EC
+                       key.export(include_private: true)
                      when OpenSSL::PKey::EC # Accept OpenSSL key as input
                        @keypair = key # Preserve the object to avoid recreation
                        parse_ec_key(key)

lib/jwt/jwk/hmac.rb

@@ -16,6 +16,8 @@ def initialize(key, params = nil, options = {})
         params = { kid: params } if params.is_a?(String)
 
         key_params = case key
+                     when JWT::JWK::HMAC
+                       key.export(include_private: true)
                      when String # Accept String key as input
                        { kty: KTY, k: key }
                      when Hash

lib/jwt/jwk/key_base.rb

@@ -33,6 +33,14 @@ def []=(key, value)
         @parameters[key.to_sym] = value
       end
 
+      def ==(other)
+        self[:kid] == other[:kid]
+      end
+
+      def <=>(other)
+        self[:kid] <=> other[:kid]
+      end
+
       private
 
       attr_reader :parameters

lib/jwt/jwk/key_finder.rb

@@ -5,57 +5,37 @@ module JWK
     class KeyFinder
       def initialize(options)
         jwks_or_loader = options[:jwks]
-        @jwks          = jwks_or_loader if jwks_or_loader.is_a?(Hash)
-        @jwk_loader    = jwks_or_loader if jwks_or_loader.respond_to?(:call)
+
+        @jwks_loader = if jwks_or_loader.respond_to?(:call)
+          jwks_or_loader
+        else
+          ->(_options) { jwks_or_loader }
+        end
       end
 
       def key_for(kid)
         raise ::JWT::DecodeError, 'No key id (kid) found from token headers' unless kid
 
         jwk = resolve_key(kid)
 
-        raise ::JWT::DecodeError, 'No keys found in jwks' if jwks_keys.empty?
+        raise ::JWT::DecodeError, 'No keys found in jwks' unless @jwks.any?
         raise ::JWT::DecodeError, "Could not find public key for kid #{kid}" unless jwk
 
-        ::JWT::JWK.import(jwk).keypair
+        jwk.keypair
       end
 
       private
 
       def resolve_key(kid)
-        jwk = find_key(kid)
+        # First try without invalidation to facilitate application caching
+        @jwks ||= JWT::JWK::Set.new(@jwks_loader.call(kid: kid))
+        jwk = @jwks.find { |key| key[:kid] == kid }
 
         return jwk if jwk
 
-        if reloadable?
-          load_keys(invalidate: true, kid_not_found: true, kid: kid) # invalidate for backwards compatibility
-          return find_key(kid)
-        end
-
-        nil
-      end
-
-      def jwks
-        return @jwks if @jwks
-
-        load_keys
-        @jwks
-      end
-
-      def load_keys(opts = {})
-        @jwks = @jwk_loader.call(opts)
-      end
-
-      def jwks_keys
-        Array(jwks[:keys] || jwks['keys'])
-      end
-
-      def find_key(kid)
-        jwks_keys.find { |key| (key[:kid] || key['kid']) == kid }
-      end
-
-      def reloadable?
-        @jwk_loader
+        # Second try, invalidate for backwards compatibility
+        @jwks = JWT::JWK::Set.new(@jwks_loader.call(invalidate: true, kid_not_found: true, kid: kid))
+        @jwks.find { |key| key[:kid] == kid }
       end
     end
   end

lib/jwt/jwk/rsa.rb

@@ -17,6 +17,8 @@ def initialize(key, params = nil, options = {})
         params = { kid: params } if params.is_a?(String)
 
         key_params = case key
+                     when JWT::JWK::RSA
+                       key.export(include_private: true)
                      when OpenSSL::PKey::RSA # Accept OpenSSL key as input
                        @keypair = key # Preserve the object to avoid recreation
                        parse_rsa_key(key)

lib/jwt/jwk/set.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module JWT
+  module JWK
+    class Set
+      include Enumerable
+
+      attr_reader :keys
+
+      def initialize(jwks)
+        jwks ||= {}
+
+        @keys = case jwks
+                when JWT::JWK::Set # Simple duplication
+                  jwks.keys
+                when JWT::JWK::KeyBase # Singleton
+                  [jwks]
+                when Hash
+                  jwks = jwks.transform_keys(&:to_sym)
+                  [*jwks[:keys]].map { |k| JWT::JWK.new k }
+                when Array
+                  jwks.map { |k| JWT::JWK.new k }
+                else
+                  raise ArgumentError, 'Can only create new JWKS from Hash, Array and JWK'
+        end
+      end
+
+      def export(options = {})
+        { keys: @keys.map { |k| k.export(options) } }
+      end
+
+      def each(&block)
+        @keys.each(&block)
+      end
+
+      def select!(&block)
+        return @keys.select! unless block
+
+        self if @keys.select!(&block)
+      end
+
+      def reject!(&block)
+        return @keys.reject! unless block
+
+        self if @keys.reject!(&block)
+      end
+
+      alias filter! select!
+
+      def size
+        @keys.size
+      end
+
+      alias length size
+
+      def merge(enum)
+        @keys += JWT::JWK::Set.new(enum.collect)
+        self
+      end
+
+      def union(enum)
+        dup.merge(enum)
+      end
+
+      def add(key)
+        @keys << JWT::JWK.new(key)
+        self
+      end
+
+      def ==(other)
+        other.is_a?(JWT::JWK::Set) && keys.sort == other.keys.sort
+      end
+
+      # For symbolic manipulation
+      alias | union
+      alias + union
+      alias << add
+    end
+  end
+end

spec/jwk/decode_with_jwk_spec.rb

@@ -4,7 +4,7 @@
   describe '.decode for JWK usecase' do
     let(:keypair)       { OpenSSL::PKey::RSA.new(2048) }
     let(:jwk)           { JWT::JWK.new(keypair) }
-    let(:public_jwks) { { keys: [jwk.export, { kid: 'not_the_correct_one' }] } }
+    let(:public_jwks) { { keys: [jwk.export, { kid: 'not_the_correct_one', kty: 'oct', k: 'secret' }] } }
     let(:token_payload) { { 'data' => 'something' } }
     let(:token_headers) { { kid: jwk.kid } }
     let(:signed_token)  { described_class.encode(token_payload, jwk.keypair, 'RS512', token_headers) }