diff --git a/docs/guides/high-availability/bblocks-multi-site.adoc b/docs/guides/high-availability/bblocks-multi-site.adoc index 5e5c0c8ce5..c8c978adf4 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 34a448959b..d88e24f720 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 36a8b0368f..4870cd2012 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 6607845fb7..450a657cc5 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 da09d3ff05..72138ef436 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 0000000000..2a953e79a4 --- /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 0000000000..43b5c7cd3d --- /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 271ec69799..88fc677e1f 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 0000000000..29d737dccc --- /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 d221e8ac7b..9ef7241586 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 78457c34b3..671beb2e43 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 f0094dc55d..ccbe4855be 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 ef33755be6..0000000000 --- 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 e217b339c0..0000000000 --- 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 ab7611905f..db3bccc5a2 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 0000000000..3203d7b40d --- /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 0000000000..a801eb5c06 --- /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 0000000000..93317e86a9 --- /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 0000000000..ceff06784d --- /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 0000000000..6289369fa1 --- /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 0000000000..2b64b5fcec --- /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 6f4546c7c4..6ecb496514 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 00f1ead81b..e20da5e885 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 9c7aa1e9f2..07c1953450 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 d29412ff1e..0458dbf517 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 18ee33ef7b..44398a6db3 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 07f27cbd67..0000000000 --- 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 1df7151e44..631b24db37 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 df91b52130..89e6c1f5ca 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 1539e32feb..e12c4f4fe7 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 d661073a3d..5212e9e378 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 fc0910e55c..c419c91f1e 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 0000000000..bb8f107a97 --- /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); + } + +}