fix: Cache keyset for decoding

.github/workflows/ci.yml

@@ -3,9 +3,10 @@ on: [push]
 jobs:
  reviewdog:
  runs-on: ubuntu-latest
+
  steps:
  - name: Check out code
- uses: actions/checkout@v2
+ uses: actions/checkout@v4
 
  - uses: ruby/setup-ruby@v1
  with:

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,46 @@
+# frozen_string_literal: true
+
+# Based on https://stackoverflow.com/a/13423584
+RSpec::Matchers.define :exceed_query_limit do |expected|
+ match do |block|
+ 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
+
+ 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])
+
+ @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