Deprecate UserSessionCrossDCManager

Fixes #31878

Signed-off-by: Pedro Ruivo <pruivo@redhat.com>

docs/documentation/upgrading/topics/changes/changes-26_0_0.adoc

@@ -125,6 +125,11 @@ It used to be difficult to regain access to a {project_name} instance when all a
 
 Consequently, the environment variables `KEYCLOAK_ADMIN` and `KEYCLOAK_ADMIN_PASSWORD` have been deprecated. You should use `KC_BOOTSTRAP_ADMIN_USERNAME` and `KC_BOOTSTRAP_ADMIN_PASSWORD` instead. These are also general options, so they may be specified via the cli or other config sources, for example `--bootstrap-admin-username=admin`. For more information, see the new https://www.keycloak.org/server/bootstrap-admin-recovery[Bootstrap admin and recovery] guide.
 
+= Deprecations in `keycloak-services` module
+
+The class `UserSessionCrossDCManager` is deprecated and planned to be removed in a future version of {project_name}.
+Read the `UserSessionCrossDCManager` Javadoc for the alternative methods to use.
+
 = Identity Providers no longer available from the realm representation
 
 As part of the improvements around the scalability of realms and organizations when they have many identity providers, the realm representation

server-spi/src/main/java/org/keycloak/models/UserSessionProvider.java

@@ -17,15 +17,14 @@
 
 package org.keycloak.models;
 
-import org.keycloak.migration.MigrationModel;
-import org.keycloak.provider.Provider;
-
 import java.util.Collection;
 import java.util.Map;
 import java.util.UUID;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
+import org.keycloak.provider.Provider;
+
 /**
  * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
  * @version $Revision: 1 $
@@ -212,4 +211,21 @@ public interface UserSessionProvider extends Provider {
 
     default void migrate(String modelVersion) {
     }
+
+    /**
+     * Returns the {@link UserSessionModel} if the user session with ID {@code userSessionId} exist, and it has an
+     * {@link AuthenticatedClientSessionModel} from a {@link ClientModel} with ID {@code clientUUID}.
+     * <p>
+     * If the {@link AuthenticatedClientSessionModel} from the client or the {@link UserSessionModel} does not exist,
+     * this method returns {@code null}.
+     *
+     * @param realm         The {@link RealmModel} where the session belongs to.
+     * @param userSessionId The ID of the {@link UserSessionModel}.
+     * @param offline       If {@code true}, it fetches an offline session and, if {@code false}, an online session.
+     * @param clientUUID    The {@link ClientModel#getId()}.
+     * @return The {@link UserSessionModel} if it has a session from the {@code clientUUID}.
+     */
+    default UserSessionModel getUserSessionIfClientExists(RealmModel realm, String userSessionId, boolean offline, String clientUUID) {
+        return getUserSessionWithPredicate(realm, userSessionId, offline, userSession -> userSession.getAuthenticatedClientSessionByClient(clientUUID) != null);
+    }
 }

services/src/main/java/org/keycloak/protocol/AuthorizationEndpointBase.java

@@ -38,7 +38,6 @@ import org.keycloak.protocol.LoginProtocol.Error;
 import org.keycloak.services.ErrorPageException;
 import org.keycloak.services.managers.AuthenticationManager;
 import org.keycloak.services.managers.AuthenticationSessionManager;
-import org.keycloak.services.managers.UserSessionCrossDCManager;
 import org.keycloak.services.messages.Messages;
 import org.keycloak.services.resources.LoginActionsService;
 import org.keycloak.sessions.AuthenticationSessionModel;
@@ -114,9 +113,7 @@ public abstract class AuthorizationEndpointBase {
             // We cancel login if any authentication action or required action is required
             try {
                 Response challenge = processor.authenticateOnly();
-                if (challenge == null) {
+                if (challenge != null) {
-                    // nothing to do - user is already authenticated;
-                } else {
                     // KEYCLOAK-8043: forward the request with prompt=none to the default provider.
                     if ("true".equals(authSession.getAuthNote(AuthenticationProcessor.FORWARDED_PASSIVE_LOGIN))) {
                         RestartLoginCookie.setRestartCookie(session, authSession);
@@ -127,16 +124,14 @@ public abstract class AuthorizationEndpointBase {
                         return challenge;
                     }
                     else {
-                        Response response = protocol.sendError(authSession, Error.PASSIVE_LOGIN_REQUIRED);
+                        return protocol.sendError(authSession, Error.PASSIVE_LOGIN_REQUIRED);
-                        return response;
                     }
                 }
 
                 AuthenticationManager.setClientScopesInSession(session, authSession);
 
                 if (processor.nextRequiredAction() != null) {
-                    Response response = protocol.sendError(authSession, Error.PASSIVE_INTERACTION_REQUIRED);
+                    return protocol.sendError(authSession, Error.PASSIVE_INTERACTION_REQUIRED);
-                    return response;
                 }
 
             } catch (Exception e) {
@@ -186,8 +181,7 @@ public abstract class AuthorizationEndpointBase {
             logger.debugf("Sent request to authz endpoint. Root authentication session with ID '%s' exists. Client is '%s' . Created new authentication session with tab ID: %s",
                     rootAuthSession.getId(), client.getClientId(), authSession.getTabId());
         } else {
-            UserSessionCrossDCManager userSessionCrossDCManager = new UserSessionCrossDCManager(session);
+            UserSessionModel userSession = manager.getUserSessionFromAuthenticationCookie(realm);
-            UserSessionModel userSession = userSessionCrossDCManager.getUserSessionIfExistsRemotely(manager, realm);
 
             if (userSession != null) {
                 UserModel user = userSession.getUser();

services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java

@@ -86,7 +86,6 @@ import org.keycloak.services.Urls;
 import org.keycloak.services.managers.AuthenticationManager;
 import org.keycloak.services.managers.AuthenticationSessionManager;
 import org.keycloak.services.managers.UserConsentManager;
-import org.keycloak.services.managers.UserSessionCrossDCManager;
 import org.keycloak.services.managers.UserSessionManager;
 import org.keycloak.services.resources.IdentityBrokerService;
 import org.keycloak.services.util.AuthorizationContextUtil;
@@ -193,7 +192,7 @@ public class TokenManager {
 
         // Can theoretically happen in cross-dc environment. Try to see if userSession with our client is available in remoteCache
         if (clientSession == null) {
-            userSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm, userSession.getId(), offline, client.getId());
+            userSession = session.sessions().getUserSessionIfClientExists(realm, userSession.getId(), offline, client.getId());
             if (userSession != null) {
                 clientSession = userSession.getAuthenticatedClientSessionByClient(client.getId());
             } else {

services/src/main/java/org/keycloak/protocol/oidc/endpoints/TokenRevocationEndpoint.java

@@ -28,8 +28,6 @@ import jakarta.ws.rs.Produces;
 import jakarta.ws.rs.core.MediaType;
 import jakarta.ws.rs.core.MultivaluedMap;
 import jakarta.ws.rs.core.Response;
-
-import org.keycloak.http.HttpRequest;
 import org.keycloak.OAuthErrorException;
 import org.keycloak.common.ClientConnection;
 import org.keycloak.common.util.Time;
@@ -38,6 +36,7 @@ import org.keycloak.events.Errors;
 import org.keycloak.events.EventBuilder;
 import org.keycloak.events.EventType;
 import org.keycloak.headers.SecurityHeadersProvider;
+import org.keycloak.http.HttpRequest;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
@@ -53,7 +52,6 @@ import org.keycloak.services.clientpolicy.context.TokenRevokeContext;
 import org.keycloak.services.clientpolicy.context.TokenRevokeResponseContext;
 import org.keycloak.services.cors.Cors;
 import org.keycloak.services.managers.UserConsentManager;
-import org.keycloak.services.managers.UserSessionCrossDCManager;
 import org.keycloak.services.managers.UserSessionManager;
 import org.keycloak.util.TokenUtil;
 
@@ -212,12 +210,11 @@ public class TokenRevocationEndpoint {
         if (token.getSessionState() == null) {
             user = TokenManager.lookupUserFromStatelessToken(session, realm, token);
         } else {
-            UserSessionModel userSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm,
+            var userSessionProvider = session.sessions();
-                    token.getSessionState(), false, client.getId());
+            UserSessionModel userSession = userSessionProvider.getUserSessionIfClientExists(realm, token.getSessionId(), false, client.getId());
 
             if (userSession == null) {
-                userSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm, token.getSessionState(), true,
+                userSession = userSessionProvider.getUserSessionIfClientExists(realm, token.getSessionId(), true, client.getId());
-                        client.getId());
 
                 if (userSession == null) {
                     event.error(Errors.USER_SESSION_NOT_FOUND);

services/src/main/java/org/keycloak/protocol/oidc/grants/device/DeviceGrantType.java

@@ -17,15 +17,18 @@
 
 package org.keycloak.protocol.oidc.grants.device;
 
-import static org.keycloak.protocol.oidc.OIDCLoginProtocolService.tokenServiceBaseUrl;
+import java.net.URI;
+import java.util.Map;
 
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.UriBuilder;
+import jakarta.ws.rs.core.UriInfo;
 import org.keycloak.OAuth2Constants;
 import org.keycloak.OAuthErrorException;
-import org.keycloak.common.Profile;
 import org.keycloak.events.Details;
 import org.keycloak.events.Errors;
-import org.keycloak.events.EventType;
 import org.keycloak.events.EventBuilder;
+import org.keycloak.events.EventType;
 import org.keycloak.models.AuthenticatedClientSessionModel;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.ClientSessionContext;
@@ -38,13 +41,11 @@ import org.keycloak.models.RealmModel;
 import org.keycloak.models.SingleUseObjectProvider;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
-import org.keycloak.provider.EnvironmentDependentProviderFactory;
 import org.keycloak.protocol.LoginProtocol;
 import org.keycloak.protocol.oidc.OIDCLoginProtocol;
 import org.keycloak.protocol.oidc.OIDCLoginProtocolService;
 import org.keycloak.protocol.oidc.TokenManager;
 import org.keycloak.protocol.oidc.endpoints.AuthorizationEndpoint;
-import org.keycloak.protocol.oidc.grants.OAuth2GrantType;
 import org.keycloak.protocol.oidc.grants.OAuth2GrantTypeBase;
 import org.keycloak.protocol.oidc.grants.device.clientpolicy.context.DeviceTokenRequestContext;
 import org.keycloak.protocol.oidc.grants.device.clientpolicy.context.DeviceTokenResponseContext;
@@ -54,17 +55,11 @@ import org.keycloak.services.CorsErrorResponseException;
 import org.keycloak.services.ErrorResponseException;
 import org.keycloak.services.clientpolicy.ClientPolicyException;
 import org.keycloak.services.managers.AuthenticationManager;
-import org.keycloak.services.managers.UserSessionCrossDCManager;
 import org.keycloak.services.resources.RealmsResource;
 import org.keycloak.services.util.DefaultClientSessionContext;
 import org.keycloak.sessions.AuthenticationSessionModel;
 
-import jakarta.ws.rs.core.Response;
+import static org.keycloak.protocol.oidc.OIDCLoginProtocolService.tokenServiceBaseUrl;
-import jakarta.ws.rs.core.UriBuilder;
-import jakarta.ws.rs.core.UriInfo;
-
-import java.net.URI;
-import java.util.Map;
 
 /**
  * OAuth 2.0 Device Authorization Grant
@@ -279,11 +274,11 @@ public class DeviceGrantType extends OAuth2GrantTypeBase {
         event.session(userSessionId);
 
         // Retrieve UserSession
-        UserSessionModel userSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm, userSessionId,
+        var userSessionProvider = session.sessions();
-            client.getId());
+        UserSessionModel userSession = userSessionProvider.getUserSessionIfClientExists(realm, userSessionId, false, client.getId());
 
         if (userSession == null) {
-            userSession = session.sessions().getUserSession(realm, userSessionId);
+            userSession = userSessionProvider.getUserSession(realm, userSessionId);
             if (userSession == null) {
                 throw new CorsErrorResponseException(cors, OAuthErrorException.AUTHORIZATION_PENDING,
                     "The authorization request is verified but can not lookup the user session yet",

services/src/main/java/org/keycloak/protocol/oidc/utils/OAuth2CodeParser.java

@@ -29,7 +29,6 @@ import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.SingleUseObjectProvider;
 import org.keycloak.models.UserSessionModel;
-import org.keycloak.services.managers.UserSessionCrossDCManager;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -99,10 +98,11 @@ public class OAuth2CodeParser {
         }
 
         // Retrieve UserSession
-        UserSessionModel userSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm, userSessionId, clientUUID);
+        var userSessionProvider = session.sessions();
+        UserSessionModel userSession = userSessionProvider.getUserSessionIfClientExists(realm, userSessionId, false, clientUUID);
         if (userSession == null) {
             // Needed to track if code is invalid or was already used.
-            userSession = session.sessions().getUserSession(realm, userSessionId);
+            userSession = userSessionProvider.getUserSession(realm, userSessionId);
             if (userSession == null) {
                 return result.illegalCode();
             }
@@ -150,16 +150,6 @@ public class OAuth2CodeParser {
         private boolean isExpiredCode = false;
 
 
-        private ParseResult(String code, OAuth2Code codeData, AuthenticatedClientSessionModel clientSession) {
-            this.code = code;
-            this.codeData = codeData;
-            this.clientSession = clientSession;
-
-            this.isIllegalCode = false;
-            this.isExpiredCode = false;
-        }
-
-
         private ParseResult(String code) {
             this.code = code;
         }

services/src/main/java/org/keycloak/protocol/saml/SamlSessionUtils.java

@@ -27,7 +27,6 @@ import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.protocol.saml.preprocessor.SamlAuthenticationPreprocessor;
-import org.keycloak.services.managers.UserSessionCrossDCManager;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -60,7 +59,7 @@ public class SamlSessionUtils {
 
         String userSessionId = parts[0];
         String clientUUID = parts[1];
-        UserSessionModel userSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm, userSessionId, false, clientUUID);
+        UserSessionModel userSession = session.sessions().getUserSessionIfClientExists(realm, userSessionId, false, clientUUID);
         if (userSession == null) {
             return null;
         }

services/src/main/java/org/keycloak/services/managers/AuthenticationSessionManager.java

@@ -17,6 +17,8 @@
 
 package org.keycloak.services.managers;
 
+import java.util.Objects;
+
 import org.jboss.logging.Logger;
 import org.keycloak.common.util.Time;
 import org.keycloak.cookie.CookieProvider;
@@ -26,12 +28,15 @@ 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.utils.SessionExpiration;
 import org.keycloak.protocol.RestartLoginCookie;
 import org.keycloak.sessions.AuthenticationSessionModel;
 import org.keycloak.sessions.RootAuthenticationSessionModel;
 import org.keycloak.sessions.StickySessionEncoderProvider;
 
+import static org.keycloak.services.managers.AuthenticationManager.authenticateIdentityCookie;
+
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -233,7 +238,7 @@ public class AuthenticationSessionManager {
 
     // Check to see if we already have authenticationSession with same ID
     public UserSessionModel getUserSession(AuthenticationSessionModel authSession) {
-        return session.sessions().getUserSession(authSession.getRealm(), authSession.getParentSession().getId());
+        return getUserSessionProvider().getUserSession(authSession.getRealm(), authSession.getParentSession().getId());
     }
 
 
@@ -242,4 +247,43 @@ public class AuthenticationSessionManager {
         RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().getRootAuthenticationSession(realm, authSessionId);
         return rootAuthSession==null ? null : rootAuthSession.getAuthenticationSession(client, tabId);
     }
+
+    public UserSessionModel getUserSessionFromAuthenticationCookie(RealmModel realm) {
+        String oldEncodedId = getAuthSessionCookies(realm);
+
+        if (oldEncodedId == null) {
+            // ideally, we should not rely on auth session id to retrieve user sessions
+            // in case the auth session was removed, we fall back to the identity cookie
+            // we are here doing the user session lookup twice, however the second lookup is going to make sure the
+            // session exists in remote caches
+            AuthenticationManager.AuthResult authResult = authenticateIdentityCookie(session, realm, true);
+
+            if (authResult != null && authResult.getSession() != null) {
+                oldEncodedId = authResult.getSession().getId();
+            } else {
+                return null;
+            }
+        }
+
+        AuthSessionId authSessionId = decodeAuthSessionId(oldEncodedId);
+        String sessionId = authSessionId.getDecodedId();
+
+        // TODO: remove this code once InfinispanUserSessionProvider is removed or no longer using any remote caches, as other implementations don't need this call.
+        // This will remove userSession "locally" if it doesn't exist on remoteCache
+        var userSessionProvider = getUserSessionProvider();
+        userSessionProvider.getUserSessionWithPredicate(realm, sessionId, false, Objects::isNull);
+
+        UserSessionModel userSession = userSessionProvider.getUserSession(realm, sessionId);
+
+        if (userSession != null) {
+            reencodeAuthSessionCookie(oldEncodedId, authSessionId, realm);
+            return userSession;
+        } else {
+            return null;
+        }
+    }
+
+    private UserSessionProvider getUserSessionProvider() {
+        return session.sessions();
+    }
 }

services/src/main/java/org/keycloak/services/managers/UserSessionCrossDCManager.java

@@ -17,22 +17,20 @@
 
 package org.keycloak.services.managers;
 
-import org.jboss.logging.Logger;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserSessionModel;
-import org.keycloak.models.AuthenticatedClientSessionModel;
+import org.keycloak.models.UserSessionProvider;
-import org.keycloak.models.ImpersonationSessionNote;
+import org.keycloak.services.util.UserSessionUtil;
-
-import static org.keycloak.services.managers.AuthenticationManager.authenticateIdentityCookie;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ *
+ * @deprecated To be removed without replacement. Check the methods documentation for alternatives.
  */
+@Deprecated(since = "26", forRemoval = true)
 public class UserSessionCrossDCManager {
 
-    private static final Logger logger = Logger.getLogger(UserSessionCrossDCManager.class);
-
     private final KeycloakSession kcSession;
 
     public UserSessionCrossDCManager(KeycloakSession session) {
@@ -41,59 +39,43 @@ public class UserSessionCrossDCManager {
 
 
     // get userSession if it has "authenticatedClientSession" of specified client attached to it. Otherwise download it from remoteCache
+
+    /**
+     * @deprecated To be removed in Keycloak 27+. Use
+     * {@link UserSessionProvider#getUserSessionIfClientExists(RealmModel, String, boolean, String)}
+     */
+    @Deprecated(since = "26", forRemoval = true)
     public UserSessionModel getUserSessionWithClient(RealmModel realm, String id, boolean offline, String clientUUID) {
-        return kcSession.sessions().getUserSessionWithPredicate(realm, id, offline, userSession -> userSession.getAuthenticatedClientSessionByClient(clientUUID) != null);
+        return kcSession.sessions().getUserSessionIfClientExists(realm, id, offline, clientUUID);
     }
 
+    /**
+     * @deprecated To be removed in Keycloak 27+. Use
+     * {@link UserSessionUtil#getUserSessionWithImpersonatorClient(KeycloakSession, RealmModel, String, boolean, String)}
+     */
+    @Deprecated(since = "26", forRemoval = true)
     public UserSessionModel getUserSessionWithImpersonatorClient(RealmModel realm, String id, boolean offline, String clientUUID) {
-       return kcSession.sessions().getUserSessionWithPredicate(realm, id, offline, userSession -> clientUUID.equals(userSession.getNote(ImpersonationSessionNote.IMPERSONATOR_CLIENT.toString())));
+        return UserSessionUtil.getUserSessionWithImpersonatorClient(kcSession, realm, id, offline, clientUUID);
     }
 
     // get userSession if it has "authenticatedClientSession" of specified client attached to it. Otherwise download it from remoteCache
     // TODO Probably remove this method once AuthenticatedClientSession.getAction is removed and information is moved to OAuth code JWT instead
+    /**
+     * @deprecated To be removed in Keycloak 27+. Use
+     * {@link UserSessionProvider#getUserSessionIfClientExists(RealmModel, String, boolean, String)}
+     */
+    @Deprecated(since = "26", forRemoval = true)
     public UserSessionModel getUserSessionWithClient(RealmModel realm, String id, String clientUUID) {
-
+        return getUserSessionWithClient(realm, id, false, clientUUID);
-        return kcSession.sessions().getUserSessionWithPredicate(realm, id, false, (UserSessionModel userSession) -> {
-
-            AuthenticatedClientSessionModel authSessions = userSession.getAuthenticatedClientSessionByClient(clientUUID);
-            return authSessions != null;
-
-        });
     }
 
-
     // Just check if userSession also exists on remoteCache. It can happen that logout happened on 2nd DC and userSession is already removed on remoteCache and this DC wasn't yet notified
+    /**
+     * @deprecated To be removed in Keycloak 27+. Use
+     * {@link AuthenticationSessionManager#getUserSessionFromAuthenticationCookie(RealmModel)}
+     */
+    @Deprecated(since = "26", forRemoval = true)
     public UserSessionModel getUserSessionIfExistsRemotely(AuthenticationSessionManager asm, RealmModel realm) {
-        String oldEncodedId = asm.getAuthSessionCookies(realm);
+        return asm.getUserSessionFromAuthenticationCookie(realm);
-
-        if (oldEncodedId == null) {
-            // ideally, we should not rely on auth session id to retrieve user sessions
-            // in case the auth session was removed, we fall back to the identity cookie
-            // we are here doing the user session lookup twice, however the second lookup is going to make sure the
-            // session exists in remote caches
-            AuthenticationManager.AuthResult authResult = authenticateIdentityCookie(kcSession, realm, true);
-
-            if (authResult != null && authResult.getSession() != null) {
-                oldEncodedId = authResult.getSession().getId();
-            } else {
-                return null;
-            }
-        }
-
-        AuthSessionId authSessionId = asm.decodeAuthSessionId(oldEncodedId);
-        String sessionId = authSessionId.getDecodedId();
-
-        // TODO: remove this code once InfinispanUserSessionProvider is removed or no longer using any remote caches, as other implementations don't need this call.
-        // This will remove userSession "locally" if it doesn't exist on remoteCache
-        kcSession.sessions().getUserSessionWithPredicate(realm, sessionId, false, (UserSessionModel userSession2) -> userSession2 == null);
-
-        UserSessionModel userSession = kcSession.sessions().getUserSession(realm, sessionId);
-
-        if (userSession != null) {
-            asm.reencodeAuthSessionCookie(oldEncodedId, authSessionId, realm);
-            return userSession;
-        } else {
-            return null;
-        }
     }
 }

services/src/main/java/org/keycloak/services/util/UserSessionUtil.java

@@ -1,5 +1,7 @@
 package org.keycloak.services.util;
 
+import java.util.Objects;
+
 import org.jboss.logging.Logger;
 import org.keycloak.common.ClientConnection;
 import org.keycloak.common.constants.ServiceAccountConstants;
@@ -7,6 +9,7 @@ import org.keycloak.events.Errors;
 import org.keycloak.events.EventBuilder;
 import org.keycloak.models.AuthenticatedClientSessionModel;
 import org.keycloak.models.ClientModel;
+import org.keycloak.models.ImpersonationSessionNote;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
@@ -17,7 +20,6 @@ import org.keycloak.protocol.oidc.TokenManager;
 import org.keycloak.representations.AccessToken;
 import org.keycloak.services.Urls;
 import org.keycloak.services.managers.AuthenticationManager;
-import org.keycloak.services.managers.UserSessionCrossDCManager;
 import org.keycloak.services.managers.UserSessionManager;
 import org.keycloak.sessions.AuthenticationSessionModel;
 import org.keycloak.sessions.RootAuthenticationSessionModel;
@@ -27,9 +29,6 @@ public class UserSessionUtil {
 
     private static final Logger logger = Logger.getLogger(UserSessionUtil.class);
 
-
-
-
     public static UserSessionModel findValidSession(KeycloakSession session, RealmModel realm, AccessToken token, EventBuilder event, ClientModel client) {
         OAuth2Error error = new OAuth2Error().json(false).realm(realm);
         return findValidSession(session, realm, token, event, client, error);
@@ -41,19 +40,20 @@ public class UserSessionUtil {
             return createTransientSessionForClient(session, realm, token, client, event);
         }
 
-        UserSessionModel userSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm, token.getSessionId(), false, client.getId());
+        var userSessionProvider = session.sessions();
+        UserSessionModel userSession = userSessionProvider.getUserSessionIfClientExists(realm, token.getSessionId(), false, client.getId());
         if (userSession == null) {
             // also try to resolve sessions created during token exchange when the user is impersonated
-            userSession = new UserSessionCrossDCManager(session).getUserSessionWithImpersonatorClient(realm, token.getSessionId(), false, client.getId());
+            userSession = getUserSessionWithImpersonatorClient(session, realm, token.getSessionId(), false, client.getId());
         }
 
-        UserSessionModel offlineUserSession = null;
+        UserSessionModel offlineUserSession;
         if (AuthenticationManager.isSessionValid(realm, userSession)) {
             checkTokenIssuedAt(realm, token, userSession, event, client);
             event.session(userSession);
             return userSession;
         } else {
-            offlineUserSession = new UserSessionCrossDCManager(session).getUserSessionWithClient(realm, token.getSessionId(), true, client.getId());
+            offlineUserSession = userSessionProvider.getUserSessionIfClientExists(realm, token.getSessionId(), true, client.getId());
             if (AuthenticationManager.isSessionValid(realm, offlineUserSession)) {
                 checkTokenIssuedAt(realm, token, offlineUserSession, event, client);
                 event.session(offlineUserSession);
@@ -67,11 +67,7 @@ public class UserSessionUtil {
             throw error.invalidToken("User session not found or doesn't have client attached on it");
         }
 
-        if (userSession != null) {
+        event.session(Objects.requireNonNullElse(userSession, offlineUserSession));
-            event.session(userSession);
-        } else {
-            event.session(offlineUserSession);
-        }
 
         logger.debug("Session expired");
         event.error(Errors.SESSION_EXPIRED);
@@ -116,4 +112,8 @@ public class UserSessionUtil {
             throw error.invalidToken("Stale token");
         }
     }
+
+    public static UserSessionModel getUserSessionWithImpersonatorClient(KeycloakSession session, RealmModel realm, String userSessionId, boolean offline, String clientUUID) {
+        return session.sessions().getUserSessionWithPredicate(realm, userSessionId, offline, userSession -> Objects.equals(clientUUID, userSession.getNote(ImpersonationSessionNote.IMPERSONATOR_CLIENT.toString())));
+    }
 }