Fetch discovered JWKS by kid so it can be cached

lib/omniauth/strategies/openid_connect.rb

@@ -307,7 +307,7 @@ def decode_id_token(id_token)
           when :HS256, :HS384, :HS512
             secret
           else
-            public_key
+            discovered_jwk(decoded.kid) || public_key
           end
 
         decoded.verify!(keyset)
@@ -327,6 +327,15 @@ def decode_id_token(id_token)
         decoded
       end
 
+      # Looks the key up by kid the way ::OpenIDConnect::ResponseObject::IdToken.decode
+      # does with a discovery config: JSON::JWK::Set::Fetcher caches the JWKS in
+      # JSON::JWK::Set::Fetcher.cache and fetches it again when a new kid appears.
+      def discovered_jwk(kid)
+        return unless options.discovery && kid
+
+        JSON::JWK::Set::Fetcher.fetch(config.jwks_uri, kid: kid)
+      end
+
       # If client_signing_alg is specified, we check that the returned JWT
       # matches the expected algorithm. If not, we reject it.
       def validate_client_algorithm!(algorithm)

test/lib/omniauth/strategies/openid_connect_test.rb

@@ -474,6 +474,44 @@ def test_callback_phase_with_discovery # rubocop:disable Metrics/AbcSize
         strategy.callback_phase
       end
 
+      def test_callback_phase_with_discovery_caches_jwks_by_kid
+        JSON::JWK::Set::Fetcher.cache = ActiveSupport::Cache::MemoryStore.new
+        key = JSON::JWK.new(private_key)
+        stub_jwks_uri(key)
+
+        2.times { discovery_id_token_callback(JSON::JWT.new(payload).sign(key, :RS256)) }
+
+        assert_requested :get, 'https://example.com/jwks', times: 1
+      ensure
+        JSON::JWK::Set::Fetcher.cache = JSON::JWK::Set::Fetcher::Cache.new
+      end
+
+      def test_callback_phase_with_discovery_after_key_rotation
+        JSON::JWK::Set::Fetcher.cache = ActiveSupport::Cache::MemoryStore.new
+        old_key = JSON::JWK.new(private_key)
+        new_key = JSON::JWK.new(OpenSSL::PKey::RSA.generate(2048))
+        stub_jwks_uri(old_key).then.to_return(jwks_response(old_key, new_key))
+
+        discovery_id_token_callback(JSON::JWT.new(payload).sign(old_key, :RS256))
+        discovery_id_token_callback(JSON::JWT.new(payload).sign(new_key, :RS256))
+
+        assert_requested :get, 'https://example.com/jwks', times: 2
+      ensure
+        JSON::JWK::Set::Fetcher.cache = JSON::JWK::Set::Fetcher::Cache.new
+      end
+
+      def test_callback_phase_with_discovery_rejects_a_kid_missing_from_the_jwks
+        JSON::JWK::Set::Fetcher.cache = ActiveSupport::Cache::MemoryStore.new
+        stub_jwks_uri(JSON::JWK.new(private_key))
+        unknown_key = JSON::JWK.new(OpenSSL::PKey::RSA.generate(2048))
+
+        assert_raises JSON::JWK::Set::KidNotFound do
+          discovery_id_token_callback(JSON::JWT.new(payload).sign(unknown_key, :RS256))
+        end
+      ensure
+        JSON::JWK::Set::Fetcher.cache = JSON::JWK::Set::Fetcher::Cache.new
+      end
+
       def test_callback_phase_with_send_state_disabled # rubocop:disable Metrics/AbcSize
         code = SecureRandom.hex(16)
 
@@ -971,6 +1009,38 @@ def test_option_pkce
         assert(strategy.authorize_uri =~ /#{Regexp.quote(strategy.options.pkce_verifier.call)}/,
                'URI must contain code challenge value')
       end
+
+      private
+
+      def stub_jwks_uri(key)
+        stub_request(:get, 'https://example.com/jwks').to_return(jwks_response(key))
+      end
+
+      def jwks_response(*keys)
+        { body: { keys: keys }.to_json, headers: { 'Content-Type' => 'application/json' } }
+      end
+
+      def discovery_id_token_callback(id_token)
+        @strategy = nil
+        state = SecureRandom.hex(16)
+        request.stubs(:params).returns('id_token' => id_token.to_s, 'state' => state)
+        request.stubs(:path_info).returns('')
+
+        config = ::OpenIDConnect::Discovery::Provider::Config::Response.new(
+          issuer: issuer,
+          authorization_endpoint: 'https://example.com/authorization',
+          jwks_uri: 'https://example.com/jwks'
+        )
+        ::OpenIDConnect::Discovery::Provider::Config.stubs(:discover!).with(issuer).returns(config)
+
+        strategy.options.issuer = issuer
+        strategy.options.discovery = true
+        strategy.options.response_type = 'id_token'
+        strategy.unstub(:user_info)
+        strategy.expects(:fail!).never
+        strategy.call!('rack.session' => { 'omniauth.state' => state, 'omniauth.nonce' => nonce })
+        strategy.callback_phase
+      end
     end
   end
 end