jwk_key_finder.rb

# 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