keycloak-aplcache

diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedFederatedIdentityLinks.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedFederatedIdentityLinks.java
new file mode 100644
index 0000000..d52f971
--- /dev/null
+++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedFederatedIdentityLinks.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2016 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.cache.infinispan.entities;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import org.keycloak.models.FederatedIdentityModel;
+import org.keycloak.models.RealmModel;
+
+/**
+ * The cache entry, which contains list of all identityProvider links for particular user. It needs to be updated every time when any
+ * federation link is added, removed or updated for the user
+ * 
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class CachedFederatedIdentityLinks extends AbstractRevisioned implements InRealm {
+
+    private final String realmId;
+    private final Set<FederatedIdentityModel> federatedIdentities = new HashSet<>();
+
+    public CachedFederatedIdentityLinks(Long revision, String id, RealmModel realm, Set<FederatedIdentityModel> federatedIdentities) {
+        super(revision, id);
+        this.realmId = realm.getId();
+        this.federatedIdentities.addAll(federatedIdentities);
+    }
+
+    @Override
+    public String getRealm() {
+        return realmId;
+    }
+
+    public Set<FederatedIdentityModel> getFederatedIdentities() {
+        return federatedIdentities;
+    }
+}

diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/UserCacheSession.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/UserCacheSession.java
index 64f4ad5..1f56be9 100755
--- a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/UserCacheSession.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/UserCacheSession.java
@@ -33,6 +33,7 @@ import org.keycloak.models.UserFederationProviderModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserProvider;
 import org.keycloak.models.cache.CacheUserProvider;
+import org.keycloak.models.cache.infinispan.entities.CachedFederatedIdentityLinks;
 import org.keycloak.models.cache.infinispan.entities.CachedUser;
 import org.keycloak.models.cache.infinispan.entities.UserListQuery;
 
@@ -85,6 +86,7 @@ public class UserCacheSession implements CacheUserProvider {
         invalidations.add(user.getId());
         if (user.getEmail() != null) invalidations.add(getUserByEmailCacheKey(realm.getId(), user.getEmail()));
         invalidations.add(getUserByUsernameCacheKey(realm.getId(), user.getUsername()));
+        if (realm.isIdentityFederationEnabled()) invalidations.add(getFederatedIdentityLinksCacheKey(user.getId()));
     }
 
     protected void runInvalidations() {
@@ -180,6 +182,14 @@ public class UserCacheSession implements CacheUserProvider {
         return realmId + ".email." + email;
     }
 
+    public String getUserByFederatedIdentityCacheKey(String realmId, FederatedIdentityModel socialLink) {
+        return realmId + ".idp." + socialLink.getIdentityProvider() + "." + socialLink.getUserId();
+    }
+
+    public String getFederatedIdentityLinksCacheKey(String userId) {
+        return userId + ".idplinks";
+    }
+
     @Override
     public UserModel getUserByUsername(String username, RealmModel realm) {
         logger.tracev("getUserByUsername: {0}", username);
@@ -283,7 +293,46 @@ public class UserCacheSession implements CacheUserProvider {
 
     @Override
     public UserModel getUserByFederatedIdentity(FederatedIdentityModel socialLink, RealmModel realm) {
-        return getDelegate().getUserByFederatedIdentity(socialLink, realm);
+        if (socialLink == null) return null;
+        if (!realm.isIdentityFederationEnabled()) return null;
+
+        if (realmInvalidations.contains(realm.getId())) {
+            return getDelegate().getUserByFederatedIdentity(socialLink, realm);
+        }
+        String cacheKey = getUserByFederatedIdentityCacheKey(realm.getId(), socialLink);
+        if (invalidations.contains(cacheKey)) {
+            return getDelegate().getUserByFederatedIdentity(socialLink, realm);
+        }
+        UserListQuery query = cache.get(cacheKey, UserListQuery.class);
+
+        String userId = null;
+        if (query == null) {
+            Long loaded = cache.getCurrentRevision(cacheKey);
+            UserModel model = getDelegate().getUserByFederatedIdentity(socialLink, realm);
+            if (model == null) return null;
+            userId = model.getId();
+            query = new UserListQuery(loaded, cacheKey, realm, userId);
+            cache.addRevisioned(query, startupRevision);
+            if (invalidations.contains(userId)) return model;
+            if (managedUsers.containsKey(userId)) return managedUsers.get(userId);
+
+            CachedUser cached = cache.get(userId, CachedUser.class);
+            if (cached == null) {
+                cached = new CachedUser(loaded, realm, model);
+                cache.addRevisioned(cached, startupRevision);
+            }
+            UserAdapter adapter = new UserAdapter(cached, this, session, realm);
+            managedUsers.put(userId, adapter);
+            return adapter;
+        } else {
+            userId = query.getUsers().iterator().next();
+            if (invalidations.contains(userId)) {
+                invalidations.add(cacheKey);
+                return getDelegate().getUserByFederatedIdentity(socialLink, realm);
+
+            }
+            return getUserById(userId, realm);
+        }
     }
 
     @Override
@@ -350,19 +399,49 @@ public class UserCacheSession implements CacheUserProvider {
 
     @Override
     public Set<FederatedIdentityModel> getFederatedIdentities(UserModel user, RealmModel realm) {
-        return getDelegate().getFederatedIdentities(user, realm);
+        logger.tracev("getFederatedIdentities: {0}", user.getUsername());
+
+        String cacheKey = getFederatedIdentityLinksCacheKey(user.getId());
+        if (realmInvalidations.contains(realm.getId()) || invalidations.contains(user.getId()) || invalidations.contains(cacheKey)) {
+            return getDelegate().getFederatedIdentities(user, realm);
+        }
+
+        CachedFederatedIdentityLinks cachedLinks = cache.get(cacheKey, CachedFederatedIdentityLinks.class);
+
+        if (cachedLinks == null) {
+            Long loaded = cache.getCurrentRevision(cacheKey);
+            Set<FederatedIdentityModel> federatedIdentities = getDelegate().getFederatedIdentities(user, realm);
+            cachedLinks = new CachedFederatedIdentityLinks(loaded, cacheKey, realm, federatedIdentities);
+            cache.addRevisioned(cachedLinks, startupRevision);
+            return federatedIdentities;
+        } else {
+            return new HashSet<>(cachedLinks.getFederatedIdentities());
+        }
     }
 
     @Override
     public FederatedIdentityModel getFederatedIdentity(UserModel user, String socialProvider, RealmModel realm) {
-        return getDelegate().getFederatedIdentity(user, socialProvider, realm);
+        logger.tracev("getFederatedIdentity: {0} {1}", user.getUsername(), socialProvider);
+
+        String cacheKey = getFederatedIdentityLinksCacheKey(user.getId());
+        if (realmInvalidations.contains(realm.getId()) || invalidations.contains(user.getId()) || invalidations.contains(cacheKey)) {
+            return getDelegate().getFederatedIdentity(user, socialProvider, realm);
+        }
+
+        Set<FederatedIdentityModel> federatedIdentities = getFederatedIdentities(user, realm);
+        for (FederatedIdentityModel socialLink : federatedIdentities) {
+            if (socialLink.getIdentityProvider().equals(socialProvider)) {
+                return socialLink;
+            }
+        }
+        return null;
     }
 
     @Override
     public UserModel addUser(RealmModel realm, String id, String username, boolean addDefaultRoles, boolean addDefaultRequiredActions) {
         UserModel user = getDelegate().addUser(realm, id, username, addDefaultRoles, addDefaultRoles);
         // just in case the transaction is rolled back you need to invalidate the user and all cache queries for that user
-        invalidateUser(realm, user);
+        invalidateUser(realm, user, false);
         managedUsers.put(user.getId(), user);
         return user;
     }
@@ -371,37 +450,60 @@ public class UserCacheSession implements CacheUserProvider {
     public UserModel addUser(RealmModel realm, String username) {
         UserModel user = getDelegate().addUser(realm, username);
         // just in case the transaction is rolled back you need to invalidate the user and all cache queries for that user
-        invalidateUser(realm, user);
+        invalidateUser(realm, user, false);
         managedUsers.put(user.getId(), user);
         return user;
     }
 
-    protected void invalidateUser(RealmModel realm, UserModel user) {
+    protected void invalidateUser(RealmModel realm, UserModel user, boolean invalidateUserLookupsByFederatedIdentity) {
         // just in case the transaction is rolled back you need to invalidate the user and all cache queries for that user
+
+        if (realm.isIdentityFederationEnabled()) {
+            // Invalidate federationLinks of user
+            invalidations.add(getFederatedIdentityLinksCacheKey(user.getId()));
+
+            // Invalidate all keys for lookup this user by any identityProvider link
+            if (invalidateUserLookupsByFederatedIdentity) {
+                Set<FederatedIdentityModel> federatedIdentities = getFederatedIdentities(user, realm);
+                for (FederatedIdentityModel socialLink : federatedIdentities) {
+                    String fedIdentityCacheKey = getUserByFederatedIdentityCacheKey(realm.getId(), socialLink);
+                    invalidations.add(fedIdentityCacheKey);
+                }
+            }
+        }
+
         invalidations.add(user.getId());
         if (user.getEmail() != null) invalidations.add(getUserByEmailCacheKey(realm.getId(), user.getEmail()));
         invalidations.add(getUserByUsernameCacheKey(realm.getId(), user.getUsername()));
-
     }
 
     @Override
     public boolean removeUser(RealmModel realm, UserModel user) {
-        invalidateUser(realm, user);
+        invalidateUser(realm, user, true);
         return getDelegate().removeUser(realm, user);
     }
 
     @Override
     public void addFederatedIdentity(RealmModel realm, UserModel user, FederatedIdentityModel socialLink) {
+        invalidations.add(getFederatedIdentityLinksCacheKey(user.getId()));
         getDelegate().addFederatedIdentity(realm, user, socialLink);
     }
 
     @Override
     public void updateFederatedIdentity(RealmModel realm, UserModel federatedUser, FederatedIdentityModel federatedIdentityModel) {
+        invalidations.add(getFederatedIdentityLinksCacheKey(federatedUser.getId()));
         getDelegate().updateFederatedIdentity(realm, federatedUser, federatedIdentityModel);
     }
 
     @Override
     public boolean removeFederatedIdentity(RealmModel realm, UserModel user, String socialProvider) {
+        // Needs to invalidate both directions
+        FederatedIdentityModel socialLink = getFederatedIdentity(user, socialProvider, realm);
+        invalidations.add(getFederatedIdentityLinksCacheKey(user.getId()));
+        if (socialLink != null) {
+            invalidations.add(getUserByFederatedIdentityCacheKey(realm.getId(), socialLink));
+        }
+
         return getDelegate().removeFederatedIdentity(realm, user, socialProvider);
     }