fix: Cache keyset for decoding

Cache keyset for decoding

keypairs.gemspec

@@ -39,6 +39,6 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency 'rubocop-performance'            # Linter for Performance optimization analysis
   spec.add_development_dependency 'rubocop-rails'                  # Linter for Rails-specific analysis
   spec.add_development_dependency 'shoulda-matchers'               # RSpec matchers
-  spec.add_development_dependency 'sqlite3'                        # Database adapter
+  spec.add_development_dependency 'sqlite3', '~> 1.4'              # Database adapter
   spec.add_development_dependency 'timecop'                        # Freeze time to test time-dependent code
 end

lib/keypair.rb

@@ -34,7 +34,7 @@
 # @attr [Time] not_before The time before which no payloads may be signed using the keypair.
 # @attr [Time] not_after The time after which no payloads may be signed using the keypair.
 # @attr [Time] expires_at The time after which the keypair may not be used for signature validation.
-class Keypair < ActiveRecord::Base
+class Keypair < ActiveRecord::Base # rubocop:disable Metrics/ClassLength
   ALGORITHM = 'RS256'
   ROTATION_INTERVAL = 1.month
 
@@ -133,13 +133,27 @@ def self.jwt_decode(id_token, options = {})
       # Change the default algorithm to match the encoding algorithm
       algorithm: ALGORITHM,
       # Load our own keyset as valid keys
-      jwks: keyset,
+      jwks: jwk_loader_cached,
       # If the `sub` is provided, validate that it matches the payload `sub`
       verify_sub: true
     )
     JWT.decode(id_token, nil, true, options).first.with_indifferent_access
   end
 
+  # options[:invalidate] will be `true` if a matching `kid` was not found
+  # https://github.com/jwt/ruby-jwt/blob/master/lib/jwt/jwk/key_finder.rb#L31
+  def self.jwk_loader_cached
+    lambda do |options|
+      cached_jwks(force: options[:invalidate]) || {}
+    end
+  end
+
+  def self.cached_jwks(force: false)
+    Rails.cache.fetch('keypairs/Keypair/jwks', force: force, skip_nil: true) do
+      keyset
+    end
+  end
+
   # JWT encodes the payload with this keypair.
   # It automatically adds the security attributes +iat+, +exp+ and +nonce+ to the payload.
   # It automatically sets the +kid+ in the header.

lib/keypairs/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Keypairs
-  VERSION = '2.0.0.develop'
+  VERSION = '1.3.4.develop'
 end

spec/models/keypair_spec.rb

@@ -470,6 +470,36 @@
           expect { subject }.to raise_error JWT::DecodeError
         end
       end
+
+      context 'should update cache if kid is not regonized' do
+        let(:other_payload) { { uuid: SecureRandom.uuid } }
+        let(:other_id_token) { other_keypair.jwt_encode(other_payload) }
+        let(:other_keypair) { described_class.create! }
+        let(:keypair) { described_class.current }
+        subject { described_class.jwt_decode(other_id_token).deep_symbolize_keys }
+
+        before do
+          old_keypair = described_class.create!
+          token = old_keypair.jwt_encode({ uuid: SecureRandom.uuid })
+          described_class.jwt_decode(token)
+        end
+
+        it { expect(subject[:uuid]).to eq other_payload[:uuid] }
+      end
+
+      context 'with query count' do
+        let!(:keypair) { described_class.current }
+
+        context 'without cache' do
+          it { expect { subject }.not_to exceed_query_limit(4) }
+        end
+
+        context 'with cache' do
+          before { subject }
+
+          it { expect { subject }.not_to exceed_query_limit(0) }
+        end
+      end
     end
 
     describe '#public_jwk_export' do

spec/support/matchers/encrypt_attribute.rb

@@ -6,9 +6,9 @@
   match do |model|
     # Correct responds to methods
     model.respond_to?(attribute) &&
-      model.respond_to?("#{attribute}=") &&
+      model.respond_to?(:"#{attribute}=") &&
       model.respond_to?(database_column_name) &&
-      model.respond_to?("#{database_column_name}=") &&
+      model.respond_to?(:"#{database_column_name}=") &&
       # Correct database columns
       model.class.column_names.exclude?(attribute.to_s) &&
       model.class.column_names.include?(database_column_name)

spec/support/matchers/exceed_query_limit.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+# Based on https://stackoverflow.com/a/13423584
+RSpec::Matchers.define :exceed_query_limit do |expected|
+  match do |block|
+    # For now we will need to exclude the keypairs since its value is not cached
+    @excluded_queries = [/SELECT "public"\."keypairs"\.\* FROM "public"\."keypairs"/]
+
+    query_count(&block) > expected
+  end
+
+  failure_message_when_negated do |_actual|
+    queries = @queries.map do |query|
+      if query[:location]
+        <<~TEXT
+          #{query[:name]}: #{query[:sql]}
+            ↳ #{query[:location]}
+        TEXT
+      else
+        <<~TEXT
+          #{query[:name]}: #{query[:sql]}
+        TEXT
+      end
+    end.join.indent(4)
+
+    <<~TEXT
+      Expected to run maximum #{expected} queries, got #{@query_count}:
+      #{queries}
+    TEXT
+  end
+
+  chain(:with_excluded_query) do |*excluded_queries|
+    @excluded_queries.push(*excluded_queries)
+  end
+
+  def query_count(&block)
+    @query_count = 0
+    @queries = []
+    ActiveSupport::Notifications.subscribed(method(:query_callback), 'sql.active_record', &block)
+    @query_count
+  end
+
+  def query_callback(_name, _start, _finish, _message_id, values)
+    return if %w[CACHE SCHEMA].include?(values[:name])
+
+    @excluded_queries&.each do |excluded_query|
+      return if excluded_query.match?(values[:sql])
+    end
+
+    @query_count += 1
+    @queries << { sql: values[:sql], name: values[:name], location: Rails.backtrace_cleaner.clean(caller).first }
+  end
+
+  def supports_block_expectations?
+    true
+  end
+end