From f2f06e112e990dc71777e11e66bd39e470290f4c Mon Sep 17 00:00:00 2001 From: Nick Nisi Date: Thu, 16 Jan 2025 09:08:09 -0700 Subject: [PATCH] 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) --- lib/workos.rb | 1 + lib/workos/cache.rb | 94 +++++++++++++++++++++++++++++++++ lib/workos/session.rb | 4 +- spec/lib/workos/cache_spec.rb | 94 +++++++++++++++++++++++++++++++++ spec/lib/workos/session_spec.rb | 46 ++++++++++++++++ 5 files changed, 238 insertions(+), 1 deletion(-) create mode 100644 lib/workos/cache.rb create mode 100644 spec/lib/workos/cache_spec.rb diff --git a/lib/workos.rb b/lib/workos.rb index 2486f9a7..f1c36fbc 100644 --- a/lib/workos.rb +++ b/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' diff --git a/lib/workos/cache.rb b/lib/workos/cache.rb new file mode 100644 index 00000000..7bd594ca --- /dev/null +++ b/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 diff --git a/lib/workos/session.rb b/lib/workos/session.rb index 11be446c..348b09fa 100644 --- a/lib/workos/session.rb +++ b/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 diff --git a/spec/lib/workos/cache_spec.rb b/spec/lib/workos/cache_spec.rb new file mode 100644 index 00000000..0c301cb7 --- /dev/null +++ b/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 diff --git a/spec/lib/workos/session_spec.rb b/spec/lib/workos/session_spec.rb index 960f8b93..08d567e5 100644 --- a/spec/lib/workos/session_spec.rb +++ b/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(