fix: Make JWKS cache shared across SDK instances

After 1123911 the JWKs cache wasn't working as expected, since the middleware constructed a new SDK instance on each request. So each request was effectively bypassing the cache, since the cache was essentially an instance variable which was re-initialized anew every time. With this change, the JWKs cache is made thread-safe and shared between all instances of the SDK. Fixes AUTH-76
clerk · Dec 22, 2022 · a875122 · a875122
1 parent 3188a66
commit a875122
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 35 deletions.
diff --git a/Gemfile.lock b/Gemfile.lock
@@ -1,14 +1,16 @@
 PATH
   remote: .
   specs:
-    clerk-sdk-ruby (2.7.0)
+    clerk-sdk-ruby (2.8.0)
+      concurrent-ruby (~> 1.1)
       faraday (~> 1.4.1)
-      jwt (~> 2.2)
+      jwt (~> 2.5)
 
 GEM
   remote: https://rubygems.org/
   specs:
     byebug (11.1.3)
+    concurrent-ruby (1.1.10)
     faraday (1.4.3)
       faraday-em_http (~> 1.0)
       faraday-em_synchrony (~> 1.0)
@@ -23,11 +25,11 @@ GEM
     faraday-net_http (1.0.1)
     faraday-net_http_persistent (1.2.0)
     jwt (2.5.0)
-    minitest (5.14.2)
+    minitest (5.16.3)
     multipart-post (2.2.3)
-    rake (13.0.3)
+    rake (13.0.6)
     ruby2_keywords (0.0.5)
-    timecop (0.9.4)
+    timecop (0.9.6)
 
 PLATFORMS
   universal-darwin-21

diff --git a/clerk-sdk-ruby.gemspec b/clerk-sdk-ruby.gemspec
@@ -28,7 +28,8 @@ Gem::Specification.new do |spec|
   spec.require_paths = ["lib"]
 
   spec.add_dependency "faraday", "~> 1.4.1"
-  spec.add_dependency "jwt", '~> 2.2'
+  spec.add_dependency "jwt", '~> 2.5'
+  spec.add_dependency "concurrent-ruby", "~> 1.1"
 
   spec.add_development_dependency "byebug", "~> 11.1"
   spec.add_development_dependency "timecop", "~> 0.9.4"

diff --git a/lib/clerk/jwks_cache.rb b/lib/clerk/jwks_cache.rb
@@ -0,0 +1,32 @@
+class JWKSCache
+  def initialize(lifetime)
+    @lifetime = lifetime
+    @jwks = nil
+    @last_update = nil
+    @lock = Concurrent::ReadWriteLock.new
+  end
+
+  def fetch(sdk, force_refresh: false, kid_not_found: false)
+    should_refresh = @lock.with_read_lock do
+      @jwks.nil? || @last_update.nil? || force_refresh ||
+        (Time.now.to_i-@last_update > @lifetime) ||
+        (kid_not_found && Time.now.to_i-@last_update > 300)
+    end
+
+    if should_refresh
+      @lock.with_write_lock do
+        @last_update = Time.now.to_i
+
+        @jwks = begin
+          sdk.jwks.all["keys"]
+        rescue Clerk::Errors::Base
+          nil
+        end
+      end
+    end
+
+    @lock.with_read_lock do
+      @jwks
+    end
+  end
+end
diff --git a/lib/clerk/sdk.rb b/lib/clerk/sdk.rb
@@ -5,6 +5,7 @@
 require "net/http"
 require "json"
 require "jwt"
+require "concurrent-ruby"
 
 require_relative "resources/allowlist_identifiers"
 require_relative "resources/allowlist"
@@ -19,6 +20,7 @@
 require_relative "resources/users"
 require_relative "resources/jwks"
 require_relative "errors"
+require_relative "jwks_cache"
 
 module Clerk
   class SDK
@@ -30,10 +32,14 @@ class SDK
     # How often (in seconds) should JWKs be refreshed
     JWKS_CACHE_LIFETIME = 3600 # 1 hour
 
+    @@jwks_cache = JWKSCache.new(JWKS_CACHE_LIFETIME)
+
+    def self.jwks_cache
+      @@jwks_cache
+    end
+
     def initialize(api_key: nil, base_url: nil, logger: nil, ssl_verify: true,
                    connection: nil)
-      @jwks_fetched_at = nil
-
       if connection # Inject a Faraday::Connection for testing or full control over Faraday
         @conn = connection
         return
@@ -170,17 +176,9 @@ def decode_token(token)
     # `timeout` argument.
     def verify_token(token, force_refresh_jwks: false, algorithms: ['RS256'], timeout: 5)
       jwk_loader = ->(options) do
-        @cached_jwks = nil if options[:invalidate] || force_refresh_jwks
-        @cached_jwks = nil if @jwks_fetched_at && Time.now.to_i - @jwks_fetched_at > JWKS_CACHE_LIFETIME
-
-        @cached_jwks ||= begin
-          keys = jwks.all["keys"]
-          @jwks_fetched_at = Time.now.to_i
-
-          # JWT.decode requires that the 'keys' key in the Hash is a symbol (as
-          # opposed to a string which our SDK returns by default)
-          { keys: keys }
-        end
+        # JWT.decode requires that the 'keys' key in the Hash is a symbol (as
+        # opposed to a string which our SDK returns by default)
+        { keys: SDK.jwks_cache.fetch(self, kid_not_found: (options[:invalidate] || options[:kid_not_found]), force_refresh: force_refresh_jwks) }
       end
 
       JWT.decode(token, nil, true, algorithms: algorithms, jwks: jwk_loader).first

diff --git a/test/sdk_test.rb b/test/sdk_test.rb
@@ -63,51 +63,51 @@ def test_verify_token_jwks_cache
     sdk = ::Clerk::SDK.new(connection: conn)
 
     Timecop.freeze(TIME_WHEN_JWT_IS_VALID) do
-      sdk.verify_token(json_fixture("jwt_valid"))
-      assert_equal jwks_endpoint_hits, 1
+      sdk.verify_token(json_fixture("jwt_valid"), force_refresh_jwks: true)
+      assert_equal 1, jwks_endpoint_hits
 
       sdk.verify_token(json_fixture("jwt_valid"))
-      assert_equal jwks_endpoint_hits, 1
+      assert_equal 1, jwks_endpoint_hits
 
       sdk.verify_token(json_fixture("jwt_valid"), force_refresh_jwks: true)
-      assert_equal jwks_endpoint_hits, 2
+      assert_equal 2, jwks_endpoint_hits
 
       sdk.verify_token(json_fixture("jwt_valid"))
-      assert_equal jwks_endpoint_hits, 2
+      assert_equal 2, jwks_endpoint_hits
     end
 
+    # cache expired
     Timecop.freeze(TIME_WHEN_JWT_IS_VALID+Clerk::SDK::JWKS_CACHE_LIFETIME+1) do
       begin
         sdk.verify_token(json_fixture("jwt_valid"))
       rescue JWT::ExpiredSignature
         # we know the token is going to be expired in this travelled to time
       end
 
-      assert_equal jwks_endpoint_hits, 3
+      assert_equal 3, jwks_endpoint_hits
     end
 
-
-    # assert that if the last server API request to the JWKS endpoint returned
-    # an error, then we're skipping the cache in the next request
+    # assert that if the server (BAPI) returns an error to the JWKS request
+    # (i.e. BAPI was down momentarily) and therefore the received token couldn't
+    # be verified, then we're going to try one more time
     jwks_endpoint_hits = 0
     conn = Faraday.new do |faraday|
       faraday.adapter :test do |stub|
         stub.get("/jwks") do
           jwks_endpoint_hits += 1
-          json_404
+          jwks_endpoint_hits < 2 ? json_404 : json_ok("jwks")
         end
       end
     end
 
     sdk = ::Clerk::SDK.new(connection: conn)
     Timecop.freeze(TIME_WHEN_JWT_IS_VALID) do
-      begin
-        sdk.verify_token(json_fixture("jwt_valid")) rescue Clerk::Errors::Fatal
-        sdk.verify_token(json_fixture("jwt_valid")) rescue Clerk::Errors::Fatal
-        sdk.verify_token(json_fixture("jwt_valid")) rescue Clerk::Errors::Fatal
-      end
+      claims = sdk.verify_token(json_fixture("jwt_valid"), force_refresh_jwks: true)
+      assert_equal 2, jwks_endpoint_hits
+      assert_equal "[email protected]", claims["email"]
 
-      assert_equal 3, jwks_endpoint_hits
+      sdk.verify_token(json_fixture("jwt_valid"))
+      assert_equal 2, jwks_endpoint_hits # cached response
     end
   end
 end