Hot Rod map storage: Authentication session no-downtime store

model/build-processor/src/main/java/org/keycloak/models/map/processor/GenerateHotRodEntityImplementationsProcessor.java

@@ -103,7 +103,8 @@ public class GenerateHotRodEntityImplementationsProcessor extends AbstractGenera
             abstractHotRodEntity = elements.getTypeElement("org.keycloak.models.map.storage.hotRod.common.AbstractHotRodEntity");
             hotRodUtils = elements.getTypeElement("org.keycloak.models.map.storage.hotRod.common.HotRodTypesUtils");
 
-            boolean hasDeepClone = allMembers.stream().filter(el -> el.getKind() == ElementKind.METHOD).anyMatch(el -> "deepClone".equals(el.getSimpleName().toString()));
+            boolean hasDeepClone = allMembers.stream()
+                    .filter(el -> el.getKind() == ElementKind.METHOD && !el.getModifiers().contains(Modifier.ABSTRACT)).anyMatch(el -> "deepClone".equals(el.getSimpleName().toString()));
             boolean needsDeepClone = fieldGetters(methodsPerAttribute)
                     .map(ExecutableElement::getReturnType)
                     .anyMatch(fieldType -> ! isKnownCollectionOfImmutableFinalTypes(fieldType) && ! isImmutableFinalType(fieldType));

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/HotRodMapStorageProviderFactory.java

@@ -28,9 +28,14 @@ import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.map.authSession.MapAuthenticationSessionEntity;
+import org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntity;
 import org.keycloak.models.map.clientscope.MapClientScopeEntity;
 import org.keycloak.models.map.group.MapGroupEntity;
 import org.keycloak.models.map.role.MapRoleEntity;
+import org.keycloak.models.map.storage.hotRod.authSession.HotRodAuthenticationSessionEntityDelegate;
+import org.keycloak.models.map.storage.hotRod.authSession.HotRodRootAuthenticationSessionEntity;
+import org.keycloak.models.map.storage.hotRod.authSession.HotRodRootAuthenticationSessionEntityDelegate;
 import org.keycloak.models.map.storage.hotRod.role.HotRodRoleEntity;
 import org.keycloak.models.map.storage.hotRod.role.HotRodRoleEntityDelegate;
 import org.keycloak.models.map.storage.hotRod.client.HotRodClientEntity;
@@ -57,6 +62,7 @@ import org.keycloak.models.map.user.MapUserCredentialEntity;
 import org.keycloak.models.map.user.MapUserEntity;
 import org.keycloak.models.map.user.MapUserFederatedIdentityEntity;
 import org.keycloak.provider.EnvironmentDependentProviderFactory;
+import org.keycloak.sessions.RootAuthenticationSessionModel;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -67,6 +73,8 @@ public class HotRodMapStorageProviderFactory implements AmphibianProviderFactory
     private static final Logger LOG = Logger.getLogger(HotRodMapStorageProviderFactory.class);
 
     private final static DeepCloner CLONER = new DeepCloner.Builder()
+            .constructor(MapRootAuthenticationSessionEntity.class,  HotRodRootAuthenticationSessionEntityDelegate::new)
+            .constructor(MapAuthenticationSessionEntity.class,      HotRodAuthenticationSessionEntityDelegate::new)
             .constructor(MapClientEntity.class,                     HotRodClientEntityDelegate::new)
             .constructor(MapProtocolMapperEntity.class,             HotRodProtocolMapperEntityDelegate::new)
             .constructor(MapClientScopeEntity.class,                HotRodClientScopeEntityDelegate::new)
@@ -80,6 +88,12 @@ public class HotRodMapStorageProviderFactory implements AmphibianProviderFactory
 
     public static final Map<Class<?>, HotRodEntityDescriptor<?, ?>> ENTITY_DESCRIPTOR_MAP = new HashMap<>();
     static {
+        // Authentication sessions descriptor
+        ENTITY_DESCRIPTOR_MAP.put(RootAuthenticationSessionModel.class,
+                new HotRodEntityDescriptor<>(RootAuthenticationSessionModel.class,
+                        HotRodRootAuthenticationSessionEntity.class,
+                        HotRodRootAuthenticationSessionEntityDelegate::new));
+
         // Clients descriptor
         ENTITY_DESCRIPTOR_MAP.put(ClientModel.class,
                 new HotRodEntityDescriptor<>(ClientModel.class,

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/authSession/HotRodAuthenticationSessionEntity.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright 2022 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.map.storage.hotRod.authSession;
+
+import org.infinispan.protostream.annotations.ProtoField;
+import org.keycloak.models.map.annotations.GenerateHotRodEntityImplementation;
+import org.keycloak.models.map.authSession.MapAuthenticationSessionEntity;
+import org.keycloak.models.map.storage.hotRod.common.AbstractHotRodEntity;
+import org.keycloak.models.map.storage.hotRod.common.HotRodPair;
+import org.keycloak.models.map.storage.hotRod.common.HotRodTypesUtils;
+import org.keycloak.models.map.storage.hotRod.common.UpdatableHotRodEntityDelegateImpl;
+import org.keycloak.sessions.AuthenticationSessionModel;
+
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+@GenerateHotRodEntityImplementation(
+        implementInterface = "org.keycloak.models.map.authSession.MapAuthenticationSessionEntity",
+        inherits = "org.keycloak.models.map.storage.hotRod.authSession.HotRodAuthenticationSessionEntity.AbstractHotRodAuthenticationSessionEntityDelegate"
+)
+public class HotRodAuthenticationSessionEntity extends AbstractHotRodEntity {
+
+    @ProtoField(number = 1)
+    public String tabId;
+
+    @ProtoField(number = 2)
+    public String clientUUID;
+
+    @ProtoField(number = 3)
+    public String authUserId;
+
+    @ProtoField(number = 4)
+    public Integer timestamp;
+
+    @ProtoField(number = 5)
+    public String redirectUri;
+
+    @ProtoField(number = 6)
+    public String action;
+
+    @ProtoField(number = 7)
+    public Set<String> clientScopes;
+
+    @ProtoField(number = 8)
+    public Set<HotRodPair<String, HotRodExecutionStatus>> executionStatuses;
+
+    @ProtoField(number = 9)
+    public String protocol;
+
+    @ProtoField(number = 10)
+    public Set<HotRodPair<String, String>> clientNotes;
+
+    @ProtoField(number = 11)
+    public Set<HotRodPair<String, String>> authNotes;
+
+    @ProtoField(number = 12)
+    public Set<String> requiredActions;
+
+    @ProtoField(number = 13)
+    public Set<HotRodPair<String, String>> userSessionNotes;
+
+    public static abstract class AbstractHotRodAuthenticationSessionEntityDelegate extends UpdatableHotRodEntityDelegateImpl<HotRodAuthenticationSessionEntity> implements MapAuthenticationSessionEntity {
+
+        @Override
+        public Map<String, AuthenticationSessionModel.ExecutionStatus> getExecutionStatuses() {
+            Set<HotRodPair<String, HotRodExecutionStatus>> executionStatuses = getHotRodEntity().executionStatuses;
+            if (executionStatuses == null) {
+                return Collections.emptyMap();
+            }
+            return executionStatuses.stream().collect(Collectors.toMap(HotRodPair::getKey,
+                    v -> AuthenticationSessionModel.ExecutionStatus.valueOf(v.getValue().name())));
+        }
+
+        @Override
+        public void setExecutionStatuses(Map<String, AuthenticationSessionModel.ExecutionStatus> executionStatus) {
+            HotRodAuthenticationSessionEntity hotRodEntity = getHotRodEntity();
+            Set<HotRodPair<String, HotRodExecutionStatus>> executionStatusSet = executionStatus == null ? null :
+                    executionStatus.entrySet().stream()
+                            .map(e -> new HotRodPair<>(e.getKey(), HotRodExecutionStatus.valueOf(e.getValue().name())))
+                            .collect(Collectors.toSet());
+            hotRodEntity.updated |= ! Objects.equals(hotRodEntity.executionStatuses, executionStatusSet);
+            hotRodEntity.executionStatuses = executionStatusSet;
+        }
+
+        @Override
+        public void setExecutionStatus(String authenticator, AuthenticationSessionModel.ExecutionStatus status) {
+            HotRodAuthenticationSessionEntity hotRodEntity = getHotRodEntity();
+            if (hotRodEntity.executionStatuses == null) {
+                hotRodEntity.executionStatuses = new HashSet<>();
+            }
+            boolean valueUndefined = status == null;
+            hotRodEntity.updated |= HotRodTypesUtils.removeFromSetByMapKey(hotRodEntity.executionStatuses, authenticator, HotRodTypesUtils::getKey);
+            hotRodEntity.updated |= !valueUndefined && hotRodEntity.executionStatuses.add(new HotRodPair<>(authenticator, HotRodExecutionStatus.valueOf(status.name())));
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return HotRodAuthenticationSessionEntityDelegate.entityEquals(this, o);
+    }
+
+    @Override
+    public int hashCode() {
+        return HotRodAuthenticationSessionEntityDelegate.entityHashCode(this);
+    }
+}

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/authSession/HotRodExecutionStatus.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2022 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.map.storage.hotRod.authSession;
+
+import org.infinispan.protostream.annotations.ProtoEnumValue;
+
+public enum HotRodExecutionStatus {
+
+    @ProtoEnumValue(number = 0)
+    FAILED,
+
+    @ProtoEnumValue(number = 1)
+    SUCCESS,
+
+    @ProtoEnumValue(number = 2)
+    SETUP_REQUIRED,
+
+    @ProtoEnumValue(number = 3)
+    ATTEMPTED,
+
+    @ProtoEnumValue(number = 4)
+    SKIPPED,
+
+    @ProtoEnumValue(number = 5)
+    CHALLENGED,
+
+    @ProtoEnumValue(number = 6)
+    EVALUATED_TRUE,
+
+    @ProtoEnumValue(number = 7)
+    EVALUATED_FALSE
+}

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/authSession/HotRodRootAuthenticationSessionEntity.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright 2022 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.map.storage.hotRod.authSession;
+
+import org.infinispan.protostream.annotations.ProtoDoc;
+import org.infinispan.protostream.annotations.ProtoField;
+import org.keycloak.models.map.annotations.GenerateHotRodEntityImplementation;
+import org.keycloak.models.map.authSession.MapAuthenticationSessionEntity;
+import org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntity;
+import org.keycloak.models.map.common.UpdatableEntity;
+import org.keycloak.models.map.storage.hotRod.common.AbstractHotRodEntity;
+import org.keycloak.models.map.storage.hotRod.common.UpdatableHotRodEntityDelegateImpl;
+
+import java.util.Collections;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+
+@GenerateHotRodEntityImplementation(
+        implementInterface = "org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntity",
+        inherits = "org.keycloak.models.map.storage.hotRod.authSession.HotRodRootAuthenticationSessionEntity.AbstractHotRodRootAuthenticationSessionEntityDelegate"
+)
+@ProtoDoc("@Indexed")
+public class HotRodRootAuthenticationSessionEntity extends AbstractHotRodEntity {
+
+    @ProtoField(number = 1, required = true)
+    public int entityVersion = 1;
+
+    @ProtoField(number = 2, required = true)
+    public String id;
+
+    @ProtoDoc("@Field(index = Index.YES, store = Store.YES)")
+    @ProtoField(number = 3)
+    public String realmId;
+
+    @ProtoDoc("@Field(index = Index.YES, store = Store.YES)")
+    @ProtoField(number = 4)
+    public Integer timestamp;
+
+    @ProtoField(number = 5)
+    public Set<HotRodAuthenticationSessionEntity> authenticationSessions;
+
+    public static abstract class AbstractHotRodRootAuthenticationSessionEntityDelegate extends UpdatableHotRodEntityDelegateImpl<HotRodRootAuthenticationSessionEntity> implements MapRootAuthenticationSessionEntity {
+
+        @Override
+        public String getId() {
+            return getHotRodEntity().id;
+        }
+
+        @Override
+        public void setId(String id) {
+            HotRodRootAuthenticationSessionEntity entity = getHotRodEntity();
+            if (entity.id != null) throw new IllegalStateException("Id cannot be changed");
+            entity.id = id;
+            entity.updated |= id != null;
+        }
+
+        @Override
+        public Optional<MapAuthenticationSessionEntity> getAuthenticationSession(String tabId) {
+            HotRodRootAuthenticationSessionEntity rootAuthSession = getHotRodEntity();
+            if (rootAuthSession.authenticationSessions == null || rootAuthSession.authenticationSessions.isEmpty()) return Optional.empty();
+
+            return rootAuthSession.authenticationSessions.stream()
+                    .filter(as -> Objects.equals(as.tabId, tabId))
+                    .findFirst()
+                    .map(HotRodAuthenticationSessionEntityDelegate::new);
+        }
+
+        @Override
+        public Boolean removeAuthenticationSession(String tabId) {
+            HotRodRootAuthenticationSessionEntity rootAuthSession = getHotRodEntity();
+            boolean removed = rootAuthSession.authenticationSessions != null &&
+                    rootAuthSession.authenticationSessions.removeIf(c -> Objects.equals(c.tabId, tabId));
+            rootAuthSession.updated |= removed;
+            return removed;
+        }
+
+        @Override
+        public boolean isUpdated() {
+            HotRodRootAuthenticationSessionEntity rootAuthSession = getHotRodEntity();
+            return rootAuthSession.updated ||
+                    Optional.ofNullable(getAuthenticationSessions()).orElseGet(Collections::emptySet).stream().anyMatch(MapAuthenticationSessionEntity::isUpdated);
+        }
+
+        @Override
+        public void clearUpdatedFlag() {
+            HotRodRootAuthenticationSessionEntity rootAuthSession = getHotRodEntity();
+            rootAuthSession.updated = false;
+            Optional.ofNullable(getAuthenticationSessions()).orElseGet(Collections::emptySet).forEach(UpdatableEntity::clearUpdatedFlag);
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return HotRodRootAuthenticationSessionEntityDelegate.entityEquals(this, o);
+    }
+
+    @Override
+    public int hashCode() {
+        return HotRodRootAuthenticationSessionEntityDelegate.entityHashCode(this);
+    }
+}

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/common/HotRodTypesUtils.java

@@ -18,6 +18,7 @@
 package org.keycloak.models.map.storage.hotRod.common;
 
 import org.keycloak.models.map.common.AbstractEntity;
+import org.keycloak.models.map.storage.hotRod.authSession.HotRodAuthenticationSessionEntity;
 import org.keycloak.models.map.storage.hotRod.user.HotRodUserConsentEntity;
 import org.keycloak.models.map.storage.hotRod.user.HotRodUserFederatedIdentityEntity;
 
@@ -110,4 +111,8 @@ public class HotRodTypesUtils {
     public static <T, V> Set<V> migrateSet(Set<T> p0, Function<T, V> migrator) {
         return p0 == null ? null : p0.stream().map(migrator).collect(Collectors.toSet());
     }
+
+    public static String getKey(HotRodAuthenticationSessionEntity hotRodAuthenticationSessionEntity) {
+        return hotRodAuthenticationSessionEntity.tabId;
+    }
 }

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/common/ProtoSchemaInitializer.java

@@ -19,6 +19,9 @@ package org.keycloak.models.map.storage.hotRod.common;
 
 import org.infinispan.protostream.GeneratedSchema;
 import org.infinispan.protostream.annotations.AutoProtoSchemaBuilder;
+import org.keycloak.models.map.storage.hotRod.authSession.HotRodAuthenticationSessionEntity;
+import org.keycloak.models.map.storage.hotRod.authSession.HotRodExecutionStatus;
+import org.keycloak.models.map.storage.hotRod.authSession.HotRodRootAuthenticationSessionEntity;
 import org.keycloak.models.map.storage.hotRod.client.HotRodClientEntity;
 import org.keycloak.models.map.storage.hotRod.client.HotRodProtocolMapperEntity;
 import org.keycloak.models.map.storage.hotRod.clientscope.HotRodClientScopeEntity;
@@ -34,6 +37,11 @@ import org.keycloak.models.map.storage.hotRod.user.HotRodUserFederatedIdentityEn
  */
 @AutoProtoSchemaBuilder(
         includeClasses = {
+                // Authentication sessions
+                HotRodRootAuthenticationSessionEntity.class,
+                HotRodAuthenticationSessionEntity.class,
+                HotRodExecutionStatus.class,
+
                 // Clients
                 HotRodClientEntity.class,
                 HotRodProtocolMapperEntity.class,

model/map-hot-rod/src/main/resources/config/cacheConfig.xml

@@ -3,6 +3,14 @@
 <infinispan>
     <cache-container>
        <!-- Specify all remote caches that should be created on the Infinispan server. -->
+        <distributed-cache name="auth-sessions" mode="SYNC">
+            <indexing>
+                <indexed-entities>
+                    <indexed-entity>kc.HotRodRootAuthenticationSessionEntity</indexed-entity>
+                </indexed-entities>
+            </indexing>
+            <encoding media-type="application/x-protostream"/>
+        </distributed-cache>
         <distributed-cache name="clients" mode="SYNC">
             <indexing>
                 <indexed-entities>

model/map-hot-rod/src/main/resources/config/infinispan.xml

@@ -5,6 +5,14 @@
         <transport stack="udp"/>
 
         <!-- Specify all remote caches that should be created on the embedded Infinispan server. -->
+        <distributed-cache name="auth-sessions" mode="SYNC">
+            <indexing>
+                <indexed-entities>
+                    <indexed-entity>kc.HotRodRootAuthenticationSessionEntity</indexed-entity>
+                </indexed-entities>
+            </indexing>
+            <encoding media-type="application/x-protostream"/>
+        </distributed-cache>
         <distributed-cache name="clients" mode="SYNC">
             <indexing>
                 <indexed-entities>

testsuite/integration-arquillian/tests/base/src/test/resources/META-INF/keycloak-server.json

@@ -266,7 +266,7 @@
             "username": "${keycloak.connectionsHotRod.username:myuser}",
             "password": "${keycloak.connectionsHotRod.password:qwer1234!}",
             "enableSecurity": "${keycloak.connectionsHotRod.enableSecurity:true}",
-            "reindexCaches": "${keycloak.connectionsHotRod.reindexCaches:clients,client-scopes,groups,users,roles}"
+            "reindexCaches": "${keycloak.connectionsHotRod.reindexCaches:auth-sessions,clients,client-scopes,groups,users,roles}"
         }
     },
 

testsuite/integration-arquillian/tests/pom.xml

@@ -1496,6 +1496,7 @@
                             <artifactId>maven-surefire-plugin</artifactId>
                             <configuration>
                                 <systemPropertyVariables>
+                                    <keycloak.authSession.map.storage.provider>hotrod</keycloak.authSession.map.storage.provider>
                                     <keycloak.client.map.storage.provider>hotrod</keycloak.client.map.storage.provider>
                                     <keycloak.clientScope.map.storage.provider>hotrod</keycloak.clientScope.map.storage.provider>
                                     <keycloak.group.map.storage.provider>hotrod</keycloak.group.map.storage.provider>

testsuite/model/src/main/resources/hotrod/infinispan.xml

@@ -1,6 +1,14 @@
 <infinispan>
     <cache-container>
         <transport stack="udp"/>
+        <distributed-cache name="auth-sessions" mode="SYNC">
+            <indexing>
+                <indexed-entities>
+                    <indexed-entity>kc.HotRodRootAuthenticationSessionEntity</indexed-entity>
+                </indexed-entities>
+            </indexing>
+            <encoding media-type="application/x-protostream"/>
+        </distributed-cache>
         <distributed-cache name="clients" mode="SYNC">
             <indexing>
                 <indexed-entities>

testsuite/model/src/test/java/org/keycloak/testsuite/model/parameters/HotRodMapStorage.java

@@ -72,7 +72,7 @@ public class HotRodMapStorage extends KeycloakModelParameters {
 
     @Override
     public void updateConfig(Config cf) {
-        cf.spi(AuthenticationSessionSpi.PROVIDER_ID).provider(MapRootAuthenticationSessionProviderFactory.PROVIDER_ID).config(STORAGE_CONFIG, ConcurrentHashMapStorageProviderFactory.PROVIDER_ID)
+        cf.spi(AuthenticationSessionSpi.PROVIDER_ID).provider(MapRootAuthenticationSessionProviderFactory.PROVIDER_ID).config(STORAGE_CONFIG, HotRodMapStorageProviderFactory.PROVIDER_ID)
           .spi("client").provider(MapClientProviderFactory.PROVIDER_ID).config(STORAGE_CONFIG, HotRodMapStorageProviderFactory.PROVIDER_ID)
           .spi("clientScope").provider(MapClientScopeProviderFactory.PROVIDER_ID).config(STORAGE_CONFIG, HotRodMapStorageProviderFactory.PROVIDER_ID)
           .spi("group").provider(MapGroupProviderFactory.PROVIDER_ID).config(STORAGE_CONFIG, HotRodMapStorageProviderFactory.PROVIDER_ID)

testsuite/model/src/test/java/org/keycloak/testsuite/model/session/AuthenticationSessionTest.java

@@ -17,6 +17,7 @@
 
 package org.keycloak.testsuite.model.session;
 
+import org.hamcrest.Matchers;
 import org.junit.Assert;
 import org.junit.Test;
 import org.keycloak.common.util.Time;
@@ -32,15 +33,17 @@ import org.keycloak.testsuite.model.KeycloakModelTest;
 import org.keycloak.testsuite.model.RequireProvider;
 
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.keycloak.testsuite.model.session.UserSessionPersisterProviderTest.createClients;
 
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
  */
-@RequireProvider(value = AuthenticationSessionProvider.class, only = InfinispanAuthenticationSessionProviderFactory.PROVIDER_ID)
+@RequireProvider(value = AuthenticationSessionProvider.class)
 public class AuthenticationSessionTest extends KeycloakModelTest {
 
     private String realmId;
@@ -49,6 +52,7 @@ public class AuthenticationSessionTest extends KeycloakModelTest {
     public void createEnvironment(KeycloakSession s) {
         RealmModel realm = s.realms().createRealm("test");
         realm.setDefaultRole(s.roles().addRealmRole(realm, Constants.DEFAULT_ROLES_ROLE_PREFIX + "-" + realm.getName()));
+        realm.setAccessCodeLifespanLogin(1800);
 
         this.realmId = realm.getId();
 
@@ -61,10 +65,12 @@ public class AuthenticationSessionTest extends KeycloakModelTest {
     }
 
     @Test
+    @RequireProvider(value = AuthenticationSessionProvider.class, only = InfinispanAuthenticationSessionProviderFactory.PROVIDER_ID)
     public void testLimitAuthSessions() {
-        RootAuthenticationSessionModel ras = withRealm(realmId, (session, realm) -> session.authenticationSessions().createRootAuthenticationSession(realm));
+        AtomicReference<String> rootAuthSessionId = new AtomicReference<>();
-
         List<String> tabIds = withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel ras = session.authenticationSessions().createRootAuthenticationSession(realm);
+            rootAuthSessionId.set(ras.getId());
             ClientModel client = realm.getClientByClientId("test-app");
             return IntStream.range(0, 300)
                     .mapToObj(i -> {
@@ -76,6 +82,7 @@ public class AuthenticationSessionTest extends KeycloakModelTest {
         });
 
         withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel ras = session.authenticationSessions().getRootAuthenticationSession(realm, rootAuthSessionId.get());
             ClientModel client = realm.getClientByClientId("test-app");
 
             // create 301st auth session
@@ -88,4 +95,101 @@ public class AuthenticationSessionTest extends KeycloakModelTest {
             return null;
         });
     }
+
+    @Test
+    public void testAuthSessions() {
+        AtomicReference<String> rootAuthSessionId = new AtomicReference<>();
+        List<String> tabIds = withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().createRootAuthenticationSession(realm);
+            rootAuthSessionId.set(rootAuthSession.getId());
+
+            ClientModel client = realm.getClientByClientId("test-app");
+            return IntStream.range(0, 5)
+                    .mapToObj(i -> {
+                        AuthenticationSessionModel authSession = rootAuthSession.createAuthenticationSession(client);
+                        authSession.setExecutionStatus("username", AuthenticationSessionModel.ExecutionStatus.ATTEMPTED);
+                        authSession.setAuthNote("foo", "bar");
+                        authSession.setClientNote("foo", "bar");
+                        return authSession;
+                    })
+                    .map(AuthenticationSessionModel::getTabId)
+                    .collect(Collectors.toList());
+        });
+
+        withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().getRootAuthenticationSession(realm, rootAuthSessionId.get());
+            Assert.assertNotNull(rootAuthSession);
+            Assert.assertEquals(rootAuthSessionId.get(), rootAuthSession.getId());
+
+            ClientModel client = realm.getClientByClientId("test-app");
+            tabIds.forEach(tabId -> {
+                AuthenticationSessionModel authSession = rootAuthSession.getAuthenticationSession(client, tabId);
+                Assert.assertNotNull(authSession);
+
+                Assert.assertEquals(AuthenticationSessionModel.ExecutionStatus.ATTEMPTED, authSession.getExecutionStatus().get("username"));
+                Assert.assertEquals("bar", authSession.getAuthNote("foo"));
+                Assert.assertEquals("bar", authSession.getClientNote("foo"));
+            });
+
+            // remove first two auth sessions
+            rootAuthSession.removeAuthenticationSessionByTabId(tabIds.get(0));
+            rootAuthSession.removeAuthenticationSessionByTabId(tabIds.get(1));
+
+            return null;
+        });
+
+        withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().getRootAuthenticationSession(realm, rootAuthSessionId.get());
+            Assert.assertNotNull(rootAuthSession);
+            Assert.assertEquals(rootAuthSessionId.get(), rootAuthSession.getId());
+
+            assertThat(rootAuthSession.getAuthenticationSessions(), Matchers.aMapWithSize(3));
+
+            Assert.assertNull(rootAuthSession.getAuthenticationSessions().get(tabIds.get(0)));
+            Assert.assertNull(rootAuthSession.getAuthenticationSessions().get(tabIds.get(1)));
+            IntStream.range(2,4).mapToObj(i -> rootAuthSession.getAuthenticationSessions().get(tabIds.get(i))).forEach(Assert::assertNotNull);
+
+            session.authenticationSessions().removeRootAuthenticationSession(realm, rootAuthSession);
+
+            return null;
+        });
+
+        withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().getRootAuthenticationSession(realm, rootAuthSessionId.get());
+            Assert.assertNull(rootAuthSession);
+
+            return null;
+        });
+    }
+
+    @Test
+    public void testRemoveExpiredAuthSessions() {
+        AtomicReference<String> rootAuthSessionId = new AtomicReference<>();
+        withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().createRootAuthenticationSession(realm);
+            ClientModel client = realm.getClientByClientId("test-app");
+            rootAuthSession.createAuthenticationSession(client);
+            rootAuthSessionId.set(rootAuthSession.getId());
+
+            return null;
+        });
+
+        withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().getRootAuthenticationSession(realm, rootAuthSessionId.get());
+            Assert.assertNotNull(rootAuthSession);
+
+            Time.setOffset(1900);
+            // not needed with Infinispan where expiration handles Infinispan itself
+            session.authenticationSessions().removeExpired(realm);
+
+            return null;
+        });
+
+        withRealm(realmId, (session, realm) -> {
+            RootAuthenticationSessionModel rootAuthSession = session.authenticationSessions().getRootAuthenticationSession(realm, rootAuthSessionId.get());
+            Assert.assertNull(rootAuthSession);
+
+            return null;
+        });
+    }
 }