Mongo related fixes

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) {
-            UserSessionModel session = realm.getUserSession(association.getSessionId());
-            result.add(session);
+        for (MongoUserSessionEntity session : sessions) {
+            result.add(new UserSessionAdapter(session, realm, invocationContext));
         }
         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.MongoRealmEntity;
-import org.keycloak.models.mongo.keycloak.entities.MongoRoleEntity;
-import org.keycloak.models.mongo.keycloak.entities.MongoUserEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoApplicationEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoOAuthClientEntity;
 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();
-        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);
+        getMongoStore().pushItemToList(entity, "associatedClientIds", client.getId(), true, invocationContext);
     }
 
     @Override
     public List<ClientModel> getClientAssociations() {
-        DBObject query = new QueryBuilder()
-                .and("sessionId").is(getId())
-                .get();
-        List<MongoClientUserSessionAssociationEntity> associations = getMongoStore().loadEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
+        List<String> associatedClientIds = getMongoEntity().getAssociatedClientIds();
 
-        List<ClientModel> result = new ArrayList<ClientModel>();
-        for (MongoClientUserSessionAssociationEntity association : associations) {
-            ClientModel client = realm.findClientById(association.getClientId());
-            result.add(client);
+        List<ClientModel> clients = new ArrayList<ClientModel>();
+        for (String clientId : associatedClientIds) {
+            // Try application first
+            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()
-                .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();
+        getMongoStore().pullItemFromList(entity, "associatedClientIds", client.getId(), invocationContext);
     }
 }

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 com.mongodb.QueryBuilder;
+import java.util.ArrayList;
+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);
     }
 
 }