fix: fix token cache key (#7)

logto-sample/config/application.rb

@@ -24,7 +24,10 @@ class Application < Rails::Application
     # config.time_zone = "Central Time (US & Canada)"
     # config.eager_load_paths << Rails.root.join("extras")
     config.session_store(
-      :cookie_store,
+      # We don't use the cookie store since the session data can be large, and
+      # cookies have a size limit of 4KB. Remember to replace it with other stores in
+      # production.
+      :cache_store,
       key: "_logto_sample",
       secure: Rails.env.production?,
       httponly: true,

logto/lib/logto/client/index.rb

@@ -110,7 +110,7 @@ def handle_sign_in_callback(url:)
     error_description = query_params[LogtoCore::QUERY_KEY[:error_description]]
     raise LogtoError::ServerCallbackError, "Error: #{error}, Description: #{error_description}" if error
 
-    current_session = SignInSession.new(@storage.get(STORAGE_KEY[:sign_in_session]))
+    current_session = data.is_a?(SignInSession) ? data : SignInSession.new(**data)
     # A loose URI check here
     raise LogtoError::SessionMismatchError, "Redirect URI mismatch" unless url.start_with?(current_session.redirect_uri)
     raise LogtoError::SessionMismatchError, "No state found in query parameters" unless query_params[LogtoCore::QUERY_KEY[:state]]
@@ -126,7 +126,7 @@ def handle_sign_in_callback(url:)
     )
 
     verify_jwt(token: token_response[:id_token])
-    handle_token_response(token_response)
+    handle_token_response(token_response, resource: nil)
     clear_sign_in_session
 
     @navigate.call(current_session.post_redirect_uri)
@@ -200,7 +200,7 @@ def access_token(resource: nil, organization_id: nil)
       resource: resource,
       organization_id: organization_id
     )
-    handle_token_response(token_response)
+    handle_token_response(token_response, resource: resource, organization_id: organization_id)
     token_response[:access_token]
   end
 
@@ -253,13 +253,13 @@ def clear_all_tokens
 
   protected
 
-  def handle_token_response(response)
+  def handle_token_response(response, resource:, organization_id: nil)
     raise ArgumentError, "Response must be a TokenResponse" unless response.is_a?(LogtoCore::TokenResponse)
     response[:refresh_token] && save_refresh_token(response[:refresh_token])
     response[:id_token] && save_id_token(response[:id_token])
     # The response should have access token
     save_access_token(
-      key: LogtoUtils.build_access_token_key(resource: nil),
+      key: LogtoUtils.build_access_token_key(resource: resource, organization_id: organization_id),
       token: LogtoCore::AccessToken.new(
         token: response[:access_token],
         scope: response[:scope],

logto/lib/logto/core/utils.rb

@@ -41,6 +41,6 @@ def self.generate_state
   end
 
   def self.build_access_token_key(resource:, organization_id: nil)
-    "#{organization_id ? "##{organization_id}" : ""}:#{resource || "openid"}"
+    "#{organization_id ? "##{organization_id}" : ""}:#{resource || "default"}"
   end
 end

logto/spec/client/index_spec.rb

@@ -113,6 +113,23 @@ def remove(key)
   end
 
   describe "#handle_sign_in_callback" do
+    let(:stub_request_proc) do
+      proc do
+        stub_request(:post, "https://example.com/oidc/token")
+          .to_return(
+            status: 200,
+            body: {
+              access_token: "access_token",
+              refresh_token: "refresh_token",
+              id_token: "id_token",
+              expires_in: 3600,
+              scope: "openid"
+            }.to_json,
+            headers: {"Content-Type" => "application/json"}
+          )
+      end
+    end
+
     it "raises an error when no sign-in session is found" do
       expect { basic_client.handle_sign_in_callback(url: "https://example.com/callback?") }.to raise_error(LogtoError::SessionNotFoundError)
     end
@@ -142,21 +159,18 @@ def remove(key)
       expect { basic_client.handle_sign_in_callback(url: "https://example.com/callback?state=state&error=error") }.to raise_error(LogtoError::ServerCallbackError)
     end
 
+    it "recoginizes the `SignInSession` struct stored in the storage" do
+      storage.set(LogtoClient::STORAGE_KEY[:sign_in_session], LogtoClient::SignInSession.new(redirect_uri: "https://example.com/callback", state: "state"))
+      allow_any_instance_of(LogtoClient).to receive(:verify_jwt).and_return(true)
+      stub_request_proc.call
+      basic_client.handle_sign_in_callback(url: "https://example.com/callback?state=state&code=code")
+      expect(storage.get(LogtoClient::STORAGE_KEY[:sign_in_session])).to be_nil
+    end
+
     it "fetches the token by the authorization code" do
       storage.set(LogtoClient::STORAGE_KEY[:sign_in_session], {redirect_uri: "https://example.com/callback", state: "state", code_verifier: "code_verifier"})
       allow_any_instance_of(LogtoClient).to receive(:verify_jwt).and_return(true)
-      stub_request(:post, "https://example.com/oidc/token")
-        .to_return(
-          status: 200,
-          body: {
-            access_token: "access_token",
-            refresh_token: "refresh_token",
-            id_token: "id_token",
-            expires_in: 3600,
-            scope: "openid"
-          }.to_json,
-          headers: {"Content-Type" => "application/json"}
-        )
+      stub_request_proc.call
       basic_client.handle_sign_in_callback(url: "https://example.com/callback?state=state&code=code")
       expect(basic_client.access_token).to eq("access_token")
       expect(basic_client.refresh_token).to eq("refresh_token")
@@ -209,14 +223,25 @@ def remove(key)
   end
 
   describe "#access_token" do
+    let(:stub_request_proc) do
+      proc do
+        stub_request(:post, "https://example.com/oidc/token")
+          .to_return(
+            status: 200,
+            body: {access_token: "new_access_token", expires_in: 3600, scope: "openid"}.to_json,
+            headers: {"Content-Type" => "application/json"}
+          )
+      end
+    end
+
     it "raises an error when not authenticated" do
       expect { basic_client.access_token(resource: nil) }.to raise_error(LogtoError::NotAuthenticatedError)
     end
 
     it "returns the access token directly when it's stored and not expired" do
       allow_any_instance_of(LogtoClient).to receive(:is_authenticated?).and_return(true)
       storage.set(LogtoClient::STORAGE_KEY[:access_token_map], {
-        ":openid" => LogtoCore::AccessToken.new(
+        ":default" => LogtoCore::AccessToken.new(
           token: "access_token",
           scope: "openid",
           expires_at: Time.now + 3600
@@ -229,22 +254,48 @@ def remove(key)
     it "uses the refresh token to fetch a new access token when the stored one is expired" do
       allow_any_instance_of(LogtoClient).to receive(:is_authenticated?).and_return(true)
       storage.set(LogtoClient::STORAGE_KEY[:access_token_map], {
-        ":openid" => LogtoCore::AccessToken.new(
+        ":default" => LogtoCore::AccessToken.new(
           token: "access_token",
           scope: "openid",
           expires_at: Time.now + 5 # Test leeway
         )
       })
       storage.set(LogtoClient::STORAGE_KEY[:refresh_token], "refresh_token")
+      stub_request_proc.call
+      expect(basic_client.access_token(resource: nil)).to eq("new_access_token")
+    end
 
-      stub_request(:post, "https://example.com/oidc/token")
-        .to_return(
-          status: 200,
-          body: {access_token: "new_access_token", expires_in: 3600, scope: "openid"}.to_json,
-          headers: {"Content-Type" => "application/json"}
-        )
+    it "uses the refresh token to fetch a new access token when no access token is stored (resource)" do
+      allow_any_instance_of(LogtoClient).to receive(:is_authenticated?).and_return(true)
+      storage.set(LogtoClient::STORAGE_KEY[:refresh_token], "refresh_token")
+      stub_request_proc.call
 
-      expect(basic_client.access_token(resource: nil)).to eq("new_access_token")
+      expect(basic_client.access_token(resource: "https://example.com/")).to eq("new_access_token")
+      token_map = storage.get(LogtoClient::STORAGE_KEY[:access_token_map])
+      expect(token_map).to include(":https://example.com/")
+      expect(token_map[":https://example.com/"].token).to eq("new_access_token")
+    end
+
+    it "uses the refresh token to fetch a new access token when no access token is stored (organization ID)" do
+      allow_any_instance_of(LogtoClient).to receive(:is_authenticated?).and_return(true)
+      storage.set(LogtoClient::STORAGE_KEY[:refresh_token], "refresh_token")
+      stub_request_proc.call
+
+      expect(basic_client.access_token(organization_id: "123")).to eq("new_access_token")
+      token_map = storage.get(LogtoClient::STORAGE_KEY[:access_token_map])
+      expect(token_map).to include("#123:default")
+      expect(token_map["#123:default"].token).to eq("new_access_token")
+    end
+
+    it "uses the refresh token to fetch a new access token when no access token is stored (resource and organization ID)" do
+      allow_any_instance_of(LogtoClient).to receive(:is_authenticated?).and_return(true)
+      storage.set(LogtoClient::STORAGE_KEY[:refresh_token], "refresh_token")
+      stub_request_proc.call
+
+      expect(basic_client.access_token(resource: "https://example.com/", organization_id: "123")).to eq("new_access_token")
+      token_map = storage.get(LogtoClient::STORAGE_KEY[:access_token_map])
+      expect(token_map).to include("#123:https://example.com/")
+      expect(token_map["#123:https://example.com/"].token).to eq("new_access_token")
     end
   end
 

logto/spec/core/utils_spec.rb

@@ -36,7 +36,7 @@
 
   describe "#build_access_token_key" do
     it "builds a key without resource and organization ID" do
-      expect(LogtoUtils.build_access_token_key(resource: nil)).to eq(":openid")
+      expect(LogtoUtils.build_access_token_key(resource: nil)).to eq(":default")
     end
 
     it "builds a key with resource and organization ID" do
@@ -48,7 +48,7 @@
     end
 
     it "builds a key without resource and with organization ID" do
-      expect(LogtoUtils.build_access_token_key(resource: nil, organization_id: "organization_id")).to eq("#organization_id:openid")
+      expect(LogtoUtils.build_access_token_key(resource: nil, organization_id: "organization_id")).to eq("#organization_id:default")
     end
   end
 end