Utilize key_fields instead of manually defining key matching order

Author: hieuk09

lib/jwt/decode.rb

@@ -60,7 +60,13 @@ def verify_algo
 
     def set_key
       @key = find_key(&@keyfinder) if @keyfinder
-      @key = ::JWT::JWK::KeyFinder.new(jwks: @options[:jwks], allow_nil_kid: @options[:allow_nil_kid]).call(token) if @options[:jwks]
+      if @options[:jwks]
+        @key = ::JWT::JWK::KeyFinder.new(
+          jwks: @options[:jwks],
+          allow_nil_kid: @options[:allow_nil_kid],
+          key_fields: @options[:key_fields]
+        ).call(token)
+      end
 
       return unless (x5c_options = @options[:x5c])
 

lib/jwt/jwk/key_finder.rb

@@ -9,6 +9,9 @@ class KeyFinder
       # @param [Hash] options the options to create a KeyFinder with
       # @option options [Proc, JWT::JWK::Set] :jwks the jwks or a loader proc
       # @option options [Boolean] :allow_nil_kid whether to allow nil kid
+      # @option options [Array] :key_fields the fields to use for key matching,
+      #                         the order of the fields are used to determine
+      #                         the priority of the keys.
       def initialize(options)
         @allow_nil_kid = options[:allow_nil_kid]
         jwks_or_loader = options[:jwks]
@@ -18,6 +21,8 @@ def initialize(options)
                        else
                          ->(_options) { jwks_or_loader }
                        end
+
+        @key_fields = options[:key_fields] || %i[kid]
       end
 
       # Returns the verification key for the given kid
@@ -36,21 +41,16 @@ def key_for(kid, key_field = :kid)
       # Returns the key for the given token
       # @param [JWT::EncodedToken] token the token
       def call(token)
-        kid = token.header['kid']
-        x5t = token.header['x5t']
-        x5c = token.header['x5c']
-
-        if kid
-          key_for(kid, :kid)
-        elsif x5t
-          key_for(x5t, :x5t)
-        elsif x5c
-          key_for(x5c, :x5c)
-        elsif @allow_nil_kid
-          key_for(kid)
-        else
-          raise ::JWT::DecodeError, 'No key id (kid) or x5t found from token headers'
+        @key_fields.each do |key_field|
+          field_value = token.header[key_field.to_s]
+
+          return key_for(field_value, key_field) if field_value
         end
+
+        raise ::JWT::DecodeError, 'No key id (kid) or x5t found from token headers' unless @allow_nil_kid
+
+        kid = token.header['kid']
+        key_for(kid)
       end
 
       private

spec/jwt/jwk/decode_with_jwk_spec.rb

@@ -44,7 +44,17 @@
         let(:valid_key) { jwk.export.merge({ x5t: x5t }) }
         let(:token_headers) { { x5t: x5t } }
         it 'is able to decode the token' do
-          payload, _header = described_class.decode(signed_token, nil, true, { algorithms: [algorithm], jwks: public_jwks })
+          payload, _header = described_class.decode(signed_token, nil, true, { algorithms: [algorithm], jwks: public_jwks, key_fields: [:x5t] })
+          expect(payload).to eq(token_payload)
+        end
+      end
+
+      context 'and both kid and x5t is in the set' do
+        let(:x5t) { Base64.urlsafe_encode64(OpenSSL::Digest::SHA1.new(keypair.to_der).digest, padding: false) }
+        let(:valid_key) { jwk.export.merge({ x5t: x5t }) }
+        let(:token_headers) { { x5t: x5t, kid: 'NOT_A_MATCH' } }
+        it 'is able to decode the token based on the priority of the key defined in key_fields' do
+          payload, _header = described_class.decode(signed_token, nil, true, { algorithms: [algorithm], jwks: public_jwks, key_fields: %i[x5t kid] })
           expect(payload).to eq(token_payload)
         end
       end