Reload jwk set if kid not found

lib/jwt/decode.rb

@@ -67,7 +67,21 @@ def find_key(&keyfinder)
 
     def find_from_jwk(jwks)
       raise JWT::DecodeError, 'No key id (kid) found from token headers' unless header['kid']
-      jwk = jwks[:keys].find { |key| key[:kid] == header['kid'] }
+
+      lazy = jwks.respond_to?(:call)
+      keys = if lazy
+               jwks.call({})
+             else
+               jwks
+             end
+
+      jwk = keys[:keys].find { |key| key[:kid] == header['kid'] }
+
+      if lazy && !jwk
+        keys = jwks.call(invalidate: true)
+        jwk = keys[:keys].find { |key| key[:kid] == header['kid'] }
+      end
+
       raise JWT::DecodeError, "Could not find public key for kid #{header['kid']}" unless jwk
 
       JWT::JWK.import(jwk).keypair

spec/decode_with_jwk_spec.rb

@@ -23,7 +23,7 @@
       end
     end
 
-    context 'when jwk keys are given' do
+    context 'when jwk keys are given as an array' do
       context 'and kid is in the set' do
         it 'is able to decode the token' do
           payload, _header = described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: public_jwks})
@@ -51,5 +51,20 @@
         end
       end
     end
+
+    context 'when jwk keys are loaded using a proc/lambda' do
+      it 'decodes the token' do
+        payload, _header = described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: lambda { |_opts| public_jwks }})
+        expect(payload).to eq(token_payload)
+      end
+    end
+
+    context 'when jwk keys are rotated' do
+      it 'decodes the token' do
+        key_loader = ->(options){ options[:invalidate] ? public_jwks : { keys: [] } }
+        payload, _header = described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: key_loader})
+        expect(payload).to eq(token_payload)
+      end
+    end
   end
 end