From cb51b3cb072369bb1c5f41e66ee753e5ca6108b3 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
---
 keypairs.gemspec                           |  2 +-
 lib/keypair.rb                             | 18 ++++++++++++++++--
 lib/keypairs/version.rb                    |  2 +-
 spec/models/keypair_spec.rb                | 16 ++++++++++++++++
 spec/support/matchers/encrypt_attribute.rb |  4 ++--
 5 files changed, 36 insertions(+), 6 deletions(-)

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/models/keypair_spec.rb b/spec/models/keypair_spec.rb
index 0676fe1..616ec4f 100644
--- a/spec/models/keypair_spec.rb
+++ b/spec/models/keypair_spec.rb
@@ -470,6 +470,22 @@
           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
     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)