Add caching to remote JWKS fetch (#342)

* feat: add in-memory cache module for storing jwk set
* add tests for cache implementation
* add test to confirm jwks is cached
* add doc comments to cache.rb
* specify the time increment (seconds)

Author: nicknisi

lib/workos.rb

@@ -45,6 +45,7 @@ def self.key
   autoload :AuthenticationFactorAndChallenge, 'workos/authentication_factor_and_challenge'
   autoload :AuthenticationResponse, 'workos/authentication_response'
   autoload :AuditLogs, 'workos/audit_logs'
+  autoload :Cache, 'workos/cache'
   autoload :Challenge, 'workos/challenge'
   autoload :Client, 'workos/client'
   autoload :Connection, 'workos/connection'

lib/workos/cache.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+module WorkOS
+  # The Cache module provides a simple in-memory cache for storing values
+  # This module is not meant to be instantiated in a user space, and is used internally by the SDK
+  module Cache
+    # The Entry class represents a cache entry with a value and an expiration time
+    class Entry
+      attr_reader :value, :expires_at
+
+      # Initializes a new cache entry
+      # @param value [Object] The value to store in the cache
+      # @param expires_in_seconds [Integer, nil] The expiration time for the value in seconds, or nil for no expiration
+      def initialize(value, expires_in_seconds)
+        @value = value
+        @expires_at = expires_in_seconds ? Time.now + expires_in_seconds : nil
+      end
+
+      # Checks if the entry has expired
+      # @return [Boolean] True if the entry has expired, false otherwise
+      def expired?
+        return false if expires_at.nil?
+
+        Time.now > @expires_at
+      end
+    end
+
+    class << self
+      # Fetches a value from the cache, or calls the block to fetch the value if it is not present
+      # @param key [String] The key to fetch the value for
+      # @param expires_in [Integer] The expiration time for the value in seconds
+      # @param force [Boolean] If true, the value will be fetched from the block even if it is present in the cache
+      # @param block [Proc] The block to call to fetch the value if it is not present in the cache
+      # @return [Object] The value fetched from the cache or the block
+      def fetch(key, expires_in: nil, force: false, &block)
+        entry = store[key]
+
+        if force || entry.nil? || entry.expired?
+          value = block.call
+          store[key] = Entry.new(value, expires_in)
+          return value
+        end
+
+        entry.value
+      end
+
+      # Reads a value from the cache
+      # @param key [String] The key to read the value for
+      # @return [Object] The value read from the cache, or nil if the value is not present or has expired
+      def read(key)
+        entry = store[key]
+        return nil if entry.nil? || entry.expired?
+
+        entry.value
+      end
+
+      # Writes a value to the cache
+      # @param key [String] The key to write the value for
+      # @param value [Object] The value to write to the cache
+      # @param expires_in [Integer] The expiration time for the value in seconds
+      # @return [Object] The value written to the cache
+      def write(key, value, expires_in: nil)
+        store[key] = Entry.new(value, expires_in)
+        value
+      end
+
+      # Deletes a value from the cache
+      # @param key [String] The key to delete the value for
+      def delete(key)
+        store.delete(key)
+      end
+
+      # Clears all values from the cache
+      def clear
+        store.clear
+      end
+
+      # Checks if a value exists in the cache
+      # @param key [String] The key to check for
+      # @return [Boolean] True if the value exists and has not expired, false otherwise
+      def exist?(key)
+        entry = store[key]
+        !(entry.nil? || entry.expired?)
+      end
+
+      private
+
+      # The in-memory store for the cache
+      def store
+        @store ||= {}
+      end
+    end
+  end
+end

lib/workos/session.rb

@@ -23,7 +23,9 @@ def initialize(user_management:, client_id:, session_data:, cookie_password:)
       @session_data = session_data
       @client_id = client_id
 
-      @jwks = create_remote_jwk_set(URI(@user_management.get_jwks_url(client_id)))
+      @jwks = Cache.fetch("jwks_#{client_id}", expires_in: 5 * 60) do
+        create_remote_jwk_set(URI(@user_management.get_jwks_url(client_id)))
+      end
       @jwks_algorithms = @jwks.map { |key| key[:alg] }.compact.uniq
     end
 

spec/lib/workos/cache_spec.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+describe WorkOS::Cache do
+  before { described_class.clear }
+
+  describe '.write and .read' do
+    it 'stores and retrieves data' do
+      described_class.write('key', 'value')
+      expect(described_class.read('key')).to eq('value')
+    end
+
+    it 'returns nil if key does not exist' do
+      expect(described_class.read('missing')).to be_nil
+    end
+  end
+
+  describe '.fetch' do
+    it 'returns cached value when present and not expired' do
+      described_class.write('key', 'value')
+      fetch_value = described_class.fetch('key') { 'new_value' }
+      expect(fetch_value).to eq('value')
+    end
+
+    it 'executes block and caches value when not present' do
+      fetch_value = described_class.fetch('key') { 'new_value' }
+      expect(fetch_value).to eq('new_value')
+    end
+
+    it 'executes block and caches value when force is true' do
+      described_class.write('key', 'value')
+      fetch_value = described_class.fetch('key', force: true) { 'new_value' }
+      expect(fetch_value).to eq('new_value')
+    end
+  end
+
+  describe 'expiration' do
+    it 'expires values after specified time' do
+      described_class.write('key', 'value', expires_in: 0.1)
+      expect(described_class.read('key')).to eq('value')
+      sleep 0.2
+      expect(described_class.read('key')).to be_nil
+    end
+
+    it 'executes block and caches new value when expired' do
+      described_class.write('key', 'old_value', expires_in: 0.1)
+      sleep 0.2
+      fetch_value = described_class.fetch('key') { 'new_value' }
+      expect(fetch_value).to eq('new_value')
+    end
+
+    it 'does not expire values when expires_in is nil' do
+      described_class.write('key', 'value', expires_in: nil)
+      sleep 0.2
+      expect(described_class.read('key')).to eq('value')
+    end
+  end
+
+  describe '.exist?' do
+    it 'returns true if key exists' do
+      described_class.write('key', 'value')
+      expect(described_class.exist?('key')).to be true
+    end
+
+    it 'returns false if expired' do
+      described_class.write('key', 'value', expires_in: 0.1)
+      sleep 0.2
+      expect(described_class.exist?('key')).to be false
+    end
+
+    it 'returns false if key does not exist' do
+      expect(described_class.exist?('missing')).to be false
+    end
+  end
+
+  describe '.delete' do
+    it 'deletes key' do
+      described_class.write('key', 'value')
+      described_class.delete('key')
+      expect(described_class.read('key')).to be_nil
+    end
+  end
+
+  describe '.clear' do
+    it 'removes all keys from the cache' do
+      described_class.write('key1', 'value1')
+      described_class.write('key2', 'value2')
+
+      described_class.clear
+
+      expect(described_class.read('key1')).to be_nil
+      expect(described_class.read('key2')).to be_nil
+    end
+  end
+end

spec/lib/workos/session_spec.rb

@@ -19,6 +19,52 @@
       allow(user_management).to receive(:get_jwks_url).with(client_id).and_return(jwks_url)
     end
 
+    describe 'JWKS caching' do
+      before do
+        WorkOS::Cache.clear
+      end
+
+      it 'caches and returns JWKS' do
+        expect(Net::HTTP).to receive(:get).once
+        session1 = WorkOS::Session.new(
+          user_management: user_management,
+          client_id: client_id,
+          session_data: session_data,
+          cookie_password: cookie_password,
+        )
+
+        session2 = WorkOS::Session.new(
+          user_management: user_management,
+          client_id: client_id,
+          session_data: session_data,
+          cookie_password: cookie_password,
+        )
+
+        expect(session1.jwks.map(&:export)).to eq(session2.jwks.map(&:export))
+      end
+
+      it 'fetches JWKS from remote when cache is expired' do
+        expect(Net::HTTP).to receive(:get).twice
+        session1 = WorkOS::Session.new(
+          user_management: user_management,
+          client_id: client_id,
+          session_data: session_data,
+          cookie_password: cookie_password,
+        )
+
+        allow(Time).to receive(:now).and_return(Time.now + 301)
+
+        session2 = WorkOS::Session.new(
+          user_management: user_management,
+          client_id: client_id,
+          session_data: session_data,
+          cookie_password: cookie_password,
+        )
+
+        expect(session1.jwks.map(&:export)).to eq(session2.jwks.map(&:export))
+      end
+    end
+
     it 'raises an error if cookie_password is nil or empty' do
       expect do
         WorkOS::Session.new(