Mongo related fixes

audit/mongo/pom.xml

@@ -37,6 +37,11 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.jboss.logging</groupId>
+            <artifactId>jboss-logging</artifactId>
+            <scope>provided</scope>
+        </dependency>
         <dependency>
             <groupId>org.mongodb</groupId>
             <artifactId>mongo-java-driver</artifactId>
@@ -85,9 +90,6 @@
                                 <keycloak.audit.mongo.db>${keycloak.audit.mongo.db}</keycloak.audit.mongo.db>
                                 <keycloak.audit.mongo.clearOnStartup>${keycloak.audit.mongo.clearOnStartup}</keycloak.audit.mongo.clearOnStartup>
                             </systemPropertyVariables>
-                            <dependenciesToScan>
-                                <dependency>org.keycloak:keycloak-model-tests</dependency>
-                            </dependenciesToScan>
                         </configuration>
                     </execution>
                     <execution>

audit/mongo/src/main/java/org/keycloak/audit/mongo/MongoAuditProviderFactory.java

@@ -5,6 +5,7 @@ import com.mongodb.MongoClient;
 import com.mongodb.MongoCredential;
 import com.mongodb.ServerAddress;
 import com.mongodb.WriteConcern;
+import org.jboss.logging.Logger;
 import org.keycloak.Config;
 import org.keycloak.audit.AuditProvider;
 import org.keycloak.audit.AuditProviderFactory;
@@ -21,6 +22,8 @@ import java.util.Set;
  */
 public class MongoAuditProviderFactory implements AuditProviderFactory {
 
+    protected static final Logger logger = Logger.getLogger(MongoAuditProviderFactory.class);
+
     public static final String ID = "mongo";
     private MongoClient client;
     private DB db;
@@ -55,6 +58,8 @@ public class MongoAuditProviderFactory implements AuditProviderFactory {
             if (clearOnStartup) {
                 db.getCollection("audit").drop();
             }
+
+            logger.infof("Initialized mongo audit. host: %s, port: %d, db: %s, clearOnStartup: %b", host, port, dbName, clearOnStartup);
         } catch (UnknownHostException e) {
             throw new RuntimeException(e);
         }

model/mongo/src/main/java/org/keycloak/models/mongo/impl/MongoStoreImpl.java

@@ -237,7 +237,7 @@ public class MongoStoreImpl implements MongoStore {
     public <T extends MongoIdentifiableEntity> T loadEntity(Class<T> type, String id, MongoStoreInvocationContext context) {
         // First look if we already read the object with this oid and type during this transaction. If yes, use it instead of DB lookup
         T cached = context.getLoadedEntity(type, id);
-        if (cached != null) return cached;
+        if (cached != null && type.isAssignableFrom(cached.getClass())) return cached;
 
         DBCollection dbCollection = getDBCollectionForType(type);
 

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ApplicationAdapter.java

@@ -22,7 +22,7 @@ import java.util.Set;
  */
 public class ApplicationAdapter extends ClientAdapter<MongoApplicationEntity> implements ApplicationModel {
 
-    public ApplicationAdapter(RealmModel realm, MongoApplicationEntity applicationEntity, MongoStoreInvocationContext invContext) {
+    public ApplicationAdapter(RealmAdapter realm, MongoApplicationEntity applicationEntity, MongoStoreInvocationContext invContext) {
         super(realm, applicationEntity, invContext);
     }
 

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ClientAdapter.java

@@ -9,12 +9,11 @@ import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.RealmModel;
-import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.entities.ClientEntity;
 import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
-import org.keycloak.models.mongo.keycloak.entities.MongoClientUserSessionAssociationEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoUserSessionEntity;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -22,9 +21,9 @@ import org.keycloak.models.mongo.keycloak.entities.MongoClientUserSessionAssocia
 public class ClientAdapter<T extends MongoIdentifiableEntity> extends AbstractMongoAdapter<T> implements ClientModel {
 
     protected final T clientEntity;
-    private final RealmModel realm;
+    private final RealmAdapter realm;
 
-    public ClientAdapter(RealmModel realm, T clientEntity, MongoStoreInvocationContext invContext) {
+    public ClientAdapter(RealmAdapter realm, T clientEntity, MongoStoreInvocationContext invContext) {
         super(invContext);
         this.clientEntity = clientEntity;
         this.realm = realm;
@@ -154,7 +153,7 @@ public class ClientAdapter<T extends MongoIdentifiableEntity> extends AbstractMo
     }
 
     @Override
-    public RealmModel getRealm() {
+    public RealmAdapter getRealm() {
         return realm;
     }
 
@@ -172,14 +171,13 @@ public class ClientAdapter<T extends MongoIdentifiableEntity> extends AbstractMo
     @Override
     public Set<UserSessionModel> getUserSessions() {
         DBObject query = new QueryBuilder()
-                .and("clientId").is(getId())
+                .and("associatedClientIds").is(getId())
                 .get();
-        List<MongoClientUserSessionAssociationEntity> associations = getMongoStore().loadEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
+        List<MongoUserSessionEntity> sessions = getMongoStore().loadEntities(MongoUserSessionEntity.class, query, invocationContext);
 
         Set<UserSessionModel> result = new HashSet<UserSessionModel>();
-        for (MongoClientUserSessionAssociationEntity association : associations) {
+        for (MongoUserSessionEntity session : sessions) {
-            UserSessionModel session = realm.getUserSession(association.getSessionId());
+            result.add(new UserSessionAdapter(session, realm, invocationContext));
-            result.add(session);
         }
         return result;
 

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSessionFactory.java

@@ -64,7 +64,7 @@ public class MongoKeycloakSessionFactory implements KeycloakSessionFactory {
         try {
             String host = config.get("host", ServerAddress.defaultHost());
             int port = config.getInt("port", ServerAddress.defaultPort());
-            String dbName = config.get("db", "keycloak-audit");
+            String dbName = config.get("db", "keycloak");
             boolean clearOnStartup = config.getBoolean("clearOnStartup", false);
 
             String user = config.get("user");
@@ -79,6 +79,8 @@ public class MongoKeycloakSessionFactory implements KeycloakSessionFactory {
             DB db = client.getDB(dbName);
 
             this.mongoStore = new MongoStoreImpl(db, clearOnStartup, MANAGED_ENTITY_TYPES);
+
+            logger.infof("Initialized mongo model. host: %s, port: %d, db: %s, clearOnStartup: %b", host, port, dbName, clearOnStartup);
         } catch (UnknownHostException e) {
             throw new RuntimeException(e);
         }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/OAuthClientAdapter.java

@@ -10,7 +10,7 @@ import org.keycloak.models.mongo.keycloak.entities.MongoOAuthClientEntity;
  */
 public class OAuthClientAdapter extends ClientAdapter<MongoOAuthClientEntity> implements OAuthClientModel {
 
-    public OAuthClientAdapter(RealmModel realm, MongoOAuthClientEntity oauthClientEntity, MongoStoreInvocationContext invContext) {
+    public OAuthClientAdapter(RealmAdapter realm, MongoOAuthClientEntity oauthClientEntity, MongoStoreInvocationContext invContext) {
         super(realm, oauthClientEntity, invContext);
     }
 

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RoleAdapter.java

@@ -7,7 +7,6 @@ import java.util.Set;
 
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
-import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleContainerModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
@@ -25,13 +24,13 @@ public class RoleAdapter extends AbstractMongoAdapter<MongoRoleEntity> implement
 
     private final MongoRoleEntity role;
     private RoleContainerModel roleContainer;
-    private RealmModel realm;
+    private RealmAdapter realm;
 
-    public RoleAdapter(RealmModel realm, MongoRoleEntity roleEntity, MongoStoreInvocationContext invContext) {
+    public RoleAdapter(RealmAdapter realm, MongoRoleEntity roleEntity, MongoStoreInvocationContext invContext) {
         this(realm, roleEntity, null, invContext);
     }
 
-    public RoleAdapter(RealmModel realm, MongoRoleEntity roleEntity, RoleContainerModel roleContainer, MongoStoreInvocationContext invContext) {
+    public RoleAdapter(RealmAdapter realm, MongoRoleEntity roleEntity, RoleContainerModel roleContainer, MongoStoreInvocationContext invContext) {
         super(invContext);
         this.role = roleEntity;
         this.roleContainer = roleContainer;

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserSessionAdapter.java

@@ -2,27 +2,27 @@ package org.keycloak.models.mongo.keycloak.adapters;
 
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
+import org.jboss.logging.Logger;
+import org.keycloak.models.ApplicationModel;
 import org.keycloak.models.ClientModel;
-import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
+import org.keycloak.models.entities.ClientEntity;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
-import org.keycloak.models.mongo.keycloak.entities.MongoClientUserSessionAssociationEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoApplicationEntity;
-import org.keycloak.models.mongo.keycloak.entities.MongoRealmEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoOAuthClientEntity;
-import org.keycloak.models.mongo.keycloak.entities.MongoRoleEntity;
-import org.keycloak.models.mongo.keycloak.entities.MongoUserEntity;
 import org.keycloak.models.mongo.keycloak.entities.MongoUserSessionEntity;
 
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 
 /**
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
  */
 public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEntity> implements UserSessionModel {
 
+    private static final Logger logger = Logger.getLogger(RealmAdapter.class);
+
     private MongoUserSessionEntity entity;
     private RealmAdapter realm;
 
@@ -46,6 +46,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setId(String id) {
         entity.setId(id);
+        updateMongoEntity();
     }
 
     @Override
@@ -56,6 +57,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setUser(UserModel user) {
         entity.setUser(user.getId());
+        updateMongoEntity();
     }
 
     @Override
@@ -66,6 +68,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setIpAddress(String ipAddress) {
         entity.setIpAddress(ipAddress);
+        updateMongoEntity();
     }
 
     @Override
@@ -76,6 +79,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setStarted(int started) {
         entity.setStarted(started);
+        updateMongoEntity();
     }
 
     @Override
@@ -86,58 +90,39 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setLastSessionRefresh(int seconds) {
         entity.setLastSessionRefresh(seconds);
+        updateMongoEntity();
     }
 
     @Override
     public void associateClient(ClientModel client) {
-        List<ClientModel> clients = getClientAssociations();
+        getMongoStore().pushItemToList(entity, "associatedClientIds", client.getId(), true, invocationContext);
-        for (ClientModel ass : clients) {
-            if (ass.getId().equals(client.getId())) return;
-        }
-
-        MongoClientUserSessionAssociationEntity association = new MongoClientUserSessionAssociationEntity();
-        association.setClientId(client.getId());
-        association.setSessionId(getId());
-
-        getMongoStore().insertEntity(association, invocationContext);
     }
 
     @Override
     public List<ClientModel> getClientAssociations() {
-        DBObject query = new QueryBuilder()
+        List<String> associatedClientIds = getMongoEntity().getAssociatedClientIds();
-                .and("sessionId").is(getId())
-                .get();
-        List<MongoClientUserSessionAssociationEntity> associations = getMongoStore().loadEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
 
-        List<ClientModel> result = new ArrayList<ClientModel>();
+        List<ClientModel> clients = new ArrayList<ClientModel>();
-        for (MongoClientUserSessionAssociationEntity association : associations) {
+        for (String clientId : associatedClientIds) {
-            ClientModel client = realm.findClientById(association.getClientId());
+            // Try application first
-            result.add(client);
+            ClientModel client = realm.getApplicationById(clientId);
+
+            // And then OAuthClient
+            if (client == null) {
+                client = realm.getOAuthClientById(clientId);
+            }
+
+            if (client != null) {
+                clients.add(client);
+            } else {
+                logger.warnf("Not found associated client with Id: %s", clientId);
+            }
         }
-        return result;
+        return clients;
     }
 
     @Override
     public void removeAssociatedClient(ClientModel client) {
-        DBObject query = new QueryBuilder()
+        getMongoStore().pullItemFromList(entity, "associatedClientIds", client.getId(), invocationContext);
-                .and("sessionId").is(getId())
-                .and("clientId").is(client.getId())
-                .get();
-        getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        UserSessionAdapter that = (UserSessionAdapter) o;
-        return getId().equals(that.getId());
-    }
-
-    @Override
-    public int hashCode() {
-        return getId().hashCode();
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoApplicationEntity.java

@@ -1,7 +1,10 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
+import java.util.List;
+
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
+import org.jboss.logging.Logger;
 import org.keycloak.models.entities.ApplicationEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
@@ -23,10 +26,13 @@ public class MongoApplicationEntity extends ApplicationEntity implements MongoId
                 .get();
         context.getMongoStore().removeEntities(MongoRoleEntity.class, query, context);
 
+        // Remove all session associations
         query = new QueryBuilder()
-                .and("clientId").is(getId())
+                .and("associatedClientIds").is(getId())
                 .get();
-        context.getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, context);
+        List<MongoUserSessionEntity> sessions = context.getMongoStore().loadEntities(MongoUserSessionEntity.class, query, context);
-
+        for (MongoUserSessionEntity session : sessions) {
+            context.getMongoStore().pullItemFromList(session, "associatedClientIds", getId(), context);
+        }
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoClientUserSessionAssociationEntity.java

@@ -1,37 +0,0 @@
-package org.keycloak.models.mongo.keycloak.entities;
-
-import org.keycloak.models.entities.AbstractIdentifiableEntity;
-import org.keycloak.models.mongo.api.MongoCollection;
-import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
-import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
-
-/**
- * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
- * @version $Revision: 1 $
- */
-@MongoCollection(collectionName = "session-client-associations")
-public class MongoClientUserSessionAssociationEntity extends AbstractIdentifiableEntity implements MongoIdentifiableEntity {
-    private String clientId;
-    private String sessionId;
-
-    public String getClientId() {
-        return clientId;
-    }
-
-    public void setClientId(String clientId) {
-        this.clientId = clientId;
-    }
-
-    public String getSessionId() {
-        return sessionId;
-    }
-
-    public void setSessionId(String sessionId) {
-        this.sessionId = sessionId;
-    }
-
-    @Override
-    public void afterRemove(MongoStoreInvocationContext invocationContext) {
-    }
-
-}

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoOAuthClientEntity.java

@@ -1,5 +1,7 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
+import java.util.List;
+
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.keycloak.models.entities.OAuthClientEntity;
@@ -16,10 +18,14 @@ import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 public class MongoOAuthClientEntity extends OAuthClientEntity implements MongoIdentifiableEntity {
 
     @Override
-    public void afterRemove(MongoStoreInvocationContext invocationContext) {
+    public void afterRemove(MongoStoreInvocationContext context) {
+        // Remove all session associations
         DBObject query = new QueryBuilder()
-                .and("clientId").is(getId())
+                .and("associatedClientIds").is(getId())
                 .get();
-        invocationContext.getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
+        List<MongoUserSessionEntity> sessions = context.getMongoStore().loadEntities(MongoUserSessionEntity.class, query, context);
+        for (MongoUserSessionEntity session : sessions) {
+            context.getMongoStore().pullItemFromList(session, "associatedClientIds", getId(), context);
+        }
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoUserSessionEntity.java

@@ -1,7 +1,8 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
-import com.mongodb.DBObject;
+import java.util.ArrayList;
-import com.mongodb.QueryBuilder;
+import java.util.List;
+
 import org.keycloak.models.entities.AbstractIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
@@ -23,6 +24,8 @@ public class MongoUserSessionEntity extends AbstractIdentifiableEntity implement
 
     private int lastSessionRefresh;
 
+    private List<String> associatedClientIds = new ArrayList<String>();
+
     public String getRealmId() {
         return realmId;
     }
@@ -63,13 +66,16 @@ public class MongoUserSessionEntity extends AbstractIdentifiableEntity implement
         this.lastSessionRefresh = lastSessionRefresh;
     }
 
+    public List<String> getAssociatedClientIds() {
+        return associatedClientIds;
+    }
+
+    public void setAssociatedClientIds(List<String> associatedClientIds) {
+        this.associatedClientIds = associatedClientIds;
+    }
+
     @Override
     public void afterRemove(MongoStoreInvocationContext context) {
-        // Remove all roles, which belongs to this application
-        DBObject query = new QueryBuilder()
-                .and("sessionId").is(getId())
-                .get();
-        context.getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, context);
     }
 
 }

model/tests/src/test/java/org/keycloak/model/test/AdapterTest.java

@@ -762,4 +762,49 @@ public class AdapterTest extends AbstractModelTest {
         assertNull(realmManager.getRealmByName("userSessions").getUserSession(userSession.getId()));
     }
 
+    @Test
+    public void userSessionAssociations() {
+        RealmModel realm = realmManager.createRealm("userSessions");
+        UserModel user = realm.addUser("userSessions1");
+        UserSessionModel userSession = realm.createUserSession(user, "127.0.0.1");
+
+        ApplicationModel app1 = realm.addApplication("app1");
+        ApplicationModel app2 = realm.addApplication("app2");
+        OAuthClientModel client1 = realm.addOAuthClient("client1");
+
+        Assert.assertEquals(0, userSession.getClientAssociations().size());
+
+        userSession.associateClient(app1);
+        userSession.associateClient(client1);
+
+        Assert.assertEquals(2, userSession.getClientAssociations().size());
+        Assert.assertTrue(app1.getUserSessions().contains(userSession));
+        Assert.assertFalse(app2.getUserSessions().contains(userSession));
+        Assert.assertTrue(client1.getUserSessions().contains(userSession));
+
+        commit();
+
+        // Refresh all
+        realm = realmManager.getRealm("userSessions");
+        userSession = realm.getUserSession(userSession.getId());
+        app1 = realm.getApplicationByName("app1");
+        client1 = realm.getOAuthClient("client1");
+
+        userSession.removeAssociatedClient(app1);
+        Assert.assertEquals(1, userSession.getClientAssociations().size());
+        Assert.assertEquals(client1, userSession.getClientAssociations().get(0));
+        Assert.assertFalse(app1.getUserSessions().contains(userSession));
+
+        commit();
+
+        // Refresh all
+        realm = realmManager.getRealm("userSessions");
+        userSession = realm.getUserSession(userSession.getId());
+        client1 = realm.getOAuthClient("client1");
+
+        userSession.removeAssociatedClient(client1);
+        Assert.assertEquals(0, userSession.getClientAssociations().size());
+        Assert.assertFalse(client1.getUserSessions().contains(userSession));
+    }
+
 }

testsuite/integration/pom.xml

@@ -439,7 +439,7 @@
             <id>mongo</id>
             <activation>
                 <property>
-                    <name>keycloak.model</name>
+                    <name>keycloak.model.provider</name>
                     <value>mongo</value>
                 </property>
             </activation>

testsuite/integration/src/main/resources/META-INF/keycloak-server.json

@@ -4,11 +4,23 @@
     },
 
     "audit": {
-        "provider": "${keycloak.audit.provider:jpa}"
+        "provider": "${keycloak.audit.provider,keycloak.model.provider:jpa}",
+        "mongo": {
+          "host": "${keycloak.audit.mongo.host:127.0.0.1}",
+          "port": "${keycloak.audit.mongo.port:27017}",
+          "db": "${keycloak.audit.mongo.db:keycloak-audit}",
+          "clearOnStartup": "${keycloak.model.mongo.clearOnStartup:false}"
+        }
     },
 
     "model": {
-        "provider": "${keycloak.model.provider:jpa}"
+        "provider": "${keycloak.model.provider:jpa}",
+        "mongo": {
+          "host": "${keycloak.model.mongo.host:127.0.0.1}",
+          "port": "${keycloak.model.mongo.port:27017}",
+          "db": "${keycloak.model.mongo.db:keycloak}",
+          "clearOnStartup": "${keycloak.model.mongo.clearOnStartup:false}"
+        }
     },
 
     "timer": {