From 1a98765fb7e2c5acf7ca6b85ebd79d86fb732f28 Mon Sep 17 00:00:00 2001
From: Michal Hajas <mhajas@redhat.com>
Date: Tue, 24 May 2022 09:43:03 +0200
Subject: [PATCH] Fix cascade removal of client session on user session removal
 for CHM

Closes #12146
---
 .../storage/chm/ConcurrentHashMapStorage.java |  7 ++++++-
 .../UserSessionConcurrentHashMapStorage.java  |  7 ++++++-
 .../session/UserSessionProviderModelTest.java | 19 +++++++++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorage.java b/model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorage.java
index a3d64ef6f9c..d7854badb0f 100644
--- a/model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorage.java
+++ b/model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorage.java
@@ -129,7 +129,12 @@ public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEnt
     @SuppressWarnings("unchecked")
     public MapKeycloakTransaction<V, M> createTransaction(KeycloakSession session) {
         MapKeycloakTransaction<V, M> sessionTransaction = session.getAttribute("map-transaction-" + hashCode(), MapKeycloakTransaction.class);
-        return sessionTransaction == null ? new ConcurrentHashMapKeycloakTransaction<>(this, keyConverter, cloner, fieldPredicates) : sessionTransaction;
+
+        if (sessionTransaction == null) {
+            sessionTransaction = new ConcurrentHashMapKeycloakTransaction<>(this, keyConverter, cloner, fieldPredicates);
+            session.setAttribute("map-transaction-" + hashCode(), sessionTransaction);
+        }
+        return sessionTransaction;
     }
 
     public MapModelCriteriaBuilder<K, V, M> createCriteriaBuilder() {
diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/chm/UserSessionConcurrentHashMapStorage.java b/model/map/src/main/java/org/keycloak/models/map/storage/chm/UserSessionConcurrentHashMapStorage.java
index 9dd8ef95673..cd198deb7a9 100644
--- a/model/map/src/main/java/org/keycloak/models/map/storage/chm/UserSessionConcurrentHashMapStorage.java
+++ b/model/map/src/main/java/org/keycloak/models/map/storage/chm/UserSessionConcurrentHashMapStorage.java
@@ -92,6 +92,11 @@ public class UserSessionConcurrentHashMapStorage<K> extends ConcurrentHashMapSto
     @SuppressWarnings("unchecked")
     public MapKeycloakTransaction<MapUserSessionEntity, UserSessionModel> createTransaction(KeycloakSession session) {
         MapKeycloakTransaction<MapUserSessionEntity, UserSessionModel> sessionTransaction = session.getAttribute("map-transaction-" + hashCode(), MapKeycloakTransaction.class);
-        return sessionTransaction == null ? new Transaction(clientSessionStore.createTransaction(session), clientSessionStore.getKeyConverter(), cloner, fieldPredicates) : sessionTransaction;
+
+        if (sessionTransaction == null) {
+            sessionTransaction = new Transaction(clientSessionStore.createTransaction(session), clientSessionStore.getKeyConverter(), cloner, fieldPredicates);
+            session.setAttribute("map-transaction-" + hashCode(), sessionTransaction);
+        }
+        return sessionTransaction;
     }
 }
diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java
index 680228ccb71..dc2db81c713 100644
--- a/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java
+++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java
@@ -20,6 +20,7 @@ import org.junit.Assert;
 import org.junit.Test;
 import org.keycloak.common.util.Time;
 import org.keycloak.models.AuthenticatedClientSessionModel;
+import org.keycloak.models.ClientModel;
 import org.keycloak.models.Constants;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
@@ -42,6 +43,8 @@ import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.keycloak.testsuite.model.session.UserSessionPersisterProviderTest.createClients;
 import static org.keycloak.testsuite.model.session.UserSessionPersisterProviderTest.createSessions;
 
@@ -202,4 +205,20 @@ public class UserSessionProviderModelTest extends KeycloakModelTest {
             }
         }
     }
+
+    @Test
+    public void testCascadeRemovalOfClientSessionOnUserSessionRemoval() {
+        UserSessionModel[] origSessions = inComittedTransaction(session -> { return createSessions(session, realmId); });
+
+        String testAppClientSessionId = withRealm(realmId, (session, realm) -> {
+            ClientModel testApp = realm.getClientByClientId("test-app");
+            UserSessionModel userSessionToBeRemoved = session.sessions().getUserSession(realm, origSessions[0].getId());
+            String returnValue = userSessionToBeRemoved.getAuthenticatedClientSessions().get(testApp.getId()).getId();
+
+            session.sessions().removeUserSession(realm, userSessionToBeRemoved);
+            return returnValue;
+        });
+
+        assertThat(withRealm(realmId, (session, realm) -> session.sessions().getClientSession(origSessions[0], realm.getClientByClientId("test-app"), testAppClientSessionId, false)), nullValue());
+    }
 }