From 202e9e8479b92a9594858ef3be0763f7b6d29a5f Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Tue, 3 Jun 2025 10:35:35 +0200 Subject: [PATCH] Cache the client session if it is missing from the cache (#40156) Closes #39785 Signed-off-by: Alexander Schwartz --- .../PersistentUserSessionProvider.java | 2 +- ...onPersistentChangelogBasedTransaction.java | 19 +++++--- .../testsuite/oauth/RefreshTokenTest.java | 46 +++++++++++++++++++ 3 files changed, 59 insertions(+), 8 deletions(-) diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/PersistentUserSessionProvider.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/PersistentUserSessionProvider.java index 47173496359..a07c217845d 100755 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/PersistentUserSessionProvider.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/PersistentUserSessionProvider.java @@ -811,7 +811,7 @@ public class PersistentUserSessionProvider implements UserSessionProvider, Sessi return sessionsById.entrySet().stream().findFirst().map(Map.Entry::getValue).orElse(null); } - private Map> importSessionsWithExpiration(Map> sessionsById, + public Map> importSessionsWithExpiration(Map> sessionsById, BasicCache> cache, SessionFunction lifespanMsCalculator, SessionFunction maxIdleTimeMsCalculator) { return sessionsById.entrySet().stream().map(entry -> { diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/ClientSessionPersistentChangelogBasedTransaction.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/ClientSessionPersistentChangelogBasedTransaction.java index 354d10ded9f..0b05bc0618c 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/ClientSessionPersistentChangelogBasedTransaction.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/ClientSessionPersistentChangelogBasedTransaction.java @@ -24,6 +24,7 @@ import org.keycloak.models.ClientModel; import org.keycloak.models.KeycloakSession; import org.keycloak.models.RealmModel; import org.keycloak.models.UserSessionModel; +import org.keycloak.models.UserSessionProvider; import org.keycloak.models.session.UserSessionPersisterProvider; import org.keycloak.models.sessions.infinispan.PersistentUserSessionProvider; import org.keycloak.models.sessions.infinispan.SessionFunction; @@ -32,8 +33,8 @@ import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessi import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionStore; import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity; import org.keycloak.models.sessions.infinispan.remotestore.RemoteCacheInvoker; -import org.keycloak.models.sessions.infinispan.util.SessionTimeouts; +import java.util.Map; import java.util.UUID; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.ConcurrentHashMap; @@ -158,13 +159,17 @@ public class ClientSessionPersistentChangelogBasedTransaction extends Persistent entity.setTimestamp(userSession.getLastSessionRefresh()); } - if (getMaxIdleMsLoader(offline).apply(realm, client, entity) == SessionTimeouts.ENTRY_EXPIRED_FLAG - || getLifespanMsLoader(offline).apply(realm, client, entity) == SessionTimeouts.ENTRY_EXPIRED_FLAG) { + final UUID clientSessionId = entity.getId(); + + SessionEntityWrapper wrapper = new SessionEntityWrapper<>(entity); + Map> imported = ((PersistentUserSessionProvider) kcSession.getProvider(UserSessionProvider.class)).importSessionsWithExpiration(Map.of(clientSessionId, wrapper), getCache(offline), + getLifespanMsLoader(offline), + getMaxIdleMsLoader(offline)); + + if (imported.isEmpty()) { return null; } - final UUID clientSessionId = entity.getId(); - SessionUpdateTask createClientSessionTask = Tasks.addIfAbsentSync(); this.addTask(entity.getId(), createClientSessionTask, entity, UserSessionModel.SessionPersistenceState.PERSISTENT); @@ -176,10 +181,10 @@ public class ClientSessionPersistentChangelogBasedTransaction extends Persistent AuthenticatedClientSessionStore clientSessions = sessionToImportInto.getEntity().getAuthenticatedClientSessions(); clientSessions.put(client.getId(), clientSessionId); - SessionUpdateTask registerClientSessionTask = new RegisterClientSessionTask(client.getId(), clientSessionId, offline); + SessionUpdateTask registerClientSessionTask = new RegisterClientSessionTask(client.getId(), clientSessionId, offline); userSessionTx.addTask(sessionToImportInto.getId(), registerClientSessionTask); - return new SessionEntityWrapper<>(entity); + return wrapper; } public static class RegisterClientSessionTask implements PersistentSessionUpdateTask { diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/oauth/RefreshTokenTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/oauth/RefreshTokenTest.java index 7cf00d25421..0444733765c 100755 --- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/oauth/RefreshTokenTest.java +++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/oauth/RefreshTokenTest.java @@ -345,6 +345,52 @@ public class RefreshTokenTest extends AbstractKeycloakTest { } } + + @Test + public void refreshingTokenLoadsSessionIntoCache() { + + ProfileAssume.assumeFeatureEnabled(Profile.Feature.PERSISTENT_USER_SESSIONS); + + oauth.doLogin("test-user@localhost", "password"); + + String code = oauth.parseLoginResponse().getCode(); + + AccessTokenResponse response = oauth.doAccessTokenRequest(code); + String refreshTokenString = response.getRefreshToken(); + + // Test when neither client nor user session is in the cache + testingClient.server().run(session -> { + session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME).clear(); + session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME).clear(); + }); + + response = oauth.doRefreshTokenRequest(refreshTokenString); + Assert.assertEquals(200, response.getStatusCode()); + + testingClient.server().run(session -> { + assertThat(session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME).size(), + greaterThan(0)); + assertThat(session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME).size(), + greaterThan(0)); + }); + + // Test is only the client session is missing + testingClient.server().run(session -> { + session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME).clear(); + }); + + response = oauth.doRefreshTokenRequest(refreshTokenString); + Assert.assertEquals(200, response.getStatusCode()); + + testingClient.server().run(session -> { + assertThat(session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME).size(), + greaterThan(0)); + assertThat(session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME).size(), + greaterThan(0)); + }); + + } + @Test public void refreshTokenWithAccessToken() { oauth.doLogin("test-user@localhost", "password");