From 07c92c85cbf685acd4bd521859df6c4e3e3fe0bb Mon Sep 17 00:00:00 2001 From: Pedro Ruivo Date: Mon, 5 Aug 2024 14:23:35 +0100 Subject: [PATCH] Drop AuthenticatedClientSessionStore from user sessions New entities for client and user sessions, more query friendly. The client sessions are found using query instead of storing them in the user session entity. Remove of sessions by its field is done based on queries. Closes #30934 Signed-off-by: Pedro Ruivo --- .../high-availability/bblocks-multi-site.adoc | 5 - .../deploy-infinispan-kubernetes-crossdc.adoc | 2 +- .../marshalling/KeycloakModelSchema.java | 6 + .../org/keycloak/marshalling/Marshalling.java | 5 + .../ByRealmIdQueryConditionalRemover.java | 4 +- .../ClientSessionQueryConditionalRemover.java | 114 ++++ .../query/MultipleConditionQueryRemover.java | 100 ++++ .../query/QueryBasedConditionalRemover.java | 30 +- .../UserSessionQueryConditionalRemover.java | 90 +++ .../remote/updater/UpdaterFactory.java | 19 +- .../AuthenticatedClientSessionUpdater.java | 84 ++- .../loginfailures/LoginFailuresUpdater.java | 9 +- .../user/ClientSessionMappingAdapter.java | 144 ----- .../updater/user/ClientSessionProvider.java | 59 -- .../updater/user/UserSessionUpdater.java | 83 +-- .../infinispan/entities/ClientSessionKey.java | 31 + ...emoteAuthenticatedClientSessionEntity.java | 194 +++++++ .../entities/RemoteUserSessionEntity.java | 213 +++++++ .../query/ClientSessionQueries.java | 76 +++ .../infinispan/query/QueryHelper.java | 207 +++++++ .../infinispan/query/UserSessionQueries.java | 68 +++ ...nAuthenticationSessionProviderFactory.java | 11 +- ...RemoteUserLoginFailureProviderFactory.java | 21 +- .../remote/RemoteUserSessionProvider.java | 383 ++++++------- .../RemoteUserSessionProviderFactory.java | 33 +- .../ClientSessionChangeLogTransaction.java | 40 +- .../transaction/RemoteCacheAndExecutor.java | 39 -- .../RemoteChangeLogTransaction.java | 8 + .../UseSessionChangeLogTransaction.java | 21 +- .../transaction/UserSessionTransaction.java | 26 +- .../infinispan/util/SessionTimeouts.java | 78 ++- .../infinispan/CacheManagerFactory.java | 35 +- .../infinispan/InfinispanIckleQueryTest.java | 538 ++++++++++++++++++ 33 files changed, 2100 insertions(+), 676 deletions(-) create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ClientSessionQueryConditionalRemover.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/MultipleConditionQueryRemover.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/UserSessionQueryConditionalRemover.java delete mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionMappingAdapter.java delete mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionProvider.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionKey.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteAuthenticatedClientSessionEntity.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteUserSessionEntity.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/ClientSessionQueries.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/QueryHelper.java create mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/UserSessionQueries.java delete mode 100644 model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteCacheAndExecutor.java create mode 100644 testsuite/model/src/test/java/org/keycloak/testsuite/model/infinispan/InfinispanIckleQueryTest.java diff --git a/docs/guides/high-availability/bblocks-multi-site.adoc b/docs/guides/high-availability/bblocks-multi-site.adoc index 5e5c0c8ce55..c8c978adf4b 100644 --- a/docs/guides/high-availability/bblocks-multi-site.adoc +++ b/docs/guides/high-availability/bblocks-multi-site.adoc @@ -48,11 +48,6 @@ A deployment of {jdgserver_name} that leverages the {jdgserver_name}'s Cross-DC *Not considered:* Direct interconnections between the Kubernetes clusters on the network layer. It might be considered in the future. -[IMPORTANT] -==== -Only {jdgserver_name} server versions 15.0.0 or greater are supported in multi-site deployments. -==== - == {project_name} A clustered deployment of {project_name} in each site, connected to an external {jdgserver_name}. diff --git a/docs/guides/high-availability/deploy-infinispan-kubernetes-crossdc.adoc b/docs/guides/high-availability/deploy-infinispan-kubernetes-crossdc.adoc index 34a448959be..d88e24f7200 100644 --- a/docs/guides/high-availability/deploy-infinispan-kubernetes-crossdc.adoc +++ b/docs/guides/high-availability/deploy-infinispan-kubernetes-crossdc.adoc @@ -19,7 +19,7 @@ See the <@links.ha id="introduction" /> {section} for an overview. [IMPORTANT] ==== -Only {jdgserver_name} server versions 15.0.0 or greater are supported for external {jdgserver_name} deployments. +Only versions based on Infinispan version ${properties["infinispan.version"]} or more recent patch releases are supported for external {jdgserver_name} deployments. ==== == Architecture diff --git a/model/infinispan/src/main/java/org/keycloak/marshalling/KeycloakModelSchema.java b/model/infinispan/src/main/java/org/keycloak/marshalling/KeycloakModelSchema.java index 36a8b0368ff..4870cd2012c 100644 --- a/model/infinispan/src/main/java/org/keycloak/marshalling/KeycloakModelSchema.java +++ b/model/infinispan/src/main/java/org/keycloak/marshalling/KeycloakModelSchema.java @@ -82,8 +82,11 @@ import org.keycloak.models.sessions.infinispan.changes.sessions.SessionData; import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionEntity; import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionStore; import org.keycloak.models.sessions.infinispan.entities.AuthenticationSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity; import org.keycloak.models.sessions.infinispan.entities.LoginFailureKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; import org.keycloak.models.sessions.infinispan.entities.RootAuthenticationSessionEntity; import org.keycloak.models.sessions.infinispan.entities.SingleUseObjectValueEntity; import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity; @@ -199,8 +202,11 @@ import org.keycloak.storage.managers.UserStorageSyncManager; AuthenticatedClientSessionStore.class, AuthenticatedClientSessionEntity.class, AuthenticationSessionEntity.class, + ClientSessionKey.class, LoginFailureEntity.class, LoginFailureKey.class, + RemoteAuthenticatedClientSessionEntity.class, + RemoteUserSessionEntity.class, RootAuthenticationSessionEntity.class, SingleUseObjectValueEntity.class, UserSessionEntity.class, diff --git a/model/infinispan/src/main/java/org/keycloak/marshalling/Marshalling.java b/model/infinispan/src/main/java/org/keycloak/marshalling/Marshalling.java index 6607845fb7b..450a657cc58 100644 --- a/model/infinispan/src/main/java/org/keycloak/marshalling/Marshalling.java +++ b/model/infinispan/src/main/java/org/keycloak/marshalling/Marshalling.java @@ -149,6 +149,11 @@ public final class Marshalling { public static final int CACHE_KEY_INVALIDATION_EVENT = 65603; public static final int CLEAR_CACHE_EVENT = 65604; + public static final int REMOTE_USER_SESSION_ENTITY = 65605; + + public static final int CLIENT_SESSION_KEY = 65606; + public static final int REMOTE_CLIENT_SESSION_ENTITY = 65607; + public static void configure(GlobalConfigurationBuilder builder) { builder.serialization() .addContextInitializer(KeycloakModelSchema.INSTANCE); diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ByRealmIdQueryConditionalRemover.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ByRealmIdQueryConditionalRemover.java index da09d3ff058..72138ef4368 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ByRealmIdQueryConditionalRemover.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ByRealmIdQueryConditionalRemover.java @@ -21,7 +21,6 @@ import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.concurrent.Executor; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -44,8 +43,7 @@ public class ByRealmIdQueryConditionalRemover extend private final String entity; private final List realms; - public ByRealmIdQueryConditionalRemover(String entity, Executor executor) { - super(executor); + public ByRealmIdQueryConditionalRemover(String entity) { this.entity = entity; this.realms = new ArrayList<>(); } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ClientSessionQueryConditionalRemover.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ClientSessionQueryConditionalRemover.java new file mode 100644 index 00000000000..2a953e79a4e --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/ClientSessionQueryConditionalRemover.java @@ -0,0 +1,114 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.changes.remote.remover.query; + +import java.util.Map; +import java.util.Objects; + +import org.keycloak.models.sessions.infinispan.changes.remote.remover.ConditionalRemover; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; +import org.keycloak.models.sessions.infinispan.query.ClientSessionQueries; + +/** + * A {@link ConditionalRemover} implementation to remove {@link RemoteAuthenticatedClientSessionEntity} based on some + * filters over its state. + *

+ * This implementation uses Infinispan Ickle Queries to perform the removal operation. Indexing is not required. + */ +public class ClientSessionQueryConditionalRemover extends MultipleConditionQueryRemover { + + public ClientSessionQueryConditionalRemover() { + super(); + } + + @Override + String getEntity() { + return ClientSessionQueries.CLIENT_SESSION; + } + + public void removeByUserSessionId(String userSessionId) { + add(new RemoveByUserSession(nextParameter(), userSessionId)); + } + + public void removeByRealmId(String realmId) { + add(new RemoveByRealm(nextParameter(), realmId)); + } + + public void removeByUserId(String realmId, String userId) { + add(new RemoveByUser(nextParameter(), realmId, nextParameter(), userId)); + } + + private record RemoveByUserSession(String userSessionParameter, + String userSessionId) implements RemoveCondition { + + @Override + public String getConditionalClause() { + return "(userSessionId = :%s)".formatted(userSessionParameter); + } + + @Override + public void addParameters(Map parameters) { + parameters.put(userSessionParameter, userSessionId); + } + + @Override + public boolean willRemove(ClientSessionKey key, RemoteAuthenticatedClientSessionEntity value) { + return Objects.equals(value.getUserSessionId(), userSessionId); + } + } + + private record RemoveByRealm(String realmParameter, + String realmId) implements RemoveCondition { + + @Override + public String getConditionalClause() { + return "(realmId = :%s)".formatted(realmParameter); + } + + @Override + public void addParameters(Map parameters) { + parameters.put(realmParameter, realmId); + } + + @Override + public boolean willRemove(ClientSessionKey key, RemoteAuthenticatedClientSessionEntity value) { + return Objects.equals(value.getRealmId(), realmId); + } + } + + private record RemoveByUser(String realmParameter, String realmId, String userParameter, + String userId) implements RemoveCondition { + + @Override + public String getConditionalClause() { + return "(userId = :%s && realmId = :%s)".formatted(userParameter, realmParameter); + } + + @Override + public void addParameters(Map parameters) { + parameters.put(realmParameter, realmId); + parameters.put(userParameter, userId); + } + + @Override + public boolean willRemove(ClientSessionKey key, RemoteAuthenticatedClientSessionEntity value) { + return Objects.equals(value.getUserId(), userId) && Objects.equals(value.getRealmId(), realmId); + } + } +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/MultipleConditionQueryRemover.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/MultipleConditionQueryRemover.java new file mode 100644 index 00000000000..43b5c7cd3db --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/MultipleConditionQueryRemover.java @@ -0,0 +1,100 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.changes.remote.remover.query; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import org.infinispan.client.hotrod.RemoteCache; + +/** + * Base class implementing {@link QueryBasedConditionalRemover} and supports multiple remove conditions. + *

+ * The remove condition can be added dynamically and, when the query is executed, they are joined together with an "or" + * operator. + * + * @param The key's type stored in the {@link RemoteCache}. + * @param The value's type stored in the {@link RemoteCache}. + */ +abstract class MultipleConditionQueryRemover extends QueryBasedConditionalRemover { + + private final List> removes; + private int parameterIndex; + + MultipleConditionQueryRemover() { + removes = new ArrayList<>(); + } + + @Override + String getQueryConditions() { + return removes.stream() + .map(RemoveCondition::getConditionalClause) + .collect(Collectors.joining(" || ")); + } + + @Override + Map getQueryParameters() { + Map parameters = new HashMap<>(); + removes.forEach(removeCondition -> removeCondition.addParameters(parameters)); + return parameters; + } + + @Override + boolean isEmpty() { + return removes.isEmpty(); + } + + @Override + public boolean willRemove(K key, V value) { + return !isEmpty() && removes.stream().anyMatch(c -> c.willRemove(key, value)); + } + + /** + * If the query has parameters, use this method to generate a new unique parameter. + */ + String nextParameter() { + return "p" + parameterIndex++; + } + + void add(RemoveCondition condition) { + removes.add(condition); + } + + /** + * A single remove condition. + */ + interface RemoveCondition { + /** + * @return The where clause with parameters. + */ + String getConditionalClause(); + + /** + * Stores this condition parameters value + */ + void addParameters(Map parameters); + + /** + * @return {@code true} if the entry wil be removed by the query. + */ + boolean willRemove(K key, V value); + } +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/QueryBasedConditionalRemover.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/QueryBasedConditionalRemover.java index 271ec69799a..88fc677e1fc 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/QueryBasedConditionalRemover.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/QueryBasedConditionalRemover.java @@ -19,11 +19,10 @@ package org.keycloak.models.sessions.infinispan.changes.remote.remover.query; import java.lang.invoke.MethodHandles; import java.util.Map; -import java.util.Objects; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executor; +import java.util.concurrent.CompletionStage; import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.client.hotrod.impl.query.RemoteQuery; import org.infinispan.commons.util.concurrent.AggregateCompletionStage; import org.jboss.logging.Logger; import org.keycloak.models.sessions.infinispan.changes.remote.remover.ConditionalRemover; @@ -44,30 +43,27 @@ abstract class QueryBasedConditionalRemover implements ConditionalRemover< private static final String QUERY_FMT = "DELETE FROM %s WHERE %s"; - private final Executor executor; - - QueryBasedConditionalRemover(Executor executor) { - this.executor = Objects.requireNonNull(executor); - } - @Override public void executeRemovals(RemoteCache cache, AggregateCompletionStage stage) { if (isEmpty()) { return; } - // TODO replace with async method: https://issues.redhat.com/browse/ISPN-16279 - stage.dependsOn(CompletableFuture.runAsync(() -> executeDeleteStatement(cache), executor)); + stage.dependsOn(executeDeleteStatement(cache)); } - private void executeDeleteStatement(RemoteCache cache) { + private CompletionStage executeDeleteStatement(RemoteCache cache) { + var isTrace = logger.isTraceEnabled(); var deleteStatement = QUERY_FMT.formatted(getEntity(), getQueryConditions()); - if (logger.isTraceEnabled()) { + if (isTrace) { logger.tracef("About to execute delete statement in cache '%s': %s", cache.getName(), deleteStatement); } - var removed = cache.query(deleteStatement) - .setParameters(getQueryParameters()) - .executeStatement(); - logger.debugf("Delete Statement removed %d entries from cache '%s'", removed, cache.getName()); + RemoteQuery query = (RemoteQuery) cache.query(deleteStatement) + .setParameters(getQueryParameters()); + var stage = query.executeStatementAsync(); + if (isTrace) { + return stage.thenAccept(removed -> logger.debugf("Delete Statement removed %d entries from cache '%s'", removed, cache.getName())); + } + return stage; } /** diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/UserSessionQueryConditionalRemover.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/UserSessionQueryConditionalRemover.java new file mode 100644 index 00000000000..29d737dccc5 --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/remover/query/UserSessionQueryConditionalRemover.java @@ -0,0 +1,90 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.changes.remote.remover.query; + +import java.util.Map; +import java.util.Objects; + +import org.keycloak.models.sessions.infinispan.changes.remote.remover.ConditionalRemover; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; +import org.keycloak.models.sessions.infinispan.query.UserSessionQueries; + +/** + * A {@link ConditionalRemover} implementation to remove {@link RemoteUserSessionEntity} based on some filters over its + * state. + *

+ * This implementation uses Infinispan Ickle Queries to perform the removal operation. Indexing is not required. + */ +public class UserSessionQueryConditionalRemover extends MultipleConditionQueryRemover { + + public UserSessionQueryConditionalRemover() { + super(); + } + + @Override + String getEntity() { + return UserSessionQueries.USER_SESSION; + } + + public void removeByRealmId(String realmId) { + add(new RemoveByRealm(nextParameter(), realmId)); + } + + public void removeByUserId(String realmId, String userId) { + add(new RemoveUser(nextParameter(), userId, nextParameter(), realmId)); + } + + private record RemoveUser(String userParameter, String userId, String realmParameter, + String realmId) implements RemoveCondition { + + @Override + public String getConditionalClause() { + return "(userId = :%s && realmId = :%s)".formatted(userParameter, realmParameter); + } + + @Override + public void addParameters(Map parameters) { + parameters.put(userParameter, userId); + parameters.put(realmParameter, realmId); + } + + @Override + public boolean willRemove(String key, RemoteUserSessionEntity value) { + return Objects.equals(value.getUserId(), userId) && Objects.equals(value.getRealmId(), realmId); + } + } + + private record RemoveByRealm(String parameter, + String realmId) implements RemoveCondition { + + @Override + public String getConditionalClause() { + return "(realmId = :%s)".formatted(parameter); + } + + @Override + public void addParameters(Map parameters) { + parameters.put(parameter, realmId); + } + + @Override + public boolean willRemove(String key, RemoteUserSessionEntity value) { + return Objects.equals(realmId, value.getRealmId()); + } + } +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/UpdaterFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/UpdaterFactory.java index d221e8ac7b7..9ef72415868 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/UpdaterFactory.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/UpdaterFactory.java @@ -16,6 +16,8 @@ */ package org.keycloak.models.sessions.infinispan.changes.remote.updater; +import java.util.Objects; + import org.infinispan.client.hotrod.MetadataValue; /** @@ -43,7 +45,21 @@ public interface UpdaterFactory> { * @param entity The Infinispan value. * @return The {@link Updater} to be used when updating the entity state. */ - T wrapFromCache(K key, MetadataValue entity); + default T wrapFromCache(K key, MetadataValue entity) { + Objects.requireNonNull(key); + Objects.requireNonNull(entity); + return wrapFromCache(key, entity.getValue(), entity.getVersion()); + } + + /** + * Wraps an entity read from the Infinispan cache. + * + * @param key The Infinispan key. + * @param value The Infinispan value. + * @param version The entry version. + * @return The {@link Updater} to be used when updating the entity state. + */ + T wrapFromCache(K key, V value, long version); /** * Deletes a entity that was not previous read by the Keycloak transaction. @@ -52,5 +68,4 @@ public interface UpdaterFactory> { * @return The {@link Updater} for a deleted entity. */ T deleted(K key); - } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/client/AuthenticatedClientSessionUpdater.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/client/AuthenticatedClientSessionUpdater.java index 78457c34b3a..671beb2e437 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/client/AuthenticatedClientSessionUpdater.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/client/AuthenticatedClientSessionUpdater.java @@ -22,10 +22,8 @@ import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.UUID; import java.util.function.Consumer; -import org.infinispan.client.hotrod.MetadataValue; import org.infinispan.client.hotrod.RemoteCache; import org.keycloak.models.AuthenticatedClientSessionModel; import org.keycloak.models.ClientModel; @@ -36,26 +34,27 @@ import org.keycloak.models.sessions.infinispan.changes.remote.updater.Expiration import org.keycloak.models.sessions.infinispan.changes.remote.updater.Updater; import org.keycloak.models.sessions.infinispan.changes.remote.updater.UpdaterFactory; import org.keycloak.models.sessions.infinispan.changes.remote.updater.helper.MapUpdater; -import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; import org.keycloak.models.sessions.infinispan.remote.transaction.ClientSessionChangeLogTransaction; import org.keycloak.models.sessions.infinispan.util.SessionTimeouts; /** * An {@link Updater} implementation that keeps track of {@link AuthenticatedClientSessionModel} changes. */ -public class AuthenticatedClientSessionUpdater extends BaseUpdater implements AuthenticatedClientSessionModel { +public class AuthenticatedClientSessionUpdater extends BaseUpdater implements AuthenticatedClientSessionModel { private static final Factory ONLINE = new Factory(false); private static final Factory OFFLINE = new Factory(true); private final MapUpdater notesUpdater; - private final List> changes; + private final List> changes; private final boolean offline; private UserSessionModel userSession; private ClientModel client; private ClientSessionChangeLogTransaction clientTransaction; - private AuthenticatedClientSessionUpdater(UUID cacheKey, AuthenticatedClientSessionEntity cacheValue, long version, boolean offline, UpdaterState initialState) { + private AuthenticatedClientSessionUpdater(ClientSessionKey cacheKey, RemoteAuthenticatedClientSessionEntity cacheValue, long version, boolean offline, UpdaterState initialState) { super(cacheKey, cacheValue, version, initialState); this.offline = offline; if (cacheValue == null) { @@ -74,35 +73,52 @@ public class AuthenticatedClientSessionUpdater extends BaseUpdater factory(boolean offline) { + public static UpdaterFactory factory(boolean offline) { return offline ? OFFLINE : ONLINE; } @Override - public AuthenticatedClientSessionEntity apply(UUID uuid, AuthenticatedClientSessionEntity entity) { + public RemoteAuthenticatedClientSessionEntity apply(ClientSessionKey uuid, RemoteAuthenticatedClientSessionEntity entity) { initNotes(entity); notesUpdater.applyChanges(entity.getNotes()); changes.forEach(change -> change.accept(entity)); + if (isCreated()) { + // The ID generation is not random + // During RefreshTokenTest, the entry is expired in KC but not in the external Infinispan. + // If it happens in production, we need to merge the timestamp and started times. + entity.setTimestamp(Math.max(entity.getTimestamp(), getTimestamp())); + entity.setStarted(Math.max(entity.getStarted(), getStarted())); + } return entity; } @Override public Expiration computeExpiration() { - long maxIdle; - long lifespan; - if (offline) { - maxIdle = SessionTimeouts.getOfflineClientSessionMaxIdleMs(userSession.getRealm(), client, getValue()); - lifespan = SessionTimeouts.getOfflineClientSessionLifespanMs(userSession.getRealm(), client, getValue()); - } else { - maxIdle = SessionTimeouts.getClientSessionMaxIdleMs(userSession.getRealm(), client, getValue()); - lifespan = SessionTimeouts.getClientSessionLifespanMs(userSession.getRealm(), client, getValue()); - } + long maxIdle = SessionTimeouts.getClientSessionMaxIdleMs(userSession.getRealm(), client, offline, isUserSessionRememberMe(), getTimestamp()); + long lifespan = SessionTimeouts.getClientSessionLifespanMs(userSession.getRealm(), client, offline, isUserSessionRememberMe(), getStarted(), getUserSessionStarted()); return new Expiration(maxIdle, lifespan); } @Override public String getId() { - return getValue().getId().toString(); + return getValue().createId(); + } + + @Override + public int getStarted() { + return getValue().getStarted(); + } + + @Override + public int getUserSessionStarted() { + checkInitialized(); + return userSession.getStarted(); + } + + @Override + public boolean isUserSessionRememberMe() { + checkInitialized(); + return userSession.isRememberMe(); } @Override @@ -112,7 +128,7 @@ public class AuthenticatedClientSessionUpdater extends BaseUpdater entity.setTimestamp(timestamp)); + addAndApplyChange(entity -> entity.setTimestamp(Math.max(timestamp, entity.getTimestamp()))); } @Override @@ -177,12 +193,17 @@ public class AuthenticatedClientSessionUpdater extends BaseUpdater entity.setAuthMethod(method)); + addAndApplyChange(entity -> entity.setProtocol(method)); + } + + @Override + public void restartClientSession() { + addAndApplyChange(RemoteAuthenticatedClientSessionEntity::restart); } @Override @@ -219,12 +240,18 @@ public class AuthenticatedClientSessionUpdater extends BaseUpdater change) { + private void addAndApplyChange(Consumer change) { changes.add(change); change.accept(getValue()); } - private static void initNotes(AuthenticatedClientSessionEntity entity) { + private void checkInitialized() { + if (!isInitialized()) { + throw new IllegalStateException(getClass().getSimpleName() + " not initialized yet!"); + } + } + + private static void initNotes(RemoteAuthenticatedClientSessionEntity entity) { var notes = entity.getNotes(); if (notes == null) { entity.setNotes(new HashMap<>()); @@ -232,21 +259,20 @@ public class AuthenticatedClientSessionUpdater extends BaseUpdater { + boolean offline) implements UpdaterFactory { @Override - public AuthenticatedClientSessionUpdater create(UUID key, AuthenticatedClientSessionEntity entity) { + public AuthenticatedClientSessionUpdater create(ClientSessionKey key, RemoteAuthenticatedClientSessionEntity entity) { return new AuthenticatedClientSessionUpdater(key, Objects.requireNonNull(entity), -1, offline, UpdaterState.CREATED); } @Override - public AuthenticatedClientSessionUpdater wrapFromCache(UUID key, MetadataValue entity) { - assert entity != null; - return new AuthenticatedClientSessionUpdater(key, Objects.requireNonNull(entity.getValue()), entity.getVersion(), offline, UpdaterState.READ); + public AuthenticatedClientSessionUpdater wrapFromCache(ClientSessionKey key, RemoteAuthenticatedClientSessionEntity value, long version) { + return new AuthenticatedClientSessionUpdater(key, Objects.requireNonNull(value), version, offline, UpdaterState.READ); } @Override - public AuthenticatedClientSessionUpdater deleted(UUID key) { + public AuthenticatedClientSessionUpdater deleted(ClientSessionKey key) { return new AuthenticatedClientSessionUpdater(key, null, -1, offline, UpdaterState.DELETED); } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/loginfailures/LoginFailuresUpdater.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/loginfailures/LoginFailuresUpdater.java index f0094dc55d2..ccbe4855bea 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/loginfailures/LoginFailuresUpdater.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/loginfailures/LoginFailuresUpdater.java @@ -21,7 +21,6 @@ import java.util.List; import java.util.Objects; import java.util.function.Consumer; -import org.infinispan.client.hotrod.MetadataValue; import org.keycloak.models.UserLoginFailureModel; import org.keycloak.models.sessions.infinispan.changes.remote.updater.BaseUpdater; import org.keycloak.models.sessions.infinispan.changes.remote.updater.Expiration; @@ -50,15 +49,15 @@ public class LoginFailuresUpdater extends BaseUpdater entity) { - return new LoginFailuresUpdater(Objects.requireNonNull(key), Objects.requireNonNull(entity.getValue()), entity.getVersion(), UpdaterState.READ); + public static LoginFailuresUpdater wrap(LoginFailureKey key, LoginFailureEntity value, long version) { + return new LoginFailuresUpdater(key, Objects.requireNonNull(value), version, UpdaterState.READ); } public static LoginFailuresUpdater delete(LoginFailureKey key) { - return new LoginFailuresUpdater(Objects.requireNonNull(key), null, -1, UpdaterState.DELETED); + return new LoginFailuresUpdater(key, null, -1, UpdaterState.DELETED); } @Override diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionMappingAdapter.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionMappingAdapter.java deleted file mode 100644 index ef33755be64..00000000000 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionMappingAdapter.java +++ /dev/null @@ -1,144 +0,0 @@ -/* - * Copyright 2024 Red Hat, Inc. and/or its affiliates - * and other contributors as indicated by the @author tags. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.keycloak.models.sessions.infinispan.changes.remote.updater.user; - -import java.util.AbstractMap; -import java.util.Collection; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Set; -import java.util.UUID; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.function.Consumer; - -import org.infinispan.client.hotrod.RemoteCache; -import org.infinispan.commons.util.concurrent.CompletionStages; -import org.keycloak.models.AuthenticatedClientSessionModel; -import org.keycloak.models.UserSessionModel; -import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionStore; -import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity; - -/** - * This class adapts and converts the {@link UserSessionEntity#getAuthenticatedClientSessions()} into - * {@link UserSessionModel#getAuthenticatedClientSessions()}. - *

- * Its implementation optimizes methods {@link #clear()}, {@link #put(String, AuthenticatedClientSessionModel)}, - * {@link #get(Object)} and {@link #remove(Object)} by avoiding download all client sessions from the - * {@link RemoteCache}. - *

- * The remaining methods are more expensive and require downloading all client sessions. The requests are done in - * concurrently to reduce the overall response time. - *

- * This class keeps track of any modification required in {@link UserSessionEntity#getAuthenticatedClientSessions()} and - * those modification can be replayed. - */ -public class ClientSessionMappingAdapter extends AbstractMap { - - private static final Consumer CLEAR = AuthenticatedClientSessionStore::clear; - - private final AuthenticatedClientSessionStore mappings; - private final ClientSessionProvider clientSessionProvider; - private final List> changes; - - public ClientSessionMappingAdapter(AuthenticatedClientSessionStore mappings, ClientSessionProvider clientSessionProvider) { - this.mappings = Objects.requireNonNull(mappings); - this.clientSessionProvider = Objects.requireNonNull(clientSessionProvider); - changes = new CopyOnWriteArrayList<>(); - } - - @Override - public void clear() { - mappings.forEach((id, uuid) -> clientSessionProvider.removeClientSession(uuid)); - changes.clear(); - addChangeAndApply(CLEAR); - } - - @Override - public AuthenticatedClientSessionModel put(String key, AuthenticatedClientSessionModel value) { - addChangeAndApply(store -> store.put(key, UUID.fromString(value.getId()))); - return clientSessionProvider.getClientSession(key, mappings.get(key)); - } - - @Override - public AuthenticatedClientSessionModel remove(Object key) { - var clientId = String.valueOf(key); - var uuid = mappings.get(clientId); - var existing = clientSessionProvider.getClientSession(clientId, uuid); - onClientRemoved(clientId, uuid); - return existing; - } - - @Override - public AuthenticatedClientSessionModel get(Object key) { - var clientId = String.valueOf(key); - return clientSessionProvider.getClientSession(clientId, mappings.get(clientId)); - } - - @SuppressWarnings("NullableProblems") - @Override - public Set> entrySet() { - Map results = new ConcurrentHashMap<>(mappings.size()); - var stage = CompletionStages.aggregateCompletionStage(); - mappings.forEach((clientId, uuid) -> stage.dependsOn(clientSessionProvider.getClientSessionAsync(clientId, uuid) - .thenAccept(updater -> { - if (updater == null) { - onClientRemoved(clientId, uuid); - return; - } - results.put(clientId, updater); - }))); - CompletionStages.join(stage.freeze()); - return results.entrySet(); - } - - boolean isUnchanged() { - return changes.isEmpty(); - } - - void removeAll(Collection removedClientUUIDS) { - if (removedClientUUIDS == null || removedClientUUIDS.isEmpty()) { - return; - } - removedClientUUIDS.forEach(this::onClientRemoved); - } - - /** - * Applies the modifications recorded by this class into a different {@link AuthenticatedClientSessionStore}. - * - * @param store The {@link AuthenticatedClientSessionStore} to update. - */ - void applyChanges(AuthenticatedClientSessionStore store) { - changes.forEach(change -> change.accept(store)); - } - - private void addChangeAndApply(Consumer change) { - change.accept(mappings); - changes.add(change); - } - - private void onClientRemoved(String clientId) { - onClientRemoved(clientId, mappings.get(clientId)); - } - - private void onClientRemoved(String clientId, UUID key) { - addChangeAndApply(store -> store.remove(clientId)); - clientSessionProvider.removeClientSession(key); - } -} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionProvider.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionProvider.java deleted file mode 100644 index e217b339c0e..00000000000 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/ClientSessionProvider.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright 2024 Red Hat, Inc. and/or its affiliates - * and other contributors as indicated by the @author tags. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.keycloak.models.sessions.infinispan.changes.remote.updater.user; - -import java.util.UUID; -import java.util.concurrent.CompletionStage; - -import org.infinispan.client.hotrod.RemoteCache; -import org.keycloak.models.AuthenticatedClientSessionModel; - -/** - * An SPI for {@link ClientSessionMappingAdapter} to interact with the {@link RemoteCache}. - */ -public interface ClientSessionProvider { - - /** - * Synchronously fetch an {@link AuthenticatedClientSessionModel} from the {@link RemoteCache}. - * - * @param clientId The client's ID. - * @param clientSessionId The {@link RemoteCache} key. - * @return The {@link AuthenticatedClientSessionModel} instance or {@code null} if the client session does not exist - * or was removed. - */ - AuthenticatedClientSessionModel getClientSession(String clientId, UUID clientSessionId); - - /** - * A non-blocking alternative to {@link #getClientSession(String, UUID)}. - * - * @see #getClientSession(String, UUID) - */ - CompletionStage getClientSessionAsync(String clientId, UUID clientSessionId); - - /** - * Removes the client session associated with the {@link RemoteCache} key. - *

- * If {@code clientSessionId} is {@code null}, nothing is removed. The methods - * {@link #getClientSession(String, UUID)} and {@link #getClientSessionAsync(String, UUID)} will return {@code null} - * for the session after this method is completed. - * - * @param clientSessionId The {@link RemoteCache} key to remove. - */ - void removeClientSession(UUID clientSessionId); - -} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/UserSessionUpdater.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/UserSessionUpdater.java index ab7611905f8..db3bccc5a20 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/UserSessionUpdater.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/remote/updater/user/UserSessionUpdater.java @@ -8,7 +8,6 @@ import java.util.Map; import java.util.Objects; import java.util.function.Consumer; -import org.infinispan.client.hotrod.MetadataValue; import org.keycloak.models.AuthenticatedClientSessionModel; import org.keycloak.models.RealmModel; import org.keycloak.models.UserModel; @@ -18,27 +17,26 @@ import org.keycloak.models.sessions.infinispan.changes.remote.updater.Expiration import org.keycloak.models.sessions.infinispan.changes.remote.updater.Updater; import org.keycloak.models.sessions.infinispan.changes.remote.updater.UpdaterFactory; import org.keycloak.models.sessions.infinispan.changes.remote.updater.helper.MapUpdater; -import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionStore; -import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; import org.keycloak.models.sessions.infinispan.util.SessionTimeouts; /** * The {@link Updater} implementation to keep track of modifications for {@link UserSessionModel}. */ -public class UserSessionUpdater extends BaseUpdater implements UserSessionModel { +public class UserSessionUpdater extends BaseUpdater implements UserSessionModel { private static final Factory ONLINE = new Factory(false); private static final Factory OFFLINE = new Factory(true); private final MapUpdater notesUpdater; - private final List> changes; + private final List> changes; private final boolean offline; private RealmModel realm; private UserModel user; - private ClientSessionMappingAdapter clientSessionMappingAdapter; + private Map clientSessions; private SessionPersistenceState persistenceState = SessionPersistenceState.PERSISTENT; - private UserSessionUpdater(String cacheKey, UserSessionEntity cacheValue, long version, boolean offline, UpdaterState initialState) { + private UserSessionUpdater(String cacheKey, RemoteUserSessionEntity cacheValue, long version, boolean offline, UpdaterState initialState) { super(cacheKey, cacheValue, version, initialState); this.offline = offline; if (cacheValue == null) { @@ -57,37 +55,28 @@ public class UserSessionUpdater extends BaseUpdater i * @param offline If {@code true}, it creates offline {@link UserSessionModel}. * @return The {@link UpdaterFactory} implementation to create instances of {@link UserSessionModel}. */ - public static UpdaterFactory factory(boolean offline) { + public static UpdaterFactory factory(boolean offline) { return offline ? OFFLINE : ONLINE; } @Override - public UserSessionEntity apply(String ignored, UserSessionEntity userSessionEntity) { + public RemoteUserSessionEntity apply(String ignored, RemoteUserSessionEntity userSessionEntity) { initNotes(userSessionEntity); - initStore(userSessionEntity); changes.forEach(change -> change.accept(userSessionEntity)); notesUpdater.applyChanges(userSessionEntity.getNotes()); - clientSessionMappingAdapter.applyChanges(userSessionEntity.getAuthenticatedClientSessions()); return userSessionEntity; } @Override public Expiration computeExpiration() { - long maxIdle; - long lifespan; - if (offline) { - maxIdle = SessionTimeouts.getOfflineSessionMaxIdleMs(realm, null, getValue()); - lifespan = SessionTimeouts.getOfflineSessionLifespanMs(realm, null, getValue()); - } else { - maxIdle = SessionTimeouts.getUserSessionMaxIdleMs(realm, null, getValue()); - lifespan = SessionTimeouts.getUserSessionLifespanMs(realm, null, getValue()); - } + long maxIdle = SessionTimeouts.getUserSessionMaxIdleMs(realm, isOffline(), getValue().isRememberMe(), getValue().getLastSessionRefresh()); + long lifespan = SessionTimeouts.getUserSessionLifespanMs(realm, isOffline(), getValue().isRememberMe(), getValue().getStarted()); return new Expiration(maxIdle, lifespan); } @Override public String getId() { - return getValue().getId(); + return getKey(); } @Override @@ -152,17 +141,20 @@ public class UserSessionUpdater extends BaseUpdater i @Override public Map getAuthenticatedClientSessions() { - return clientSessionMappingAdapter; + return clientSessions; } @Override public void removeAuthenticatedClientSessions(Collection removedClientUUIDS) { - clientSessionMappingAdapter.removeAll(removedClientUUIDS); + if (removedClientUUIDS == null || removedClientUUIDS.isEmpty()) { + return; + } + removedClientUUIDS.forEach(clientSessions::remove); } @Override public AuthenticatedClientSessionModel getAuthenticatedClientSessionByClient(String clientUUID) { - return clientSessionMappingAdapter.get(clientUUID); + return clientSessions.get(clientUUID); } @Override @@ -205,8 +197,8 @@ public class UserSessionUpdater extends BaseUpdater i this.user = user; changes.clear(); notesUpdater.clear(); - clientSessionMappingAdapter.clear(); - addAndApplyChange(userSessionEntity -> UserSessionEntity.updateSessionEntity(userSessionEntity, realm, user, loginUsername, ipAddress, authMethod, rememberMe, brokerSessionId, brokerUserId)); + clientSessions.clear(); + addAndApplyChange(userSessionEntity -> userSessionEntity.restart(realm.getId(), user.getId(), loginUsername, ipAddress, authMethod, rememberMe, brokerSessionId, brokerUserId)); } @Override @@ -221,7 +213,7 @@ public class UserSessionUpdater extends BaseUpdater i @Override protected boolean isUnchanged() { - return changes.isEmpty() && notesUpdater.isUnchanged() && clientSessionMappingAdapter.isUnchanged(); + return changes.isEmpty() && notesUpdater.isUnchanged(); } /** @@ -230,15 +222,13 @@ public class UserSessionUpdater extends BaseUpdater i * @param persistenceState The {@link SessionPersistenceState}. * @param realm The {@link RealmModel} to where this user session belongs. * @param user The {@link UserModel} associated to this user session. - * @param factory The {@link ClientSessionAdapterFactory} to create the {@link ClientSessionMappingAdapter} - * to track modifications into the client sessions. + * @param clientSessions The {@link Map} associated to this use session. */ - public synchronized void initialize(SessionPersistenceState persistenceState, RealmModel realm, UserModel user, ClientSessionAdapterFactory factory) { - initStore(getValue()); + public synchronized void initialize(SessionPersistenceState persistenceState, RealmModel realm, UserModel user, Map clientSessions) { this.realm = Objects.requireNonNull(realm); this.user = Objects.requireNonNull(user); this.persistenceState = Objects.requireNonNull(persistenceState); - clientSessionMappingAdapter = factory.create(getValue().getAuthenticatedClientSessions()); + this.clientSessions = Objects.requireNonNull(clientSessions); } /** @@ -248,44 +238,29 @@ public class UserSessionUpdater extends BaseUpdater i return realm != null; } - private void addAndApplyChange(Consumer change) { + private void addAndApplyChange(Consumer change) { change.accept(getValue()); changes.add(change); } - private static void initNotes(UserSessionEntity entity) { + private static void initNotes(RemoteUserSessionEntity entity) { var notes = entity.getNotes(); if (notes == null) { entity.setNotes(new HashMap<>()); } } - private static void initStore(UserSessionEntity entity) { - var store = entity.getAuthenticatedClientSessions(); - if (store == null) { - entity.setAuthenticatedClientSessions(new AuthenticatedClientSessionStore()); - } - } - - /** - * Factory SPI to create {@link ClientSessionMappingAdapter} for the {@link AuthenticatedClientSessionStore} used by - * this instance. - */ - public interface ClientSessionAdapterFactory { - ClientSessionMappingAdapter create(AuthenticatedClientSessionStore clientSessionStore); - } - - private record Factory(boolean offline) implements UpdaterFactory { + private record Factory( + boolean offline) implements UpdaterFactory { @Override - public UserSessionUpdater create(String key, UserSessionEntity entity) { + public UserSessionUpdater create(String key, RemoteUserSessionEntity entity) { return new UserSessionUpdater(key, Objects.requireNonNull(entity), -1, offline, UpdaterState.CREATED); } @Override - public UserSessionUpdater wrapFromCache(String key, MetadataValue entity) { - assert entity != null; - return new UserSessionUpdater(key, Objects.requireNonNull(entity.getValue()), entity.getVersion(), offline, UpdaterState.READ); + public UserSessionUpdater wrapFromCache(String key, RemoteUserSessionEntity value, long version) { + return new UserSessionUpdater(key, value, version, offline, UpdaterState.READ); } @Override diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionKey.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionKey.java new file mode 100644 index 00000000000..3203d7b40d0 --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionKey.java @@ -0,0 +1,31 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.entities; + +import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.protostream.annotations.Proto; +import org.infinispan.protostream.annotations.ProtoTypeId; +import org.keycloak.marshalling.Marshalling; + +/** + * The key stored in the {@link RemoteCache} for {@link RemoteAuthenticatedClientSessionEntity}. + */ +@ProtoTypeId(Marshalling.CLIENT_SESSION_KEY) +@Proto +public record ClientSessionKey(String userSessionId, String clientId) { +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteAuthenticatedClientSessionEntity.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteAuthenticatedClientSessionEntity.java new file mode 100644 index 00000000000..a801eb5c06f --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteAuthenticatedClientSessionEntity.java @@ -0,0 +1,194 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.entities; + +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; + +import org.infinispan.api.annotations.indexing.Basic; +import org.infinispan.api.annotations.indexing.Indexed; +import org.infinispan.protostream.annotations.ProtoFactory; +import org.infinispan.protostream.annotations.ProtoField; +import org.infinispan.protostream.annotations.ProtoTypeId; +import org.keycloak.common.util.Time; +import org.keycloak.marshalling.Marshalling; +import org.keycloak.models.AuthenticatedClientSessionModel; +import org.keycloak.models.UserSessionModel; + +@ProtoTypeId(Marshalling.REMOTE_CLIENT_SESSION_ENTITY) +@Indexed +public class RemoteAuthenticatedClientSessionEntity { + + // immutable state + private final String userSessionId; + private final String clientId; + private final String userId; + private final String realmId; + + // mutable state + private int started; + private String protocol; + private String redirectUri; + private String action; + private Map notes; + private int timestamp; + + private RemoteAuthenticatedClientSessionEntity(String userSessionId, String clientId, String userId, String realmId) { + this.userSessionId = Objects.requireNonNull(userSessionId); + this.clientId = Objects.requireNonNull(clientId); + this.userId = Objects.requireNonNull(userId); + this.realmId = Objects.requireNonNull(realmId); + } + + @ProtoFactory + RemoteAuthenticatedClientSessionEntity(String clientId, String userId, String userSessionId, String realmId, Map notes, String action, String protocol, String redirectUri, int timestamp, int started) { + this.userSessionId = userSessionId; + this.clientId = clientId; + this.userId = userId; + this.realmId = realmId; + this.action = action; + this.protocol = protocol; + this.redirectUri = redirectUri; + this.notes = notes; + this.timestamp = timestamp; + this.started = started; + } + + public static RemoteAuthenticatedClientSessionEntity create(ClientSessionKey id, String realmId, UserSessionModel userSession) { + var e = new RemoteAuthenticatedClientSessionEntity(id.userSessionId(), id.clientId(), userSession.getUser().getId(), realmId); + e.timestamp = e.started = Time.currentTime(); + e.notes = new HashMap<>(); + return e; + } + + public static RemoteAuthenticatedClientSessionEntity createFromModel(ClientSessionKey id, AuthenticatedClientSessionModel model) { + var e = new RemoteAuthenticatedClientSessionEntity(id.userSessionId(), id.clientId(), model.getUserSession().getUser().getId(), model.getRealm().getId()); + e.timestamp = e.started = Time.currentTime(); + e.notes = model.getNotes() == null || model.getNotes().isEmpty() ? + new HashMap<>() : + new HashMap<>(model.getNotes()); + return e; + } + + // for testing purposes only! + public static RemoteAuthenticatedClientSessionEntity mockEntity(String userSessionId, String userId, String realmId) { + return mockEntity(userSessionId, "client", userId, realmId); + } + + // for testing purposes only! + public static RemoteAuthenticatedClientSessionEntity mockEntity(String userSessionId, String clientId, String userId, String realmId) { + return new RemoteAuthenticatedClientSessionEntity(userSessionId, clientId, userId, realmId); + } + + @ProtoField(1) + @Basic(projectable = true, sortable = true) + public String getClientId() { + return clientId; + } + + @ProtoField(2) + @Basic + public String getUserId() { + return userId; + } + + @ProtoField(3) + @Basic(projectable = true, sortable = true) + public String getUserSessionId() { + return userSessionId; + } + + @ProtoField(4) + @Basic + public String getRealmId() { + return realmId; + } + + @ProtoField(value = 5, mapImplementation = HashMap.class) + public Map getNotes() { + return notes; + } + + public void setNotes(Map notes) { + this.notes = notes; + } + + @ProtoField(6) + public String getAction() { + return action; + } + + public void setAction(String action) { + this.action = action; + } + + @ProtoField(7) + public String getProtocol() { + return protocol; + } + + public void setProtocol(String protocol) { + this.protocol = protocol; + } + + @ProtoField(8) + public String getRedirectUri() { + return redirectUri; + } + + public void setRedirectUri(String redirectUri) { + this.redirectUri = redirectUri; + } + + @ProtoField(9) + public int getTimestamp() { + return timestamp; + } + + public void setTimestamp(int timestamp) { + this.timestamp = timestamp; + } + + @ProtoField(10) + public int getStarted() { + return started; + } + + public void setStarted(int started) { + this.started = started; + } + + public void restart() { + action = null; + redirectUri = null; + timestamp = started = Time.currentTime(); + notes.clear(); + } + + public ClientSessionKey createCacheKey() { + return new ClientSessionKey(userSessionId, clientId); + } + + public String createId() { + return UUID.nameUUIDFromBytes((userSessionId + clientId).getBytes(StandardCharsets.UTF_8)).toString(); + } + +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteUserSessionEntity.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteUserSessionEntity.java new file mode 100644 index 00000000000..93317e86a9a --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/RemoteUserSessionEntity.java @@ -0,0 +1,213 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.entities; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import org.infinispan.api.annotations.indexing.Basic; +import org.infinispan.api.annotations.indexing.Indexed; +import org.infinispan.protostream.annotations.ProtoFactory; +import org.infinispan.protostream.annotations.ProtoField; +import org.infinispan.protostream.annotations.ProtoTypeId; +import org.keycloak.common.util.Time; +import org.keycloak.marshalling.Marshalling; +import org.keycloak.models.OfflineUserSessionModel; +import org.keycloak.models.RealmModel; +import org.keycloak.models.UserModel; +import org.keycloak.models.UserSessionModel; + +@ProtoTypeId(Marshalling.REMOTE_USER_SESSION_ENTITY) +@Indexed +public class RemoteUserSessionEntity { + + // immutable state + private final String userSessionId; + + // mutable state + private String realmId; + private String userId; + private String brokerSessionId; + private String brokerUserId; + private String loginUsername; + private String ipAddress; + private String authMethod; + private boolean rememberMe; + private int started; + private int lastSessionRefresh; + private UserSessionModel.State state; + private Map notes; + + private RemoteUserSessionEntity(String userSessionId) { + this.userSessionId = Objects.requireNonNull(userSessionId); + } + + public static RemoteUserSessionEntity create(String id, RealmModel realm, UserModel user, String loginUsername, String ipAddress, String authMethod, boolean rememberMe, String brokerSessionId, String brokerUserId) { + var e = new RemoteUserSessionEntity(id); + e.restart(realm.getId(), user.getId(), loginUsername, ipAddress, authMethod, rememberMe, brokerSessionId, brokerUserId); + return e; + } + + public static RemoteUserSessionEntity createFromModel(UserSessionModel model) { + String userId; + String loginUsername = null; + if (model instanceof OfflineUserSessionModel offline) { + // this is a hack so that UserModel doesn't have to be available when offline token is imported. + // see related JIRA - KEYCLOAK-5350 and corresponding test + userId = offline.getUserId(); + // NOTE: Hack + // We skip calling entity.setLoginUsername(userSession.getLoginUsername()) + } else { + userId = model.getUser().getId(); + loginUsername = model.getLoginUsername(); + } + var e = new RemoteUserSessionEntity(model.getId()); + e.restart(model.getRealm().getId(), userId, loginUsername, model.getIpAddress(), model.getAuthMethod(), model.isRememberMe(), model.getBrokerSessionId(), model.getBrokerUserId()); + var notes = model.getNotes(); + if (notes != null && !notes.isEmpty()) { + e.notes = new HashMap<>(notes); + } + e.state = model.getState(); + return e; + } + + // for testing purposes only! + public static RemoteUserSessionEntity mockEntity(String id, String realmId, String userId) { + return mockEntity(id, realmId, userId, null, null); + } + + // for testing purposes only! + public static RemoteUserSessionEntity mockEntity(String id, String realmId, String userId, String brokerSessionId, String brokerUserId) { + var e = new RemoteUserSessionEntity(id); + e.realmId = realmId; + e.userId = userId; + e.brokerSessionId = brokerSessionId; + e.brokerUserId = brokerUserId; + return e; + } + + @ProtoFactory + static RemoteUserSessionEntity protoFactory(String userSessionId, String authMethod, String brokerSessionId, String brokerUserId, String ipAddress, int lastSessionRefresh, String loginUsername, Map notes, String realmId, boolean rememberMe, int started, UserSessionModel.State state, String userId) { + var e = new RemoteUserSessionEntity(userSessionId); + e.applyState(authMethod, brokerSessionId, brokerUserId, ipAddress, lastSessionRefresh, loginUsername, notes, realmId, rememberMe, started, state, userId); + return e; + } + + @ProtoField(1) + @Basic(sortable = true) + public String getUserSessionId() { + return userSessionId; + } + + @ProtoField(2) + public String getAuthMethod() { + return authMethod; + } + + @ProtoField(3) + @Basic + public String getBrokerSessionId() { + return brokerSessionId; + } + + @ProtoField(4) + @Basic + public String getBrokerUserId() { + return brokerUserId; + } + + @ProtoField(5) + public String getIpAddress() { + return ipAddress; + } + + @ProtoField(6) + public int getLastSessionRefresh() { + return lastSessionRefresh; + } + + public void setLastSessionRefresh(int lastSessionRefresh) { + this.lastSessionRefresh = Math.max(this.lastSessionRefresh, lastSessionRefresh); + } + + @ProtoField(7) + public String getLoginUsername() { + return loginUsername; + } + + @ProtoField(value = 8, mapImplementation = HashMap.class) + public Map getNotes() { + return notes; + } + + public void setNotes(Map notes) { + this.notes = notes; + } + + @ProtoField(9) + @Basic + public String getRealmId() { + return realmId; + } + + @ProtoField(10) + public boolean isRememberMe() { + return rememberMe; + } + + @ProtoField(11) + public int getStarted() { + return started; + } + + @ProtoField(12) + public UserSessionModel.State getState() { + return state; + } + + public void setState(UserSessionModel.State state) { + this.state = state; + } + + @ProtoField(13) + @Basic + public String getUserId() { + return userId; + } + + public void restart(String realmId, String userId, String loginUsername, String ipAddress, String authMethod, boolean rememberMe, String brokerSessionId, String brokerUserId) { + var currentTime = Time.currentTime(); + applyState(authMethod, brokerSessionId, brokerUserId, ipAddress, currentTime, loginUsername, null, realmId, rememberMe, currentTime, null, userId); + } + + private void applyState(String authMethod, String brokerSessionId, String brokerUserId, String ipAddress, int lastSessionRefresh, String loginUsername, Map notes, String realmId, boolean rememberMe, int started, UserSessionModel.State state, String userId) { + this.realmId = realmId; + this.userId = userId; + this.loginUsername = loginUsername; + this.ipAddress = ipAddress; + this.authMethod = authMethod; + this.rememberMe = rememberMe; + this.brokerSessionId = brokerSessionId; + this.brokerUserId = brokerUserId; + this.started = started; + this.lastSessionRefresh = lastSessionRefresh; + this.notes = notes; + this.state = state; + } +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/ClientSessionQueries.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/ClientSessionQueries.java new file mode 100644 index 00000000000..ceff06784d1 --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/ClientSessionQueries.java @@ -0,0 +1,76 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.query; + +import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.commons.api.query.Query; +import org.keycloak.marshalling.Marshalling; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; + +/** + * Util class with Infinispan Ickle Queries for {@link RemoteAuthenticatedClientSessionEntity}. + */ +public final class ClientSessionQueries { + + private ClientSessionQueries() { + } + + public static final String CLIENT_SESSION = Marshalling.protoEntity(RemoteAuthenticatedClientSessionEntity.class); + + private static final String FETCH_USER_SESSION_ID = "SELECT e.userSessionId FROM %s as e WHERE e.realmId = :realmId && e.clientId = :clientId ORDER BY e.userSessionId".formatted(CLIENT_SESSION); + private static final String PER_CLIENT_COUNT = "SELECT e.clientId, count(e.clientId) FROM %s as e GROUP BY e.clientId ORDER BY e.clientId".formatted(CLIENT_SESSION); + private static final String CLIENT_SESSION_COUNT = "SELECT count(e) FROM %s as e WHERE e.realmId = :realmId && e.clientId = :clientId".formatted(CLIENT_SESSION); + private static final String FROM_USER_SESSION = "SELECT e, version(e) FROM %s as e WHERE e.userSessionId = :userSessionId ORDER BY e.clientId".formatted(CLIENT_SESSION); + + /** + * Returns a projection with the user session ID for client sessions from the client {@code clientId}. + */ + public static Query fetchUserSessionIdForClientId(RemoteCache cache, String realmId, String clientId) { + return cache.query(FETCH_USER_SESSION_ID) + .setParameter("realmId", realmId) + .setParameter("clientId", clientId); + } + + /** + * Returns a projection with the client ID and its number of active client sessions. + */ + public static Query activeClientCount(RemoteCache cache) { + return cache.query(PER_CLIENT_COUNT); + } + + /** + * Returns a projection with the sum of all client session belonging to the client ID. + */ + public static Query countClientSessions(RemoteCache cache, String realmId, String clientId) { + return cache.query(CLIENT_SESSION_COUNT) + .setParameter("realmId", realmId) + .setParameter("clientId", clientId); + } + + /** + * Returns a projection with the client session, and the version of all client sessions belonging to the user + * session ID. + */ + public static Query fetchClientSessions(RemoteCache cache, String userSessionId) { + return cache.query(FROM_USER_SESSION) + .setParameter("userSessionId", userSessionId); + } + + +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/QueryHelper.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/QueryHelper.java new file mode 100644 index 00000000000..6289369fa10 --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/QueryHelper.java @@ -0,0 +1,207 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.query; + +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Optional; +import java.util.Spliterators; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import java.util.stream.StreamSupport; + +import org.infinispan.client.hotrod.impl.query.RemoteQuery; +import org.infinispan.commons.api.query.Query; +import org.infinispan.query.dsl.QueryResult; + +public final class QueryHelper { + + /** + * Converts a single projection results into a long value. + */ + public static final Function SINGLE_PROJECTION_TO_LONG = projection -> { + assert projection.length == 1; + return (long) projection[0]; + }; + + /** + * Converts a single projection value into a {@link String}. + */ + public static final Function SINGLE_PROJECTION_TO_STRING = projection -> { + assert projection.length == 1; + return String.valueOf(projection[0]); + }; + + /** + * Converts a projection with two values into a {@link Map.Entry} of {@link String} and {@link Long}, where the key + * is the first projection, and the second is the second project. + */ + public static final Function> PROJECTION_TO_STRING_LONG_ENTRY = projection -> { + assert projection.length == 2; + return Map.entry((String) projection[0], (long) projection[1]); + }; + + private QueryHelper() { + } + + /** + * Fetches a single value from the query. + *

+ * This method changes the {@link Query} state to return just a single value. + * + * @param query The {@link Query} instance. + * @param mapping The {@link Function} that maps the query results (projection) into the result. + * @param The {@link Query} response type. + * @param The {@link Optional} type. + * @return An {@link Optional} with the {@link Query} results mapped. + */ + public static Optional fetchSingle(Query query, Function mapping) { + query.hitCountAccuracy(1).maxResults(1); + try (var iterator = query.iterator()) { + return iterator.hasNext() ? Optional.ofNullable(mapping.apply(iterator.next())) : Optional.empty(); + } + } + + /** + * Streams using batching over all results from the {@link Query}. + *

+ * If a large result set is expected, this method is recommended to avoid loading downloading a lot of data in a + * single request. + *

+ * The results are fetched on demand. + *

+ * Warning: This method changes ignores the start offset and the max results. It will return everything. + * + * @param query The {@link Query} instance. + * @param batchSize The number of results to fetch for each remote request. + * @param mapping The {@link Function} that maps the query results (projection) into the result. + * @param The {@link Query} response type. + * @param The {@link Stream} type. + * @return A {@link Stream} with the results. + */ + public static Stream streamAll(Query query, int batchSize, Function mapping) { + return StreamSupport.stream(Spliterators.spliteratorUnknownSize(new BatchingIterator<>(query, batchSize, mapping), 0), false); + } + + /** + * Performs the {@link Query} and returns the results. + *

+ * This method is preferred to {@link Query#list()} since it does not have to compute an accurate hit count (affects + * Indexed query performance). + *

+ * If a large dataset is expected, use {@link #streamAll(Query, int, Function)}. + * + * @param query The {@link Query} instance. + * @param mapping The {@link Function} that maps the query results (projection) into the result. + * @param The {@link Query} response type. + * @param The {@link Collection} type. + * @return A {@link Collection} with the results. + */ + public static Collection toCollection(Query query, Function mapping) { + try (var iterator = query.iterator()) { + return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator, 0), false) + .map(mapping) + .collect(Collectors.toList()); + } + } + + // TODO to be removed. A publisher was added to the Infinispan API since version 15.1. + private static class BatchingIterator implements Iterator { + + private final RemoteQuery query; + private final int batchSize; + private final Function mapping; + private int currentOffset; + private Iterator currentResults; + private CompletableFuture> nextResults; + private R next; + private boolean completed; + + private BatchingIterator(Query query, int batchSize, Function mapping) { + assert query instanceof RemoteQuery; + this.query = (RemoteQuery) query.startOffset(0).hitCountAccuracy(batchSize).maxResults(batchSize); + this.batchSize = batchSize; + this.mapping = mapping; + currentResults = Collections.emptyIterator(); + executeQueryAsync(); + fetchNext(); + } + + @Override + public boolean hasNext() { + return next != null; + } + + @Override + public R next() { + if (next == null) { + throw new NoSuchElementException(); + } + var result = next; + fetchNext(); + return result; + } + + private void executeQueryAsync() { + nextResults = query.executeAsync().toCompletableFuture(); + } + + private void fetchNext() { + while (true) { + while (currentResults.hasNext()) { + next = mapping.apply(currentResults.next()); + if (next != null) { + return; + } + } + if (completed) { + next = null; + return; + } + useNextResultsAndRequestMore(); + } + } + + private void useNextResultsAndRequestMore() { + var rsp = nextResults.join(); + var resultList = rsp.list(); + if (resultList.isEmpty()) { + completed = true; + return; + } + currentResults = resultList.iterator(); + if (resultList.size() < batchSize) { + completed = true; + return; + } + currentOffset += resultList.size(); + if (rsp.count().isExact() && currentOffset >= rsp.count().value()) { + completed = true; + return; + } + query.startOffset(currentOffset); + executeQueryAsync(); + } + } + +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/UserSessionQueries.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/UserSessionQueries.java new file mode 100644 index 00000000000..2b64b5fcec2 --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/query/UserSessionQueries.java @@ -0,0 +1,68 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.models.sessions.infinispan.query; + +import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.commons.api.query.Query; +import org.keycloak.marshalling.Marshalling; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; + +/** + * Util class with Infinispan Ickle Queries for {@link RemoteUserSessionEntity}. + */ +public final class UserSessionQueries { + + private UserSessionQueries() { + } + + public static final String USER_SESSION = Marshalling.protoEntity(RemoteUserSessionEntity.class); + + private static final String BASE_QUERY = "SELECT e, version(e) FROM %s as e ".formatted(USER_SESSION); + private static final String BY_BROKER_SESSION_ID = BASE_QUERY + "WHERE e.realmId = :realmId && e.brokerSessionId = :brokerSessionId ORDER BY e.userSessionId"; + private static final String BY_USER_ID = BASE_QUERY + "WHERE e.realmId = :realmId && e.userId = :userId ORDER BY e.userSessionId"; + private static final String BY_BROKER_USER_ID = BASE_QUERY + "WHERE e.realmId = :realmId && e.brokerUserId = :brokerUserId ORDER BY e.userSessionId"; + + /** + * Returns a projection with the user session, and the version of all user sessions belonging to the broker session + * ID. + */ + public static Query searchByBrokerSessionId(RemoteCache cache, String realmId, String brokerSessionId) { + return cache.query(BY_BROKER_SESSION_ID) + .setParameter("realmId", realmId) + .setParameter("brokerSessionId", brokerSessionId); + } + + /** + * Returns a projection with the user session, and the version of all user sessions belonging to the user ID. + */ + public static Query searchByUserId(RemoteCache cache, String realmId, String userId) { + return cache.query(BY_USER_ID) + .setParameter("realmId", realmId) + .setParameter("userId", userId); + } + + /** + * Returns a projection with the user session, and the version of all user sessions belonging to the broker user + * ID. + */ + public static Query searchByBrokerUserId(RemoteCache cache, String realmId, String brokerUserId) { + return cache.query(BY_BROKER_USER_ID) + .setParameter("realmId", realmId) + .setParameter("brokerUserId", brokerUserId); + } +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanAuthenticationSessionProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanAuthenticationSessionProviderFactory.java index 6f4546c7c45..6ecb496514c 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanAuthenticationSessionProviderFactory.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanAuthenticationSessionProviderFactory.java @@ -20,6 +20,7 @@ package org.keycloak.models.sessions.infinispan.remote; import java.lang.invoke.MethodHandles; import java.util.List; +import org.infinispan.client.hotrod.RemoteCache; import org.jboss.logging.Logger; import org.keycloak.Config; import org.keycloak.infinispan.util.InfinispanUtils; @@ -30,13 +31,13 @@ import org.keycloak.models.sessions.infinispan.InfinispanAuthenticationSessionPr import org.keycloak.models.sessions.infinispan.changes.remote.remover.query.ByRealmIdQueryConditionalRemover; import org.keycloak.models.sessions.infinispan.entities.RootAuthenticationSessionEntity; import org.keycloak.models.sessions.infinispan.remote.transaction.AuthenticationSessionTransaction; -import org.keycloak.models.sessions.infinispan.remote.transaction.RemoteCacheAndExecutor; import org.keycloak.provider.EnvironmentDependentProviderFactory; import org.keycloak.provider.ProviderConfigProperty; import org.keycloak.provider.ProviderConfigurationBuilder; import org.keycloak.sessions.AuthenticationSessionProviderFactory; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.AUTHENTICATION_SESSIONS_CACHE_NAME; +import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.getRemoteCache; import static org.keycloak.models.sessions.infinispan.InfinispanAuthenticationSessionProviderFactory.DEFAULT_AUTH_SESSIONS_LIMIT; public class RemoteInfinispanAuthenticationSessionProviderFactory implements AuthenticationSessionProviderFactory, EnvironmentDependentProviderFactory { @@ -45,7 +46,7 @@ public class RemoteInfinispanAuthenticationSessionProviderFactory implements Aut private static final String PROTO_ENTITY = Marshalling.protoEntity(RootAuthenticationSessionEntity.class); private int authSessionsLimit; - private volatile RemoteCacheAndExecutor cacheHolder; + private volatile RemoteCache cache; @Override public boolean isSupported(Config.Scope config) { @@ -64,13 +65,13 @@ public class RemoteInfinispanAuthenticationSessionProviderFactory implements Aut @Override public void postInit(KeycloakSessionFactory factory) { - cacheHolder = RemoteCacheAndExecutor.create(factory, AUTHENTICATION_SESSIONS_CACHE_NAME); + cache = getRemoteCache(factory, AUTHENTICATION_SESSIONS_CACHE_NAME); logger.debugf("Provided initialized. session limit=%s", authSessionsLimit); } @Override public void close() { - cacheHolder = null; + cache = null; } @Override @@ -96,7 +97,7 @@ public class RemoteInfinispanAuthenticationSessionProviderFactory implements Aut } private AuthenticationSessionTransaction createAndEnlistTransaction(KeycloakSession session) { - var tx = new AuthenticationSessionTransaction(cacheHolder.cache(), new ByRealmIdQueryConditionalRemover<>(PROTO_ENTITY, cacheHolder.executor())); + var tx = new AuthenticationSessionTransaction(cache, new ByRealmIdQueryConditionalRemover<>(PROTO_ENTITY)); session.getTransactionManager().enlistAfterCompletion(tx); return tx; } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserLoginFailureProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserLoginFailureProviderFactory.java index 00f1ead81b4..e20da5e8858 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserLoginFailureProviderFactory.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserLoginFailureProviderFactory.java @@ -18,7 +18,7 @@ package org.keycloak.models.sessions.infinispan.remote; import java.lang.invoke.MethodHandles; -import org.infinispan.client.hotrod.MetadataValue; +import org.infinispan.client.hotrod.RemoteCache; import org.jboss.logging.Logger; import org.keycloak.Config; import org.keycloak.infinispan.util.InfinispanUtils; @@ -34,17 +34,17 @@ import org.keycloak.models.sessions.infinispan.changes.remote.updater.loginfailu import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity; import org.keycloak.models.sessions.infinispan.entities.LoginFailureKey; import org.keycloak.models.sessions.infinispan.remote.transaction.LoginFailureChangeLogTransaction; -import org.keycloak.models.sessions.infinispan.remote.transaction.RemoteCacheAndExecutor; import org.keycloak.provider.EnvironmentDependentProviderFactory; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.LOGIN_FAILURE_CACHE_NAME; +import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.getRemoteCache; public class RemoteUserLoginFailureProviderFactory implements UserLoginFailureProviderFactory, UpdaterFactory, EnvironmentDependentProviderFactory { private static final Logger log = Logger.getLogger(MethodHandles.lookup().lookupClass()); - private static final String PROTO_ENTITY = Marshalling.protoEntity(LoginFailureEntity.class); + public static final String PROTO_ENTITY = Marshalling.protoEntity(LoginFailureEntity.class); - private volatile RemoteCacheAndExecutor cacheHolder; + private volatile RemoteCache cache; @Override public RemoteUserLoginFailureProvider create(KeycloakSession session) { @@ -57,19 +57,19 @@ public class RemoteUserLoginFailureProviderFactory implements UserLoginFailurePr @Override public void postInit(final KeycloakSessionFactory factory) { - cacheHolder = RemoteCacheAndExecutor.create(factory, LOGIN_FAILURE_CACHE_NAME); + cache = getRemoteCache(factory, LOGIN_FAILURE_CACHE_NAME); factory.register(event -> { if (event instanceof UserModel.UserRemovedEvent userRemovedEvent) { UserLoginFailureProvider provider = userRemovedEvent.getKeycloakSession().getProvider(UserLoginFailureProvider.class, getId()); provider.removeUserLoginFailure(userRemovedEvent.getRealm(), userRemovedEvent.getUser().getId()); } }); - log.debugf("Post Init. Cache=%s", cacheHolder.cache().getName()); + log.debugf("Post Init. Cache=%s", cache.getName()); } @Override public void close() { - cacheHolder = null; + cache = null; } @Override @@ -93,9 +93,8 @@ public class RemoteUserLoginFailureProviderFactory implements UserLoginFailurePr } @Override - public LoginFailuresUpdater wrapFromCache(LoginFailureKey key, MetadataValue entity) { - assert entity != null; - return LoginFailuresUpdater.wrap(key, entity); + public LoginFailuresUpdater wrapFromCache(LoginFailureKey key, LoginFailureEntity value, long version) { + return LoginFailuresUpdater.wrap(key, value, version); } @Override @@ -104,7 +103,7 @@ public class RemoteUserLoginFailureProviderFactory implements UserLoginFailurePr } private LoginFailureChangeLogTransaction createAndEnlistTransaction(KeycloakSession session) { - var tx = new LoginFailureChangeLogTransaction(this, cacheHolder.cache(), new ByRealmIdQueryConditionalRemover<>(PROTO_ENTITY, cacheHolder.executor())); + var tx = new LoginFailureChangeLogTransaction(this, cache, new ByRealmIdQueryConditionalRemover<>(PROTO_ENTITY)); session.getTransactionManager().enlistAfterCompletion(tx); return tx; } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProvider.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProvider.java index 9c7aa1e9f27..07c1953450c 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProvider.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProvider.java @@ -18,30 +18,26 @@ package org.keycloak.models.sessions.infinispan.remote; import java.lang.invoke.MethodHandles; +import java.util.AbstractMap; import java.util.ArrayList; -import java.util.Collection; import java.util.Collections; -import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Optional; -import java.util.UUID; -import java.util.concurrent.CompletionStage; +import java.util.Set; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Predicate; import java.util.stream.Collectors; import java.util.stream.Stream; import io.reactivex.rxjava3.core.Flowable; -import org.infinispan.client.hotrod.MetadataValue; +import io.reactivex.rxjava3.core.Maybe; import org.infinispan.client.hotrod.RemoteCache; -import org.infinispan.commons.util.concurrent.CompletableFutures; import org.infinispan.commons.util.concurrent.CompletionStages; import org.jboss.logging.Logger; import org.keycloak.cluster.ClusterProvider; import org.keycloak.common.Profile; -import org.keycloak.common.util.Time; import org.keycloak.models.AuthenticatedClientSessionModel; import org.keycloak.models.ClientModel; import org.keycloak.models.KeycloakSession; @@ -54,12 +50,13 @@ import org.keycloak.models.light.LightweightUserAdapter; import org.keycloak.models.session.UserSessionPersisterProvider; import org.keycloak.models.sessions.infinispan.changes.remote.updater.BaseUpdater; import org.keycloak.models.sessions.infinispan.changes.remote.updater.client.AuthenticatedClientSessionUpdater; -import org.keycloak.models.sessions.infinispan.changes.remote.updater.user.ClientSessionMappingAdapter; -import org.keycloak.models.sessions.infinispan.changes.remote.updater.user.ClientSessionProvider; import org.keycloak.models.sessions.infinispan.changes.remote.updater.user.UserSessionUpdater; -import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionEntity; -import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionStore; -import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; +import org.keycloak.models.sessions.infinispan.query.ClientSessionQueries; +import org.keycloak.models.sessions.infinispan.query.QueryHelper; +import org.keycloak.models.sessions.infinispan.query.UserSessionQueries; import org.keycloak.models.sessions.infinispan.remote.transaction.ClientSessionChangeLogTransaction; import org.keycloak.models.sessions.infinispan.remote.transaction.UseSessionChangeLogTransaction; import org.keycloak.models.sessions.infinispan.remote.transaction.UserSessionTransaction; @@ -74,6 +71,7 @@ import static org.keycloak.models.Constants.SESSION_NOTE_LIGHTWEIGHT_USER; public class RemoteUserSessionProvider implements UserSessionProvider { private static final Logger log = Logger.getLogger(MethodHandles.lookup().lookupClass()); + private static final int MAX_CONCURRENT_REQUESTS = 16; private final KeycloakSession session; private final UserSessionTransaction transaction; @@ -87,14 +85,13 @@ public class RemoteUserSessionProvider implements UserSessionProvider { @Override public AuthenticatedClientSessionModel createClientSession(RealmModel realm, ClientModel client, UserSessionModel userSession) { - var transaction = getClientSessionTransaction(false); - var clientSessionId = UUID.randomUUID(); - var entity = AuthenticatedClientSessionEntity.create(clientSessionId, realm, client, userSession); - var model = transaction.create(clientSessionId, entity); + var clientTx = getClientSessionTransaction(false); + var key = new ClientSessionKey(userSession.getId(), client.getId()); + var entity = RemoteAuthenticatedClientSessionEntity.create(key, realm.getId(), userSession); + var model = clientTx.create(key, entity); if (!model.isInitialized()) { - model.initialize(userSession, client, transaction); + model.initialize(userSession, client, clientTx); } - userSession.getAuthenticatedClientSessions().put(client.getId(), model); return model; } @@ -103,13 +100,13 @@ public class RemoteUserSessionProvider implements UserSessionProvider { if (clientSessionId == null) { return null; } - var transaction = getClientSessionTransaction(offline); - var updater = transaction.get(UUID.fromString(clientSessionId)); + var clientTx = getClientSessionTransaction(offline); + var updater = clientTx.get(new ClientSessionKey(userSession.getId(), client.getId())); if (updater == null) { return null; } if (!updater.isInitialized()) { - updater.initialize(userSession, client, transaction); + updater.initialize(userSession, client, clientTx); } return updater; } @@ -120,8 +117,8 @@ public class RemoteUserSessionProvider implements UserSessionProvider { id = KeycloakModelUtils.generateId(); } - var entity = UserSessionEntity.create(id, realm, user, loginUsername, ipAddress, authMethod, rememberMe, brokerSessionId, brokerUserId); - var updater = transaction.getUserSessions().create(id, entity); + var entity = RemoteUserSessionEntity.create(id, realm, user, loginUsername, ipAddress, authMethod, rememberMe, brokerSessionId, brokerUserId); + var updater = getUserSessionTransaction(false).create(id, entity); return initUserSessionUpdater(updater, persistenceState, realm, user, false); } @@ -132,28 +129,30 @@ public class RemoteUserSessionProvider implements UserSessionProvider { @Override public Stream getUserSessionsStream(RealmModel realm, UserModel user) { - return StreamsUtil.closing(streamUserSessions(new UserAndRealmPredicate(realm.getId(), user.getId()), realm, user, false)); + return StreamsUtil.closing(streamUserSessionByUserId(realm, user, false)); } @Override public Stream getUserSessionsStream(RealmModel realm, ClientModel client) { - return StreamsUtil.closing(streamUserSessions(new ClientAndRealmPredicate(realm.getId(), client.getId()), realm, null, false)); + return StreamsUtil.closing(streamUserSessionByClientId(realm, client.getId(), false, null, null)); } @Override public Stream getUserSessionsStream(RealmModel realm, ClientModel client, Integer firstResult, Integer maxResults) { - return StreamsUtil.paginatedStream(getUserSessionsStream(realm, client).sorted(Comparator.comparing(UserSessionModel::getLastSessionRefresh)), firstResult, maxResults); + return StreamsUtil.closing(streamUserSessionByClientId(realm, client.getId(), false, firstResult, maxResults)); } @Override public Stream getUserSessionByBrokerUserIdStream(RealmModel realm, String brokerUserId) { - return StreamsUtil.closing(streamUserSessions(new BrokerUserIdAndRealmPredicate(realm.getId(), brokerUserId), realm, null, false)); + return StreamsUtil.closing(streamUserSessionByBrokerUserId(realm, brokerUserId, false)); } @Override public UserSessionModel getUserSessionByBrokerSessionId(RealmModel realm, String brokerSessionId) { - return StreamsUtil.closing(streamUserSessions(new BrokerSessionIdAndRealmPredicate(realm.getId(), brokerSessionId), realm, null, false)) - .findFirst() + var userTx = getUserSessionTransaction(false); + var query = UserSessionQueries.searchByBrokerSessionId(userTx.getCache(), realm.getId(), brokerSessionId); + return QueryHelper.fetchSingle(query, userTx::wrapFromProjection) + .map(session -> initUserSessionFromQuery(session, realm, null, false)) .orElse(null); } @@ -165,22 +164,14 @@ public class RemoteUserSessionProvider implements UserSessionProvider { @Override public long getActiveUserSessions(RealmModel realm, ClientModel client) { - return StreamsUtil.closing(getUserSessionsStream(realm, client)).count(); + return computeUserSessionCount(realm, client, false); } @Override public Map getActiveClientSessionStats(RealmModel realm, boolean offline) { - var userSessions = getUserSessionTransaction(offline); - return Flowable.fromPublisher(userSessions.getCache().publishEntriesWithMetadata(null, batchSize)) - .filter(new RealmPredicate(realm.getId())) - .map(Map.Entry::getValue) - .map(MetadataValue::getValue) - .map(UserSessionEntity::getAuthenticatedClientSessions) - .map(AuthenticatedClientSessionStore::keySet) - .map(Collection::stream) - .flatMap(Flowable::fromStream) - .collect(Collectors.groupingBy(Function.identity(), Collectors.counting())) - .blockingGet(); + var query = ClientSessionQueries.activeClientCount(getClientSessionTransaction(offline).getCache()); + return QueryHelper.streamAll(query, batchSize, QueryHelper.PROJECTION_TO_STRING_LONG_ENTRY) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } @Override @@ -190,7 +181,7 @@ public class RemoteUserSessionProvider implements UserSessionProvider { @Override public void removeUserSessions(RealmModel realm, UserModel user) { - getUserSessionsStream(realm, user).forEach(s -> removeUserSession(realm, s)); + transaction.removeAllSessionByUserId(realm.getId(), user.getId()); } @Override @@ -227,13 +218,8 @@ public class RemoteUserSessionProvider implements UserSessionProvider { @Override public UserSessionModel createOfflineUserSession(UserSessionModel userSession) { - var entity = UserSessionEntity.createFromModel(userSession); - - int currentTime = Time.currentTime(); - entity.setStarted(currentTime); - entity.setLastSessionRefresh(currentTime); - - var updater = getUserSessionTransaction(true).create(entity.getId(), entity); + var entity = RemoteUserSessionEntity.createFromModel(userSession); + var updater = getUserSessionTransaction(true).create(userSession.getId(), entity); return initUserSessionUpdater(updater, userSession.getPersistenceState(), userSession.getRealm(), userSession.getUser(), true); } @@ -249,34 +235,34 @@ public class RemoteUserSessionProvider implements UserSessionProvider { @Override public AuthenticatedClientSessionModel createOfflineClientSession(AuthenticatedClientSessionModel clientSession, UserSessionModel offlineUserSession) { - var transaction = getClientSessionTransaction(true); - var entity = AuthenticatedClientSessionEntity.createFromModel(clientSession); - var model = transaction.create(entity.getId(), entity); + var clientTx = getClientSessionTransaction(true); + var key = new ClientSessionKey(offlineUserSession.getId(), clientSession.getClient().getId()); + var entity = RemoteAuthenticatedClientSessionEntity.createFromModel(key, clientSession); + var model = clientTx.create(key, entity); if (!model.isInitialized()) { - model.initialize(offlineUserSession, clientSession.getClient(), transaction); + model.initialize(offlineUserSession, clientSession.getClient(), clientTx); } - offlineUserSession.getAuthenticatedClientSessions().put(clientSession.getClient().getId(), model); return model; } @Override public Stream getOfflineUserSessionsStream(RealmModel realm, UserModel user) { - return StreamsUtil.closing(streamUserSessions(new UserAndRealmPredicate(realm.getId(), user.getId()), realm, user, true)); + return StreamsUtil.closing(streamUserSessionByUserId(realm, user, true)); } @Override public Stream getOfflineUserSessionByBrokerUserIdStream(RealmModel realm, String brokerUserId) { - return StreamsUtil.closing(streamUserSessions(new BrokerUserIdAndRealmPredicate(realm.getId(), brokerUserId), realm, null, true)); + return StreamsUtil.closing(streamUserSessionByBrokerUserId(realm, brokerUserId, true)); } @Override public long getOfflineSessionsCount(RealmModel realm, ClientModel client) { - return StreamsUtil.closing(streamUserSessions(new ClientAndRealmPredicate(realm.getId(), client.getId()), realm, null, true)).count(); + return computeUserSessionCount(realm, client, true); } @Override public Stream getOfflineUserSessionsStream(RealmModel realm, ClientModel client, Integer firstResult, Integer maxResults) { - return StreamsUtil.closing(StreamsUtil.paginatedStream(streamUserSessions(new ClientAndRealmPredicate(realm.getId(), client.getId()), realm, null, true), firstResult, maxResults)); + return StreamsUtil.closing(streamUserSessionByClientId(realm, client.getId(), true, firstResult, maxResults)); } @Override @@ -327,13 +313,14 @@ public class RemoteUserSessionProvider implements UserSessionProvider { var stage = CompletionStages.aggregateCompletionStage(); database.loadUserSessionsStream(-1, batchSize, offline, "") .forEach(userSessionModel -> { - var userSessionEntity = UserSessionEntity.createFromModel(userSessionModel); + var userSessionEntity = RemoteUserSessionEntity.createFromModel(userSessionModel); stage.dependsOn(userSessionCache.putIfAbsentAsync(userSessionModel.getId(), userSessionEntity)); userSessionBuffer.add(userSessionModel.getId()); for (var clientSessionModel : userSessionModel.getAuthenticatedClientSessions().values()) { + var clientSessionKey = new ClientSessionKey(userSessionModel.getId(), clientSessionModel.getClient().getId()); clientSessionBuffer.add(Map.entry(userSessionModel.getId(), clientSessionModel.getId())); - var clientSessionEntity = AuthenticatedClientSessionEntity.createFromModel(clientSessionModel); - stage.dependsOn(clientSessionCache.putIfAbsentAsync(clientSessionEntity.getId(), clientSessionEntity)); + var clientSessionEntity = RemoteAuthenticatedClientSessionEntity.createFromModel(clientSessionKey, clientSessionModel); + stage.dependsOn(clientSessionCache.putIfAbsentAsync(clientSessionKey, clientSessionEntity)); } }); CompletionStages.join(stage.freeze()); @@ -365,64 +352,51 @@ public class RemoteUserSessionProvider implements UserSessionProvider { if (updater.isInitialized()) { return updater; } - UserModel user = session.users().getUserById(realm, updater.getValue().getUser()); + UserModel user = session.users().getUserById(realm, updater.getValue().getUserId()); return initUserSessionUpdater(updater, UserSessionModel.SessionPersistenceState.PERSISTENT, realm, user, offline); } private void internalRemoveUserSession(UserSessionModel userSession, boolean offline) { - var clientSessionTransaction = getClientSessionTransaction(offline); - var userSessionTransaction = getUserSessionTransaction(offline); - userSession.getAuthenticatedClientSessions().values() - .stream() - .filter(Objects::nonNull) // we need to filter, it may not be a UserSessionUpdater class. - .map(AuthenticatedClientSessionModel::getId) - .filter(Objects::nonNull) // we need to filter, it may not be a AuthenticatedClientSessionUpdater class. - .map(UUID::fromString) - .forEach(clientSessionTransaction::remove); - userSessionTransaction.remove(userSession.getId()); - } - - private Stream streamUserSessions(InternalUserSessionPredicate predicate, RealmModel realm, UserModel user, boolean offline) { - var userSessions = getUserSessionTransaction(offline); - return Flowable.fromPublisher(userSessions.getCache().publishEntriesWithMetadata(null, batchSize)) - .filter(predicate) - .map(userSessions::wrap) - .map(s -> initFromStream(s, realm, user, offline)) - .filter(Optional::isPresent) - .map(Optional::get) - .map(UserSessionModel.class::cast) - .blockingStream(batchSize); + transaction.removeUserSessionById(userSession.getId(), offline); } private UseSessionChangeLogTransaction getUserSessionTransaction(boolean offline) { - return offline ? transaction.getOfflineUserSessions() : transaction.getUserSessions(); + return transaction.getUserSessions(offline); } private ClientSessionChangeLogTransaction getClientSessionTransaction(boolean offline) { - return offline ? transaction.getOfflineClientSessions() : transaction.getClientSessions(); + return transaction.getClientSessions(offline); } - private Optional initFromStream(UserSessionUpdater updater, RealmModel realm, UserModel user, boolean offline) { - if (updater.isInitialized()) { - return Optional.of(updater); - } + private UserSessionUpdater initUserSessionFromQuery(UserSessionUpdater updater, RealmModel realm, UserModel user, boolean offline) { + assert updater != null; assert realm != null; - if (user == null) { - user = session.users().getUserById(realm, updater.getValue().getUser()); + if (updater.isDeleted()) { + return null; } - return Optional.ofNullable(initUserSessionUpdater(updater, UserSessionModel.SessionPersistenceState.PERSISTENT, realm, user, offline)); + if (updater.isInitialized()) { + return updater; + } + if (user == null) { + user = session.users().getUserById(realm, updater.getValue().getUserId()); + } + return initUserSessionUpdater(updater, UserSessionModel.SessionPersistenceState.PERSISTENT, realm, user, offline); + } + + private Maybe maybeInitUserSessionFromQuery(UserSessionUpdater updater, RealmModel realm, boolean offline) { + var model = initUserSessionFromQuery(updater, realm, null, offline); + return model == null ? Maybe.empty() : Maybe.just(model); } private UserSessionUpdater initUserSessionUpdater(UserSessionUpdater updater, UserSessionModel.SessionPersistenceState persistenceState, RealmModel realm, UserModel user, boolean offline) { - var provider = new RemoteClientSessionAdapterProvider(getClientSessionTransaction(offline), updater); if (user instanceof LightweightUserAdapter) { - updater.initialize(persistenceState, realm, user, provider); + updater.initialize(persistenceState, realm, user, new ClientSessionMapping(updater)); return checkExpiration(updater); } // copied from org.keycloak.models.sessions.infinispan.InfinispanUserSessionProvider if (Profile.isFeatureEnabled(Profile.Feature.TRANSIENT_USERS) && updater.getNotes().containsKey(SESSION_NOTE_LIGHTWEIGHT_USER)) { LightweightUserAdapter lua = LightweightUserAdapter.fromString(session, realm, updater.getNotes().get(SESSION_NOTE_LIGHTWEIGHT_USER)); - updater.initialize(persistenceState, realm, lua, provider); + updater.initialize(persistenceState, realm, lua, new ClientSessionMapping(updater)); lua.setUpdateHandler(lua1 -> { if (lua == lua1) { // Ensure there is no conflicting user model, only the latest lightweight user can be used updater.setNote(SESSION_NOTE_LIGHTWEIGHT_USER, lua1.serialize()); @@ -436,11 +410,64 @@ public class RemoteUserSessionProvider implements UserSessionProvider { internalRemoveUserSession(updater, offline); return null; } - updater.initialize(persistenceState, realm, user, provider); + updater.initialize(persistenceState, realm, user, new ClientSessionMapping(updater)); return checkExpiration(updater); } - private > T checkExpiration(T updater) { + private AuthenticatedClientSessionModel initClientSessionUpdater(AuthenticatedClientSessionUpdater updater, UserSessionUpdater userSession) { + if (updater == null || updater.isDeleted()) { + return null; + } + var client = userSession.getRealm().getClientById(updater.getKey().clientId()); + if (client == null) { + updater.markDeleted(); + return null; + } + if (updater.isInitialized()) { + return updater; + } + updater.initialize(userSession, client, getClientSessionTransaction(userSession.isOffline())); + return checkExpiration(updater); + } + + private long computeUserSessionCount(RealmModel realm, ClientModel client, boolean offline) { + var query = ClientSessionQueries.countClientSessions(getClientSessionTransaction(offline).getCache(), realm.getId(), client.getId()); + return QueryHelper.fetchSingle(query, QueryHelper.SINGLE_PROJECTION_TO_LONG).orElse(0L); + } + + private Stream streamUserSessionByUserId(RealmModel realm, UserModel user, boolean offline) { + var userTx = getUserSessionTransaction(offline); + var query = UserSessionQueries.searchByUserId(userTx.getCache(), realm.getId(), user.getId()); + return QueryHelper.streamAll(query, batchSize, userTx::wrapFromProjection) + .map(session -> initUserSessionFromQuery(session, realm, user, offline)) + .filter(Objects::nonNull) + .map(UserSessionModel.class::cast); + } + + private Stream streamUserSessionByBrokerUserId(RealmModel realm, String brokerUserId, boolean offline) { + var userTx = getUserSessionTransaction(offline); + var query = UserSessionQueries.searchByBrokerUserId(userTx.getCache(), realm.getId(), brokerUserId); + return QueryHelper.streamAll(query, batchSize, userTx::wrapFromProjection) + .map(session -> initUserSessionFromQuery(session, realm, null, offline)) + .filter(Objects::nonNull) + .map(UserSessionModel.class::cast); + } + + private Stream streamUserSessionByClientId(RealmModel realm, String clientId, boolean offline, Integer offset, Integer maxResults) { + var userSessionIdQuery = ClientSessionQueries.fetchUserSessionIdForClientId(getClientSessionTransaction(offline).getCache(), realm.getId(), clientId); + if (offset != null) { + userSessionIdQuery.startOffset(offset); + } + userSessionIdQuery.maxResults(maxResults == null ? Integer.MAX_VALUE : maxResults); + var userSessionTx = getUserSessionTransaction(offline); + return Flowable.fromIterable(QueryHelper.toCollection(userSessionIdQuery, QueryHelper.SINGLE_PROJECTION_TO_STRING)) + .flatMapMaybe(userSessionTx::maybeGet, false, MAX_CONCURRENT_REQUESTS) + .concatMapMaybe(session -> maybeInitUserSessionFromQuery(session, realm, offline)) + .map(UserSessionModel.class::cast) + .blockingStream(batchSize); + } + + private static > T checkExpiration(T updater) { var expiration = updater.computeExpiration(); if (expiration.isExpired()) { updater.markDeleted(); @@ -449,115 +476,89 @@ public class RemoteUserSessionProvider implements UserSessionProvider { return updater; } - private record RealmPredicate(String realmId) implements InternalUserSessionPredicate { + private class ClientSessionMapping extends AbstractMap implements Consumer { - @Override - public boolean testUserSession(UserSessionEntity userSession) { - return Objects.equals(userSession.getRealmId(), realmId); - } - } - - private interface InternalUserSessionPredicate extends io.reactivex.rxjava3.functions.Predicate>> { - - @Override - default boolean test(Map.Entry> e) { - return testUserSession(e.getValue().getValue()); - } - - boolean testUserSession(UserSessionEntity userSession); - } - - private record UserAndRealmPredicate(String realmId, String userId) implements InternalUserSessionPredicate { - - @Override - public boolean testUserSession(UserSessionEntity userSession) { - return Objects.equals(userSession.getRealmId(), realmId) && Objects.equals(userSession.getUser(), userId); - } - - } - - private record ClientAndRealmPredicate(String realmId, String clientId) implements InternalUserSessionPredicate { - - @Override - public boolean testUserSession(UserSessionEntity userSession) { - return Objects.equals(userSession.getRealmId(), realmId) && userSession.getAuthenticatedClientSessions().containsKey(clientId); - } - } - - private record BrokerUserIdAndRealmPredicate(String realmId, String brokerUserId) implements InternalUserSessionPredicate { - - @Override - public boolean testUserSession(UserSessionEntity userSession) { - return Objects.equals(userSession.getRealmId(), realmId) && Objects.equals(userSession.getBrokerUserId(), brokerUserId); - } - } - - private record BrokerSessionIdAndRealmPredicate(String realmId, - String brokeSessionId) implements InternalUserSessionPredicate { - - @Override - public boolean testUserSession(UserSessionEntity userSession) { - return Objects.equals(userSession.getRealmId(), realmId) && Objects.equals(userSession.getBrokerSessionId(), brokeSessionId); - } - } - - private class RemoteClientSessionAdapterProvider implements ClientSessionProvider, UserSessionUpdater.ClientSessionAdapterFactory { - - private final ClientSessionChangeLogTransaction transaction; private final UserSessionUpdater userSession; + private boolean coldCache = true; - private RemoteClientSessionAdapterProvider(ClientSessionChangeLogTransaction transaction, UserSessionUpdater userSession) { - this.transaction = transaction; + ClientSessionMapping(UserSessionUpdater userSession) { this.userSession = userSession; } @Override - public AuthenticatedClientSessionModel getClientSession(String clientId, UUID clientSessionId) { - if (clientId == null || clientSessionId == null) { - return null; - } - var client = userSession.getRealm().getClientById(clientId); - if (client == null) { - return null; - } - return initialize(client, transaction.get(clientSessionId)); + public void clear() { + getTransaction().removeByUserSessionId(getUserSessionId()); } @Override - public CompletionStage getClientSessionAsync(String clientId, UUID clientSessionId) { - if (clientId == null || clientSessionId == null) { - return CompletableFutures.completedNull(); - } - var client = userSession.getRealm().getClientById(clientId); - if (client == null) { - return CompletableFutures.completedNull(); - } - return transaction.getAsync(clientSessionId).thenApply(updater -> initialize(client, updater)); + public AuthenticatedClientSessionModel get(Object key) { + var updater = getTransaction().get(keyForClientId(key)); + return initClientSessionUpdater(updater, userSession); } @Override - public void removeClientSession(UUID clientSessionId) { - if (clientSessionId == null) { - return; - } - transaction.remove(clientSessionId); - } - - private AuthenticatedClientSessionModel initialize(ClientModel client, AuthenticatedClientSessionUpdater updater) { - if (updater == null) { - return null; - } - if (updater.isInitialized()) { - return updater; - } - updater.initialize(userSession, client, transaction); - return checkExpiration(updater); + public AuthenticatedClientSessionModel remove(Object key) { + getTransaction().remove(keyForClientId(key)); + return null; } @Override - public ClientSessionMappingAdapter create(AuthenticatedClientSessionStore clientSessionStore) { - return new ClientSessionMappingAdapter(clientSessionStore, this); + public boolean containsKey(Object key) { + return get(key) != null; + } + + @SuppressWarnings("NullableProblems") + @Override + public Set> entrySet() { + if (coldCache) { + fetchAndCacheClientSessions(); + coldCache = false; + } + // iterate from the locally cached data. + return getTransaction().getClientSessions() + .filter(this::isFromUserSession) + .map(this::initialize) + .filter(Objects::nonNull) + .map(RemoteUserSessionProvider::toMapEntry) + .collect(Collectors.toSet()); + } + + private ClientSessionKey keyForClientId(String clientId) { + return new ClientSessionKey(getUserSessionId(), clientId); + } + + private ClientSessionKey keyForClientId(Object clientId) { + return keyForClientId(String.valueOf(clientId)); + } + + private void fetchAndCacheClientSessions() { + var query = ClientSessionQueries.fetchClientSessions(getTransaction().getCache(), getUserSessionId()); + QueryHelper.streamAll(query, batchSize, Function.identity()).forEach(this); + } + + @Override + public void accept(Object[] projections) { + getTransaction().wrapFromProjection(projections); + } + + private ClientSessionChangeLogTransaction getTransaction() { + return getClientSessionTransaction(userSession.isOffline()); + } + + private String getUserSessionId() { + return userSession.getKey(); + } + + private boolean isFromUserSession(AuthenticatedClientSessionUpdater updater) { + return Objects.equals(getUserSessionId(), updater.getValue().getUserSessionId()); + } + + private AuthenticatedClientSessionModel initialize(AuthenticatedClientSessionUpdater updater) { + return initClientSessionUpdater(updater, userSession); } } + private static Map.Entry toMapEntry(AuthenticatedClientSessionModel model) { + return Map.entry(model.getClient().getId(), model); + } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProviderFactory.java index d29412ff1e2..0458dbf5173 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProviderFactory.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteUserSessionProviderFactory.java @@ -1,7 +1,7 @@ package org.keycloak.models.sessions.infinispan.remote; import java.util.List; -import java.util.UUID; +import java.util.concurrent.Executor; import org.infinispan.client.hotrod.RemoteCache; import org.keycloak.Config; @@ -16,8 +16,9 @@ import org.keycloak.models.UserSessionProviderFactory; import org.keycloak.models.session.UserSessionPersisterProvider; import org.keycloak.models.sessions.infinispan.changes.remote.updater.client.AuthenticatedClientSessionUpdater; import org.keycloak.models.sessions.infinispan.changes.remote.updater.user.UserSessionUpdater; -import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionEntity; -import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; import org.keycloak.models.sessions.infinispan.remote.transaction.ClientSessionChangeLogTransaction; import org.keycloak.models.sessions.infinispan.remote.transaction.UseSessionChangeLogTransaction; import org.keycloak.models.sessions.infinispan.remote.transaction.UserSessionTransaction; @@ -45,7 +46,7 @@ public class RemoteUserSessionProviderFactory implements UserSessionProviderFact @Override public void init(Config.Scope config) { - batchSize = config.getInt(CONFIG_MAX_BATCH_SIZE, DEFAULT_BATCH_SIZE); + batchSize = Math.max(1, config.getInt(CONFIG_MAX_BATCH_SIZE, DEFAULT_BATCH_SIZE)); } @Override @@ -101,11 +102,12 @@ public class RemoteUserSessionProviderFactory implements UserSessionProviderFact return; } InfinispanConnectionProvider connections = session.getProvider(InfinispanConnectionProvider.class); - RemoteCache userSessionCache = connections.getRemoteCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME); - RemoteCache offlineUserSessionsCache = connections.getRemoteCache(InfinispanConnectionProvider.OFFLINE_USER_SESSION_CACHE_NAME); - RemoteCache clientSessionCache = connections.getRemoteCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME); - RemoteCache offlineClientSessionsCache = connections.getRemoteCache(InfinispanConnectionProvider.OFFLINE_CLIENT_SESSION_CACHE_NAME); - remoteCacheHolder = new RemoteCacheHolder(userSessionCache, offlineUserSessionsCache, clientSessionCache, offlineClientSessionsCache); + RemoteCache userSessionCache = connections.getRemoteCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME); + RemoteCache offlineUserSessionsCache = connections.getRemoteCache(InfinispanConnectionProvider.OFFLINE_USER_SESSION_CACHE_NAME); + RemoteCache clientSessionCache = connections.getRemoteCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME); + RemoteCache offlineClientSessionsCache = connections.getRemoteCache(InfinispanConnectionProvider.OFFLINE_CLIENT_SESSION_CACHE_NAME); + var executor = connections.getExecutor("query-delete"); + remoteCacheHolder = new RemoteCacheHolder(userSessionCache, offlineUserSessionsCache, clientSessionCache, offlineClientSessionsCache, executor); } private UserSessionTransaction createTransaction(KeycloakSession session) { @@ -127,16 +129,17 @@ public class RemoteUserSessionProviderFactory implements UserSessionProviderFact } private record RemoteCacheHolder( - RemoteCache userSession, - RemoteCache offlineUserSession, - RemoteCache clientSession, - RemoteCache offlineClientSession) { + RemoteCache userSession, + RemoteCache offlineUserSession, + RemoteCache clientSession, + RemoteCache offlineClientSession, + Executor executor) { - RemoteCache userSessionCache(boolean offline) { + RemoteCache userSessionCache(boolean offline) { return offline ? offlineUserSession : userSession; } - RemoteCache clientSessionCache(boolean offline) { + RemoteCache clientSessionCache(boolean offline) { return offline ? offlineClientSession : clientSession; } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/ClientSessionChangeLogTransaction.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/ClientSessionChangeLogTransaction.java index 18ee33ef7b8..44398a6db35 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/ClientSessionChangeLogTransaction.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/ClientSessionChangeLogTransaction.java @@ -17,23 +17,51 @@ package org.keycloak.models.sessions.infinispan.remote.transaction; -import java.util.UUID; +import java.util.stream.Stream; import org.infinispan.client.hotrod.RemoteCache; -import org.keycloak.models.sessions.infinispan.changes.remote.remover.iteration.ByRealmIdConditionalRemover; +import org.keycloak.models.sessions.infinispan.changes.remote.remover.query.ClientSessionQueryConditionalRemover; +import org.keycloak.models.sessions.infinispan.changes.remote.updater.BaseUpdater; import org.keycloak.models.sessions.infinispan.changes.remote.updater.UpdaterFactory; import org.keycloak.models.sessions.infinispan.changes.remote.updater.client.AuthenticatedClientSessionUpdater; -import org.keycloak.models.sessions.infinispan.entities.AuthenticatedClientSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; /** * Syntactic sugar for * {@code RemoteChangeLogTransaction>} */ -public class ClientSessionChangeLogTransaction extends RemoteChangeLogTransaction> { +public class ClientSessionChangeLogTransaction extends RemoteChangeLogTransaction { - public ClientSessionChangeLogTransaction(UpdaterFactory factory, RemoteCache cache) { - super(factory, cache, new ByRealmIdConditionalRemover<>()); + public ClientSessionChangeLogTransaction(UpdaterFactory factory, RemoteCache cache) { + super(factory, cache, new ClientSessionQueryConditionalRemover()); } + /** + * Wraps a Query project results, where the first argument is the entity, and the second the version. + */ + public void wrapFromProjection(Object[] projection) { + assert projection.length == 2; + RemoteAuthenticatedClientSessionEntity entity = (RemoteAuthenticatedClientSessionEntity) projection[0]; + wrap(entity.createCacheKey(), entity, (long) projection[1]); + } + + /** + * Remove all client sessions belonging to the user session. + */ + public void removeByUserSessionId(String userSessionId) { + getConditionalRemover().removeByUserSessionId(userSessionId); + // make cached entities as deleted too + getClientSessions() + .filter(getConditionalRemover()::willRemove) + .forEach(BaseUpdater::markDeleted); + } + + /** + * @return A stream with all currently cached {@link AuthenticatedClientSessionUpdater} in this transaction. + */ + public Stream getClientSessions() { + return getCachedEntities().values().stream(); + } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteCacheAndExecutor.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteCacheAndExecutor.java deleted file mode 100644 index 07f27cbd679..00000000000 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteCacheAndExecutor.java +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright 2024 Red Hat, Inc. and/or its affiliates - * and other contributors as indicated by the @author tags. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.keycloak.models.sessions.infinispan.remote.transaction; - -import java.util.concurrent.Executor; - -import org.infinispan.client.hotrod.RemoteCache; -import org.keycloak.connections.infinispan.InfinispanConnectionProvider; -import org.keycloak.models.KeycloakSession; -import org.keycloak.models.KeycloakSessionFactory; - -public record RemoteCacheAndExecutor(RemoteCache cache, Executor executor) { - - public static RemoteCacheAndExecutor create(KeycloakSession session, String cacheName) { - var connection = session.getProvider(InfinispanConnectionProvider.class); - return new RemoteCacheAndExecutor<>(connection.getRemoteCache(cacheName), connection.getExecutor(cacheName + "-query-delete")); - } - - public static RemoteCacheAndExecutor create(KeycloakSessionFactory factory, String cacheName) { - try (var session = factory.create()) { - return create(session, cacheName); - } - } -} diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteChangeLogTransaction.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteChangeLogTransaction.java index 1df7151e442..631b24db376 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteChangeLogTransaction.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/RemoteChangeLogTransaction.java @@ -182,6 +182,14 @@ public class RemoteChangeLogTransaction, R extends return entityChanges.computeIfAbsent(entry.getKey(), k -> factory.wrapFromCache(k, entry.getValue())); } + public T wrap(K key, V value, long version) { + return entityChanges.computeIfAbsent(key, k -> factory.wrapFromCache(k, value, version)); + } + + protected Map getCachedEntities() { + return entityChanges; + } + private T onEntityFromCache(K key, MetadataValue entity) { if (entity == null) { return null; diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UseSessionChangeLogTransaction.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UseSessionChangeLogTransaction.java index df91b521302..89e6c1f5ca3 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UseSessionChangeLogTransaction.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UseSessionChangeLogTransaction.java @@ -17,21 +17,32 @@ package org.keycloak.models.sessions.infinispan.remote.transaction; +import io.reactivex.rxjava3.core.Maybe; import org.infinispan.client.hotrod.RemoteCache; -import org.keycloak.models.sessions.infinispan.changes.remote.remover.iteration.ByRealmIdConditionalRemover; +import org.keycloak.models.sessions.infinispan.changes.remote.remover.query.UserSessionQueryConditionalRemover; import org.keycloak.models.sessions.infinispan.changes.remote.updater.UpdaterFactory; import org.keycloak.models.sessions.infinispan.changes.remote.updater.user.UserSessionUpdater; -import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; /** * Syntactic sugar for * {@code RemoteChangeLogTransaction>} */ -public class UseSessionChangeLogTransaction extends RemoteChangeLogTransaction> { +public class UseSessionChangeLogTransaction extends RemoteChangeLogTransaction { - public UseSessionChangeLogTransaction(UpdaterFactory factory, RemoteCache cache) { - super(factory, cache, new ByRealmIdConditionalRemover<>()); + public UseSessionChangeLogTransaction(UpdaterFactory factory, RemoteCache cache) { + super(factory, cache, new UserSessionQueryConditionalRemover()); + } + + public UserSessionUpdater wrapFromProjection(Object[] projection) { + assert projection.length == 2; + RemoteUserSessionEntity entity = (RemoteUserSessionEntity) projection[0]; + return wrap(entity.getUserSessionId(), entity, (long) projection[1]); + } + + public Maybe maybeGet(String userSessionId) { + return Maybe.fromCompletionStage(getAsync(userSessionId)); } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UserSessionTransaction.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UserSessionTransaction.java index 1539e32febe..e12c4f4fe7e 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UserSessionTransaction.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/transaction/UserSessionTransaction.java @@ -69,20 +69,12 @@ public class UserSessionTransaction extends AbstractKeycloakTransaction { offlineClientSessions.rollback(); } - public ClientSessionChangeLogTransaction getClientSessions() { - return clientSessions; + public ClientSessionChangeLogTransaction getClientSessions(boolean offline) { + return offline ? offlineClientSessions : clientSessions; } - public UseSessionChangeLogTransaction getUserSessions() { - return userSessions; - } - - public ClientSessionChangeLogTransaction getOfflineClientSessions() { - return offlineClientSessions; - } - - public UseSessionChangeLogTransaction getOfflineUserSessions() { - return offlineUserSessions; + public UseSessionChangeLogTransaction getUserSessions(boolean offline) { + return offline ? offlineUserSessions : userSessions; } public void removeAllSessionsByRealmId(String realmId) { @@ -96,4 +88,14 @@ public class UserSessionTransaction extends AbstractKeycloakTransaction { clientSessions.getConditionalRemover().removeByRealmId(realmId); userSessions.getConditionalRemover().removeByRealmId(realmId); } + + public void removeAllSessionByUserId(String realmId, String userId) { + userSessions.getConditionalRemover().removeByUserId(realmId, userId); + clientSessions.getConditionalRemover().removeByUserId(realmId, userId); + } + + public void removeUserSessionById(String userSessionId, boolean offline) { + getUserSessions(offline).remove(userSessionId); + getClientSessions(offline).removeByUserSessionId(userSessionId); + } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/util/SessionTimeouts.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/util/SessionTimeouts.java index d661073a3df..5212e9e3788 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/util/SessionTimeouts.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/util/SessionTimeouts.java @@ -19,6 +19,7 @@ package org.keycloak.models.sessions.infinispan.util; import java.util.concurrent.TimeUnit; + import org.keycloak.common.util.Time; import org.keycloak.models.ClientModel; import org.keycloak.models.RealmModel; @@ -49,8 +50,15 @@ public class SessionTimeouts { * @return */ public static long getUserSessionLifespanMs(RealmModel realm, ClientModel client, UserSessionEntity userSessionEntity) { - long lifespan = SessionExpirationUtils.calculateUserSessionMaxLifespanTimestamp(false, userSessionEntity.isRememberMe(), - TimeUnit.SECONDS.toMillis(userSessionEntity.getStarted()), realm); + return getUserSessionLifespanMs(realm, false, userSessionEntity.isRememberMe(), userSessionEntity.getStarted()); + } + + public static long getUserSessionLifespanMs(RealmModel realm, boolean offline, boolean rememberMe, int started) { + long lifespan = SessionExpirationUtils.calculateUserSessionMaxLifespanTimestamp(offline, rememberMe, + TimeUnit.SECONDS.toMillis(started), realm); + if (offline && lifespan == IMMORTAL_FLAG) { + return IMMORTAL_FLAG; + } lifespan = lifespan - Time.currentTimeMillis(); if (lifespan <= 0) { return ENTRY_EXPIRED_FLAG; @@ -68,8 +76,11 @@ public class SessionTimeouts { * @return */ public static long getUserSessionMaxIdleMs(RealmModel realm, ClientModel client, UserSessionEntity userSessionEntity) { - long idle = SessionExpirationUtils.calculateUserSessionIdleTimestamp(false, userSessionEntity.isRememberMe(), - TimeUnit.SECONDS.toMillis(userSessionEntity.getLastSessionRefresh()), realm); + return getUserSessionMaxIdleMs(realm, false, userSessionEntity.isRememberMe(), userSessionEntity.getLastSessionRefresh()); + } + + public static long getUserSessionMaxIdleMs(RealmModel realm, boolean offline, boolean rememberMe, int lastSessionRefresh) { + long idle = SessionExpirationUtils.calculateUserSessionIdleTimestamp(offline, rememberMe, TimeUnit.SECONDS.toMillis(lastSessionRefresh), realm); idle = idle - Time.currentTimeMillis(); if (idle <= 0) { return ENTRY_EXPIRED_FLAG; @@ -88,9 +99,15 @@ public class SessionTimeouts { * @return */ public static long getClientSessionLifespanMs(RealmModel realm, ClientModel client, AuthenticatedClientSessionEntity clientSessionEntity) { - long lifespan = SessionExpirationUtils.calculateClientSessionMaxLifespanTimestamp(false, clientSessionEntity.isUserSessionRememberMe(), - TimeUnit.SECONDS.toMillis(clientSessionEntity.getStarted()), TimeUnit.SECONDS.toMillis(clientSessionEntity.getUserSessionStarted()), - realm, client); + return getClientSessionLifespanMs(realm, client, false, clientSessionEntity.isUserSessionRememberMe(), clientSessionEntity.getStarted(), clientSessionEntity.getUserSessionStarted()); + } + + public static long getClientSessionLifespanMs(RealmModel realm, ClientModel client, boolean offline, boolean isUserSessionRememberMe, int started, int userSessionStarted) { + long lifespan = SessionExpirationUtils.calculateClientSessionMaxLifespanTimestamp(offline, isUserSessionRememberMe, + TimeUnit.SECONDS.toMillis(started), TimeUnit.SECONDS.toMillis(userSessionStarted), realm, client); + if (offline && lifespan == IMMORTAL_FLAG) { + return IMMORTAL_FLAG; + } lifespan = lifespan - Time.currentTimeMillis(); if (lifespan <= 0) { return ENTRY_EXPIRED_FLAG; @@ -109,8 +126,12 @@ public class SessionTimeouts { * @return */ public static long getClientSessionMaxIdleMs(RealmModel realm, ClientModel client, AuthenticatedClientSessionEntity clientSessionEntity) { - long idle = SessionExpirationUtils.calculateClientSessionIdleTimestamp(false, clientSessionEntity.isUserSessionRememberMe(), - TimeUnit.SECONDS.toMillis(clientSessionEntity.getTimestamp()), realm, client); + return getClientSessionMaxIdleMs(realm, client, false, clientSessionEntity.isUserSessionRememberMe(), clientSessionEntity.getTimestamp()); + } + + public static long getClientSessionMaxIdleMs(RealmModel realm, ClientModel client, boolean offline, boolean isUserSessionRememberMe, int timestamp) { + long idle = SessionExpirationUtils.calculateClientSessionIdleTimestamp(offline, isUserSessionRememberMe, + TimeUnit.SECONDS.toMillis(timestamp), realm, client); idle = idle - Time.currentTimeMillis(); if (idle <= 0) { return ENTRY_EXPIRED_FLAG; @@ -129,16 +150,7 @@ public class SessionTimeouts { * @return */ public static long getOfflineSessionLifespanMs(RealmModel realm, ClientModel client, UserSessionEntity userSessionEntity) { - long lifespan = SessionExpirationUtils.calculateUserSessionMaxLifespanTimestamp(true, userSessionEntity.isRememberMe(), - TimeUnit.SECONDS.toMillis(userSessionEntity.getStarted()), realm); - if (lifespan == -1L) { - return lifespan; - } - lifespan = lifespan - Time.currentTimeMillis(); - if (lifespan <= 0) { - return ENTRY_EXPIRED_FLAG; - } - return lifespan; + return getUserSessionLifespanMs(realm, true, userSessionEntity.isRememberMe(), userSessionEntity.getStarted()); } @@ -152,13 +164,7 @@ public class SessionTimeouts { * @return */ public static long getOfflineSessionMaxIdleMs(RealmModel realm, ClientModel client, UserSessionEntity userSessionEntity) { - long idle = SessionExpirationUtils.calculateUserSessionIdleTimestamp(true, userSessionEntity.isRememberMe(), - TimeUnit.SECONDS.toMillis(userSessionEntity.getLastSessionRefresh()), realm); - idle = idle - Time.currentTimeMillis(); - if (idle <= 0) { - return ENTRY_EXPIRED_FLAG; - } - return idle; + return getUserSessionMaxIdleMs(realm, true, userSessionEntity.isRememberMe(), userSessionEntity.getLastSessionRefresh()); } /** @@ -171,17 +177,7 @@ public class SessionTimeouts { * @return */ public static long getOfflineClientSessionLifespanMs(RealmModel realm, ClientModel client, AuthenticatedClientSessionEntity authenticatedClientSessionEntity) { - long lifespan = SessionExpirationUtils.calculateClientSessionMaxLifespanTimestamp(true, authenticatedClientSessionEntity.isUserSessionRememberMe(), - TimeUnit.SECONDS.toMillis(authenticatedClientSessionEntity.getStarted()), TimeUnit.SECONDS.toMillis(authenticatedClientSessionEntity.getUserSessionStarted()), - realm, client); - if (lifespan == -1L) { - return lifespan; - } - lifespan = lifespan - Time.currentTimeMillis(); - if (lifespan <= 0) { - return ENTRY_EXPIRED_FLAG; - } - return lifespan; + return getClientSessionLifespanMs(realm, client, true, authenticatedClientSessionEntity.isUserSessionRememberMe(), authenticatedClientSessionEntity.getStarted(), authenticatedClientSessionEntity.getUserSessionStarted()); } /** @@ -194,13 +190,7 @@ public class SessionTimeouts { * @return */ public static long getOfflineClientSessionMaxIdleMs(RealmModel realm, ClientModel client, AuthenticatedClientSessionEntity authenticatedClientSessionEntity) { - long idle = SessionExpirationUtils.calculateClientSessionIdleTimestamp(true, authenticatedClientSessionEntity.isUserSessionRememberMe(), - TimeUnit.SECONDS.toMillis(authenticatedClientSessionEntity.getTimestamp()), realm, client); - idle = idle - Time.currentTimeMillis(); - if (idle <= 0) { - return ENTRY_EXPIRED_FLAG; - } - return idle; + return getClientSessionMaxIdleMs(realm, client, true, authenticatedClientSessionEntity.isUserSessionRememberMe(), authenticatedClientSessionEntity.getTimestamp()); } diff --git a/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/storage/legacy/infinispan/CacheManagerFactory.java b/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/storage/legacy/infinispan/CacheManagerFactory.java index fc0910e55c1..c419c91f1e8 100644 --- a/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/storage/legacy/infinispan/CacheManagerFactory.java +++ b/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/storage/legacy/infinispan/CacheManagerFactory.java @@ -20,14 +20,12 @@ package org.keycloak.quarkus.runtime.storage.legacy.infinispan; import java.security.KeyManagementException; import java.security.NoSuchAlgorithmException; import java.util.Arrays; -import java.util.List; import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import java.util.stream.Stream; import io.micrometer.core.instrument.Metrics; import org.infinispan.client.hotrod.RemoteCache; @@ -67,7 +65,6 @@ import org.keycloak.marshalling.KeycloakModelSchema; import org.keycloak.marshalling.Marshalling; import org.keycloak.models.sessions.infinispan.RootAuthenticationSessionAdapter; import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity; -import org.keycloak.models.sessions.infinispan.entities.RootAuthenticationSessionEntity; import org.keycloak.quarkus.runtime.configuration.Configuration; import javax.net.ssl.SSLContext; @@ -80,7 +77,6 @@ import static org.keycloak.config.CachingOptions.CACHE_REMOTE_HOST_PROPERTY; import static org.keycloak.config.CachingOptions.CACHE_REMOTE_PASSWORD_PROPERTY; import static org.keycloak.config.CachingOptions.CACHE_REMOTE_PORT_PROPERTY; import static org.keycloak.config.CachingOptions.CACHE_REMOTE_USERNAME_PROPERTY; -import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.ACTION_TOKEN_CACHE; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.AUTHENTICATION_SESSIONS_CACHE_NAME; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.CLUSTERED_CACHE_NAMES; @@ -88,7 +84,6 @@ import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.L import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.OFFLINE_CLIENT_SESSION_CACHE_NAME; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.OFFLINE_USER_SESSION_CACHE_NAME; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.USER_SESSION_CACHE_NAME; -import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.WORK_CACHE_NAME; import static org.wildfly.security.sasl.util.SaslMechanismInformation.Names.SCRAM_SHA_512; public class CacheManagerFactory { @@ -192,17 +187,8 @@ public class CacheManagerFactory { logger.warn("Creating remote cache in external Infinispan server. It should not be used in production!"); var baseConfig = defaultRemoteCacheBuilder().build(); - Stream.of(USER_SESSION_CACHE_NAME, OFFLINE_USER_SESSION_CACHE_NAME, CLIENT_SESSION_CACHE_NAME, OFFLINE_CLIENT_SESSION_CACHE_NAME, ACTION_TOKEN_CACHE, WORK_CACHE_NAME) + Arrays.stream(CLUSTERED_CACHE_NAMES) .forEach(name -> builder.remoteCache(name).configuration(baseConfig.toStringConfiguration(name))); - - // We need indexed caches because the delete statement fails for non-indexed cache. - createIndexedRemoteCache(builder, LOGIN_FAILURE_CACHE_NAME, List.of(LoginFailureEntity.class)); - createIndexedRemoteCache(builder, AUTHENTICATION_SESSIONS_CACHE_NAME, List.of(RootAuthenticationSessionEntity.class)); - } - - private static void createIndexedRemoteCache(org.infinispan.client.hotrod.configuration.ConfigurationBuilder builder, String name, List> entities) { - var config = indexedRemoteCacheBuilder(entities).build(); - builder.remoteCache(name).configuration(config.toStringConfiguration(name)); } private static ConfigurationBuilder defaultRemoteCacheBuilder() { @@ -212,15 +198,6 @@ public class CacheManagerFactory { return builder; } - private static ConfigurationBuilder indexedRemoteCacheBuilder(List> entities) { - var builder = defaultRemoteCacheBuilder(); - var indexBuilder = builder.indexing().enable(); - entities.stream() - .map(Marshalling::protoEntity) - .forEach(indexBuilder::addIndexedEntity); - return builder; - } - private void updateProtoSchema(RemoteCacheManager remoteCacheManager) { var key = KeycloakModelSchema.INSTANCE.getProtoFileName(); var current = KeycloakModelSchema.INSTANCE.getProtoFile(); @@ -229,22 +206,22 @@ public class CacheManagerFactory { var stored = protostreamMetadataCache.getWithMetadata(key); if (stored == null) { if (protostreamMetadataCache.putIfAbsent(key, current) == null) { - logger.info("Infinispan Protostream schema uploaded for the first time."); + logger.info("Infinispan ProtoStream schema uploaded for the first time."); } else { - logger.info("Failed to update Infinispan Protostream schema. Assumed it was updated by other Keycloak server."); + logger.info("Failed to update Infinispan ProtoStream schema. Assumed it was updated by other Keycloak server."); } checkForProtoSchemaErrors(protostreamMetadataCache); return; } if (Objects.equals(stored.getValue(), current)) { - logger.info("Infinispan Protostream schema is up to date!"); + logger.info("Infinispan ProtoStream schema is up to date!"); return; } if (protostreamMetadataCache.replaceWithVersion(key, current, stored.getVersion())) { - logger.info("Infinispan Protostream schema successful updated."); + logger.info("Infinispan ProtoStream schema successful updated."); reindexCaches(remoteCacheManager, stored.getValue(), current); } else { - logger.info("Failed to update Infinispan Protostream schema. Assumed it was updated by other Keycloak server."); + logger.info("Failed to update Infinispan ProtoStream schema. Assumed it was updated by other Keycloak server."); } checkForProtoSchemaErrors(protostreamMetadataCache); } diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/infinispan/InfinispanIckleQueryTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/infinispan/InfinispanIckleQueryTest.java new file mode 100644 index 00000000000..bb8f107a976 --- /dev/null +++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/infinispan/InfinispanIckleQueryTest.java @@ -0,0 +1,538 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.keycloak.testsuite.model.infinispan; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.Executor; +import java.util.concurrent.ThreadLocalRandom; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.commons.api.query.Query; +import org.infinispan.commons.util.concurrent.CompletionStages; +import org.infinispan.util.concurrent.WithinThreadExecutor; +import org.junit.Assert; +import org.junit.Assume; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.rules.TestRule; +import org.keycloak.connections.infinispan.InfinispanConnectionProvider; +import org.keycloak.infinispan.util.InfinispanUtils; +import org.keycloak.models.KeycloakSession; +import org.keycloak.models.RealmProvider; +import org.keycloak.models.UserLoginFailureProvider; +import org.keycloak.models.UserProvider; +import org.keycloak.models.UserSessionProvider; +import org.keycloak.models.sessions.infinispan.changes.remote.remover.ConditionalRemover; +import org.keycloak.models.sessions.infinispan.changes.remote.remover.query.ByRealmIdQueryConditionalRemover; +import org.keycloak.models.sessions.infinispan.changes.remote.remover.query.ClientSessionQueryConditionalRemover; +import org.keycloak.models.sessions.infinispan.changes.remote.remover.query.UserSessionQueryConditionalRemover; +import org.keycloak.models.sessions.infinispan.entities.ClientSessionKey; +import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity; +import org.keycloak.models.sessions.infinispan.entities.LoginFailureKey; +import org.keycloak.models.sessions.infinispan.entities.RemoteAuthenticatedClientSessionEntity; +import org.keycloak.models.sessions.infinispan.entities.RemoteUserSessionEntity; +import org.keycloak.models.sessions.infinispan.query.ClientSessionQueries; +import org.keycloak.models.sessions.infinispan.query.QueryHelper; +import org.keycloak.models.sessions.infinispan.query.UserSessionQueries; +import org.keycloak.models.sessions.infinispan.remote.RemoteUserLoginFailureProviderFactory; +import org.keycloak.testsuite.model.KeycloakModelTest; +import org.keycloak.testsuite.model.RequireProvider; + +@RequireProvider(UserLoginFailureProvider.class) +@RequireProvider(UserSessionProvider.class) +@RequireProvider(UserProvider.class) +@RequireProvider(RealmProvider.class) +public class InfinispanIckleQueryTest extends KeycloakModelTest { + + private static final Executor EXECUTOR = new WithinThreadExecutor(); + private static final List REALMS = IntStream.range(0, 2).mapToObj(value -> "realm" + value).toList(); + private static final List USERS = IntStream.range(0, 2).mapToObj(value -> "user" + value).toList(); + private static final List BROKER_SESSIONS = IntStream.range(0, 2).mapToObj(value -> "brokerSession" + value).toList(); + private static final List BROKER_USERS = IntStream.range(0, 2).mapToObj(value -> "brokerUser" + value).toList(); + private static final List USER_SESSIONS = IntStream.range(0, 2).mapToObj(value -> "userSession" + value).toList(); + private static final List CLIENTS = IntStream.range(0, 2).mapToObj(value -> "client" + value).toList(); + + @ClassRule + public static final TestRule SKIPPED_PROFILES = (base, description) -> { + Assume.assumeTrue(InfinispanUtils.isRemoteInfinispan()); + return base; + }; + + @Test + public void testByRealmIdQueryConditionalRemover() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.LOGIN_FAILURE_CACHE_NAME); + + var realm0Key = new LoginFailureKey("realm0", "a"); + var realm1Key = new LoginFailureKey("realm1", "a"); + var realm2Key = new LoginFailureKey("realm2", "a"); + + Map data = new HashMap<>(); + + // create and store users + Stream.of(realm0Key, realm1Key, realm2Key).forEach(key -> data.put(key, new LoginFailureEntity(key.realmId(), key.userId()))); + cache.putAll(data); + assertCacheSize(cache, 3); + + ByRealmIdQueryConditionalRemover remover = new ByRealmIdQueryConditionalRemover<>(RemoteUserLoginFailureProviderFactory.PROTO_ENTITY); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 3); + + // remove single realm + remover.removeByRealmId("realm0"); + assertRemove(remover, realm0Key, data.get(realm0Key), true); + assertRemove(remover, realm1Key, data.get(realm1Key), false); + assertRemove(remover, realm2Key, data.get(realm2Key), false); + executeRemover(remover, cache); + assertCacheSize(cache, 2); + Assert.assertFalse(cache.containsKey(realm0Key)); + + // remove all realms + remover.removeByRealmId("realm1"); + remover.removeByRealmId("realm2"); + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testUserSessionRemoveByRealm() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME); + + var realm0Key = "a"; + var realm1Key = "b"; + var realm2Key = "c"; + + Map data = Map.of( + realm0Key, RemoteUserSessionEntity.mockEntity(realm0Key, "realm0", "user0"), + realm1Key, RemoteUserSessionEntity.mockEntity(realm1Key, "realm1", "user0"), + realm2Key, RemoteUserSessionEntity.mockEntity(realm2Key, "realm2", "user0") + ); + cache.putAll(data); + assertCacheSize(cache, 3); + + var remover = new UserSessionQueryConditionalRemover(); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 3); + + // remove single realm + remover.removeByRealmId("realm0"); + assertRemove(remover, realm0Key, data.get(realm0Key), true); + assertRemove(remover, realm1Key, data.get(realm1Key), false); + assertRemove(remover, realm2Key, data.get(realm2Key), false); + executeRemover(remover, cache); + assertCacheSize(cache, 2); + Assert.assertFalse(cache.containsKey(realm0Key)); + + // remove all realms + remover.removeByRealmId("realm1"); + remover.removeByRealmId("realm2"); + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testUserSessionRemoveByUser() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME); + + var user0Key = "a"; + var user1Key = "b"; + var user2Key = "c"; + + Map data = Map.of( + user0Key, RemoteUserSessionEntity.mockEntity(user0Key, "realm0", "user0"), + user1Key, RemoteUserSessionEntity.mockEntity(user1Key, "realm0", "user1"), + user2Key, RemoteUserSessionEntity.mockEntity(user2Key, "realm1", "user2") + ); + cache.putAll(data); + assertCacheSize(cache, 3); + + var remover = new UserSessionQueryConditionalRemover(); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 3); + + // remove single user session + remover.removeByUserId("realm0", "user1"); + assertRemove(remover, user0Key, data.get(user0Key), false); + assertRemove(remover, user1Key, data.get(user1Key), true); + assertRemove(remover, user2Key, data.get(user2Key), false); + executeRemover(remover, cache); + assertCacheSize(cache, 2); + Assert.assertFalse(cache.containsKey(user1Key)); + + // remove all user sessions + remover.removeByUserId("realm0", "user0"); + remover.removeByUserId("realm1", "user2"); + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testUserSessionRemoveMultiple() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME); + + var k0 = "a"; + var k1 = "b"; + var k2 = "c"; + var k3 = "d"; + + Map data = Map.of( + k0, RemoteUserSessionEntity.mockEntity(k0, "realm0", "user0"), + k1, RemoteUserSessionEntity.mockEntity(k1, "realm0", "user1"), + k2, RemoteUserSessionEntity.mockEntity(k2, "realm1", "user2"), + k3, RemoteUserSessionEntity.mockEntity(k3, "realm2", "user3") + ); + cache.putAll(data); + assertCacheSize(cache, 4); + + var remover = new UserSessionQueryConditionalRemover(); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 4); + + // remove all + remover.removeByRealmId("realm0"); // removes k0, k1 + remover.removeByUserId("realm1", "user2"); // removes k2 + remover.removeByUserId("realm2", "user3"); // removes k3 + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testClientSessionRemoveByRealm() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME); + + var realm0Key = new ClientSessionKey("a", "a"); + var realm1Key = new ClientSessionKey("b", "b"); + var realm2Key = new ClientSessionKey("c", "c"); + + Map data = Map.of( + realm0Key, RemoteAuthenticatedClientSessionEntity.mockEntity("a", "a", "realm0"), + realm1Key, RemoteAuthenticatedClientSessionEntity.mockEntity("a", "a", "realm1"), + realm2Key, RemoteAuthenticatedClientSessionEntity.mockEntity("a", "a", "realm2") + ); + cache.putAll(data); + assertCacheSize(cache, 3); + + var remover = new ClientSessionQueryConditionalRemover(); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 3); + + // remove single realm + remover.removeByRealmId("realm0"); + assertRemove(remover, realm0Key, data.get(realm0Key), true); + assertRemove(remover, realm1Key, data.get(realm1Key), false); + assertRemove(remover, realm2Key, data.get(realm2Key), false); + executeRemover(remover, cache); + assertCacheSize(cache, 2); + Assert.assertFalse(cache.containsKey(realm0Key)); + + // remove all realms + remover.removeByRealmId("realm1"); + remover.removeByRealmId("realm2"); + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testClientSessionRemoveByUser() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME); + + var user0Key = new ClientSessionKey("a", "a"); + var user1Key = new ClientSessionKey("b", "b"); + var user2Key = new ClientSessionKey("c", "c"); + + Map data = Map.of( + user0Key, RemoteAuthenticatedClientSessionEntity.mockEntity("a", "user0", "realm0"), + user1Key, RemoteAuthenticatedClientSessionEntity.mockEntity("a", "user1", "realm0"), + user2Key, RemoteAuthenticatedClientSessionEntity.mockEntity("a", "user2", "realm1") + ); + cache.putAll(data); + assertCacheSize(cache, 3); + + var remover = new ClientSessionQueryConditionalRemover(); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 3); + + // remove client session + remover.removeByUserId("realm0", "user1"); + assertRemove(remover, user0Key, data.get(user0Key), false); + assertRemove(remover, user1Key, data.get(user1Key), true); + assertRemove(remover, user2Key, data.get(user2Key), false); + executeRemover(remover, cache); + assertCacheSize(cache, 2); + Assert.assertFalse(cache.containsKey(user1Key)); + + // remove client sessions + remover.removeByUserId("realm0", "user0"); + remover.removeByUserId("realm1", "user2"); + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testClientSessionRemoveByUserSession() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME); + + var userSession0Key = new ClientSessionKey("a", "a"); + var userSession1Key = new ClientSessionKey("b", "b"); + var userSession2Key = new ClientSessionKey("c", "c"); + var userSession3Key = new ClientSessionKey("d", "d"); + + Map data = Map.of( + userSession0Key, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession0", "a", "a"), + userSession1Key, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession1", "a", "a"), + userSession2Key, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession1", "a", "a"), + userSession3Key, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession2", "a", "a") + ); + cache.putAll(data); + assertCacheSize(cache, 4); + + var remover = new ClientSessionQueryConditionalRemover(); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 4); + + // remove single client session + remover.removeByUserSessionId("userSession0"); + assertRemove(remover, userSession0Key, data.get(userSession0Key), true); + assertRemove(remover, userSession1Key, data.get(userSession1Key), false); + assertRemove(remover, userSession2Key, data.get(userSession2Key), false); + assertRemove(remover, userSession3Key, data.get(userSession3Key), false); + executeRemover(remover, cache); + assertCacheSize(cache, 3); + Assert.assertFalse(cache.containsKey(userSession0Key)); + + // remove all client sessions + remover.removeByUserSessionId("userSession1"); + remover.removeByUserSessionId("userSession2"); + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testClientSessionRemoveMultiple() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME); + + var key0 = new ClientSessionKey("a", "a"); + var key1 = new ClientSessionKey("b", "b"); + var key2 = new ClientSessionKey("c", "c"); + var key3 = new ClientSessionKey("d", "d"); + var key4 = new ClientSessionKey("e", "e"); + var key5 = new ClientSessionKey("f", "f"); + + Map data = Map.of( + key0, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession0", "user0", "realm0"), + key1, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession1", "user1", "realm0"), + key2, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession2", "user2", "realm1"), + key3, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession3", "user2", "realm1"), + key4, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession4", "user2", "realm2"), + key5, RemoteAuthenticatedClientSessionEntity.mockEntity("userSession4", "user2", "realm2") + ); + cache.putAll(data); + assertCacheSize(cache, 6); + + var remover = new ClientSessionQueryConditionalRemover(); + + // nothing should be removed + data.forEach((k, v) -> assertRemove(remover, k, v, false)); + executeRemover(remover, cache); + assertCacheSize(cache, 6); + + // remove all users + remover.removeByRealmId("realm0"); // key0 & key1 + remover.removeByUserId("realm1", "user2"); // key2 & key3 + remover.removeByUserSessionId("userSession4"); // key4 && key5 + data.forEach((k, v) -> assertRemove(remover, k, v, true)); + executeRemover(remover, cache); + assertCacheSize(cache, 0); + } + + @Test + public void testUserSessionQueries() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME); + + for (var realmId : REALMS) { + for (var userId : USERS) { + for (var brokerSessionId : BROKER_SESSIONS) { + for (var brokerUserId : BROKER_USERS) { + var id = String.format("%s-%s-%s-%s", realmId, userId, brokerSessionId, brokerUserId); + cache.put(id, RemoteUserSessionEntity.mockEntity(id, realmId, userId, brokerSessionId, brokerUserId)); + } + } + } + } + + var realm = random(REALMS); + var brokerSession = random(BROKER_SESSIONS); + var user = random(USERS); + var brokerUser = random(BROKER_USERS); + + var query = UserSessionQueries.searchByBrokerSessionId(cache, realm, brokerSession); + var expectedResults = expectUserSessionId(realm, USERS, List.of(brokerSession), BROKER_USERS); + assertQuery(query, objects -> ((RemoteUserSessionEntity) objects[0]).getUserSessionId(), expectedResults); + + query = UserSessionQueries.searchByUserId(cache, realm, user); + expectedResults = expectUserSessionId(realm, List.of(user), BROKER_SESSIONS, BROKER_USERS); + assertQuery(query, objects -> ((RemoteUserSessionEntity) objects[0]).getUserSessionId(), expectedResults); + + query = UserSessionQueries.searchByBrokerUserId(cache, realm, brokerUser); + expectedResults = expectUserSessionId(realm, USERS, BROKER_SESSIONS, List.of(brokerUser)); + assertQuery(query, objects -> ((RemoteUserSessionEntity) objects[0]).getUserSessionId(), expectedResults); + } + + @Test + public void testClientSessionQueries() { + RemoteCache cache = assumeAndReturnCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME); + + + for (var realmId : REALMS) { + for (var clientId : CLIENTS) { + for (var userSessionId : USER_SESSIONS) { + var id = new ClientSessionKey(userSessionId + "-" + realmId, clientId); + cache.put(id, RemoteAuthenticatedClientSessionEntity.mockEntity(userSessionId + "-" + realmId, clientId, "user", realmId)); + } + } + } + + var realm = random(REALMS); + var client = random(CLIENTS); + var userSession = random(USER_SESSIONS) + "-" + realm; + + var query = ClientSessionQueries.countClientSessions(cache, realm, client); + var expectedResults = Set.of(String.valueOf(USER_SESSIONS.size())); + assertQuery(query, objects -> String.valueOf(objects[0]), expectedResults); + var optCount = QueryHelper.fetchSingle(query, QueryHelper.SINGLE_PROJECTION_TO_LONG); + Assert.assertTrue(optCount.isPresent()); + Assert.assertEquals(USER_SESSIONS.size(), (long) optCount.get()); + + query = ClientSessionQueries.fetchUserSessionIdForClientId(cache, realm, client); + expectedResults = USER_SESSIONS.stream().map(s -> s + "-" + realm).collect(Collectors.toSet()); + assertQuery(query, objects -> String.valueOf(objects[0]), expectedResults); + + query = ClientSessionQueries.fetchClientSessions(cache, userSession); + expectedResults = CLIENTS.stream().map(s -> new ClientSessionKey(userSession, s)).map(Objects::toString).collect(Collectors.toSet()); + assertQuery(query, objects -> ((RemoteAuthenticatedClientSessionEntity) objects[0]).createCacheKey().toString(), expectedResults); + + // each client has user-session * realms number of active client sessions + query = ClientSessionQueries.activeClientCount(cache); + expectedResults = CLIENTS.stream().map(s -> String.format("%s-%s", s, USER_SESSIONS.size() * REALMS.size())).collect(Collectors.toSet()); + assertQuery(query, objects -> String.format("%s-%s", objects[0], objects[1]), expectedResults); + } + + private static void assertQuery(Query query, Function resultMapping, Set expectedResults) { + var results = new HashSet(); + + // test streaming with batchSize = 1 + QueryHelper.streamAll(query, 1, resultMapping).forEach(results::add); + Assert.assertEquals(expectedResults, results); + results.clear(); + + // test streaming with batchSize = results.size + QueryHelper.streamAll(query, expectedResults.size(), resultMapping).forEach(results::add); + Assert.assertEquals(expectedResults, results); + results.clear(); + + // test streaming with batchSize > results.size + QueryHelper.streamAll(query, expectedResults.size() * 2, resultMapping).forEach(results::add); + Assert.assertEquals(expectedResults, results); + results.clear(); + + query.startOffset(0).maxResults(Integer.MAX_VALUE); + Assert.assertEquals(expectedResults, new HashSet<>(QueryHelper.toCollection(query, resultMapping))); + } + + + private static String random(List elements) { + return elements.get(ThreadLocalRandom.current().nextInt(elements.size())); + } + + private static Set expectUserSessionId(String realmId, List users, List brokerSessions, List brokerUsers) { + var results = new HashSet(); + for (var userId : users) { + for (var brokerSessionId : brokerSessions) { + for (var brokerUserId : brokerUsers) { + results.add(String.format("%s-%s-%s-%s", realmId, userId, brokerSessionId, brokerUserId)); + } + } + } + return results; + } + + private RemoteCache assumeAndReturnCache(String cacheName) { + var cache = getInfinispanConnectionProvider().getRemoteCache(cacheName); + cache.clear(); + return cache; + } + + private static void executeRemover(ConditionalRemover remover, RemoteCache cache) { + var stage = CompletionStages.aggregateCompletionStage(); + remover.executeRemovals(cache, stage); + CompletionStages.join(stage.freeze()); + } + + private static void assertRemove(ConditionalRemover remover, K key, V value, boolean willRemove) { + Assert.assertEquals(willRemove, remover.willRemove(key, value)); + } + + private static void assertCacheSize(RemoteCache cache, int expectedSize) { + Assert.assertEquals(expectedSize, cache.size()); + } + + private InfinispanConnectionProvider getInfinispanConnectionProvider() { + return inComittedTransaction(InfinispanIckleQueryTest::getInfinispanConnectionProviderWithSession); + } + + private static InfinispanConnectionProvider getInfinispanConnectionProviderWithSession(KeycloakSession session) { + return session.getProvider(InfinispanConnectionProvider.class); + } + +}