From 8b45133b0865c1a96531dfd2ac894a9e60a7095d Mon Sep 17 00:00:00 2001
From: Koen Sengers <koen.sengers@drieam.com>
Date: Fri, 2 Aug 2024 09:30:34 +0200
Subject: [PATCH] fix: Cache keyset for decoding

Cache keyset for decoding
---
 .github/workflows/ci.yml                    |   3 +-
 keypairs.gemspec                            |   2 +-
 lib/keypair.rb                              |  18 +++++++-
 lib/keypairs/version.rb                     |   2 +-
 spec/internal/db/keypairs_test.sqlite3-shm  | Bin 0 -> 32768 bytes
 spec/internal/db/keypairs_test.sqlite3-wal  | Bin 0 -> 16512 bytes
 spec/models/keypair_spec.rb                 |  30 +++++++++++++
 spec/support/matchers/encrypt_attribute.rb  |   4 +-
 spec/support/matchers/exceed_query_limit.rb |  46 ++++++++++++++++++++
 9 files changed, 98 insertions(+), 7 deletions(-)
 create mode 100644 spec/internal/db/keypairs_test.sqlite3-shm
 create mode 100644 spec/internal/db/keypairs_test.sqlite3-wal
 create mode 100644 spec/support/matchers/exceed_query_limit.rb

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 9e72e0e..01b9e05 100644
--- a/.github/workflows/ci.yml
+++ b/.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:
diff --git a/keypairs.gemspec b/keypairs.gemspec
index b8a6d72..3e46dad 100644
--- a/keypairs.gemspec
+++ b/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
diff --git a/lib/keypair.rb b/lib/keypair.rb
index 38e19c0..d2a8cdb 100644
--- a/lib/keypair.rb
+++ b/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.
diff --git a/lib/keypairs/version.rb b/lib/keypairs/version.rb
index 1d26233..d354a9b 100644
--- a/lib/keypairs/version.rb
+++ b/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
diff --git a/spec/internal/db/keypairs_test.sqlite3-shm b/spec/internal/db/keypairs_test.sqlite3-shm
new file mode 100644
index 0000000000000000000000000000000000000000..75609407622ba18fa71cff44a4c951f02be22686
GIT binary patch
literal 32768
zcmeI)y9ok85C-7$fsf4aumv%*3nK$TF0l?PFxJG(z!n6{F}Nl)(^-+j)KK^PVc20=
zcG<6hDeGyG(vE5s#i$$WFzS3<%<A!SvODb8*Uh}DZl}%tvf7qko{y_uN%uZJ^4@Ni
zUUY+~^mf15)es;+fB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk
z1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t9|jAm%&@
zflLL`BvK)er$CzXD+KZs=;Zlrza~I{009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs
O0RjXF5FkL{F9n{<xhHP`

literal 0
HcmV?d00001

diff --git a/spec/internal/db/keypairs_test.sqlite3-wal b/spec/internal/db/keypairs_test.sqlite3-wal
new file mode 100644
index 0000000000000000000000000000000000000000..0298ddeb8d2c8408156e8fb6709ccf0d9c987c6c
GIT binary patch
literal 16512
zcmeI(F$)1<9LMp~)wv_F&~;HZgBWyWF_<hSgOWv*K^9T!Cgl}iQItVOSu7T--C#A`
zZ1V=(L#KNUci;c?^si_5ea2tUQ=TbvsbK4pl#s+M*Q?Ics<%yRbEq9(IcKGK_vuz;
zdU|Y2F|qtz5hn%#1Q0*~0R#|0009ILKmY**{3)P{SZ-giPs%+xd%Hx$l3qY4FbE)k
z00IagfB*srAb<b@2>4k*bHZtEFJM#`#p_*9d3piS8+_js_!)t02q1s}0tg_000Iag
zfB*sr1XLjG<o@>pW@C5GHx1kUM-YDpRE5|GAb<b@2q1s}0tg_000Id3Qy?0WGwQH!
Jjk<&J<O9|&EOGz<

literal 0
HcmV?d00001

diff --git a/spec/models/keypair_spec.rb b/spec/models/keypair_spec.rb
index 0676fe1..6cd3952 100644
--- a/spec/models/keypair_spec.rb
+++ b/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
diff --git a/spec/support/matchers/encrypt_attribute.rb b/spec/support/matchers/encrypt_attribute.rb
index b046207..e0b823c 100644
--- a/spec/support/matchers/encrypt_attribute.rb
+++ b/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)
diff --git a/spec/support/matchers/exceed_query_limit.rb b/spec/support/matchers/exceed_query_limit.rb
new file mode 100644
index 0000000..136d6ca
--- /dev/null
+++ b/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