Moved the JWK key resolving to own class

lib/jwt.rb

@@ -6,6 +6,8 @@
 require 'jwt/encode'
 require 'jwt/error'
 require 'jwt/jwk'
+require 'jwt/jwk_key_finder'
+
 # JSON Web Token implementation
 #
 # Should be up to date with the latest spec:

lib/jwt/decode.rb

@@ -39,7 +39,7 @@ def decode_segments
 
     def verify_signature
       @key = find_key(&@keyfinder) if @keyfinder
-      @key = find_from_jwk(@options[:jwks]) if @options[:jwks]
+      @key = ::JWT::JWKKeyFinder.new(jwks: @options[:jwks]).key_for(header['kid']) if @options[:jwks]
 
       raise(JWT::IncorrectAlgorithm, 'An algorithm must be specified') if allowed_algorithms.empty?
       raise(JWT::IncorrectAlgorithm, 'Expected a different algorithm') unless options_includes_algo_in_header?
@@ -65,29 +65,6 @@ def find_key(&keyfinder)
       key
     end
 
-    def find_from_jwk(jwks)
-      kid = header['kid']
-      raise JWT::DecodeError, 'No key id (kid) found from token headers' unless kid
-
-      lazy = jwks.respond_to?(:call)
-      keys = if lazy
-        jwks.call({})
-      else
-        jwks
-      end
-
-      jwk = keys[:keys].find { |key| key[:kid] == kid }
-
-      if lazy && !jwk
-        keys = jwks.call(invalidate: true)
-        jwk = keys[:keys].find { |key| key[:kid] == kid }
-      end
-
-      raise JWT::DecodeError, "Could not find public key for kid #{kid}" unless jwk
-
-      JWT::JWK.import(jwk).keypair
-    end
-
     def verify_claims
       Verify.verify_claims(payload, @options)
     end

lib/jwt/jwk.rb

@@ -15,6 +15,7 @@ def initialize(keypair)
 
     def supported_key!(keypair)
       return if keypair.is_a?(OpenSSL::PKey::RSA)
+
       raise JWT::JWKError, "Key of type #{keypair.class.name} not supported"
     end
 
@@ -29,8 +30,8 @@ def export
       when OpenSSL::PKey::RSA
         {
           kty: 'RSA',
-          n: self.class.to_base64(public_key.n.to_s(BINARY)),
-          e: self.class.to_base64(public_key.e.to_s(BINARY)),
+          n: Base64.urlsafe_encode64(public_key.n.to_s(BINARY), padding: false),
+          e: Base64.urlsafe_encode64(public_key.e.to_s(BINARY), padding: false),
           kid: kid
         }
       end
@@ -41,22 +42,14 @@ def import(jwk_data)
         case jwk_data[:kty]
         when 'RSA'
           imported_key = OpenSSL::PKey::RSA.new
-          imported_key.set_key(OpenSSL::BN.new(from_base64(jwk_data[:n]), BINARY), 
-                               OpenSSL::BN.new(from_base64(jwk_data[:e]), BINARY),
-                               nil)
+          imported_key.set_key(OpenSSL::BN.new(Base64.urlsafe_decode64(jwk_data[:n]), BINARY),
+            OpenSSL::BN.new(Base64.urlsafe_decode64(jwk_data[:e]), BINARY),
+            nil)
           self.new(imported_key)
         else
           raise JWT::JWKError, "Key type #{jwk_data[:kty]} not supported"
         end
       end
-
-      def to_base64(input)
-        Base64.urlsafe_encode64(input, padding: false)
-      end
-
-      def from_base64(input)
-        Base64.urlsafe_decode64(input)
-      end
     end
   end
 end

lib/jwt/jwk_key_finder.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module JWT
+  class JWKKeyFinder
+    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)
+    end
+
+    def key_for(kid)
+      raise JWT::DecodeError, 'No key id (kid) found from token headers' unless kid
+
+      jwk = find_key(kid)
+
+      if !jwk && reloadable?
+        load_keys(invalidate: true)
+        jwk = find_key(kid)
+      end
+
+      raise JWT::DecodeError, "Could not find public key for kid #{kid}" unless jwk
+
+      JWT::JWK.import(jwk).keypair
+    end
+
+    private
+
+    def jwks
+      return @jwks if @jwks
+
+      load_keys
+      @jwks
+    end
+
+    def load_keys(opts = {})
+      @jwks = @jwk_loader.call(opts)
+    end
+
+    def find_key(kid)
+      Array(jwks[:keys]).find { |key| key[:kid] == kid }
+    end
+
+    def reloadable?
+      @jwk_loader
+    end
+  end
+end