keycloak-memoizeit

diff --git a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/LoginFailureEntity.java b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/LoginFailureEntity.java
index 09bfca6..e742bbc 100644
--- a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/LoginFailureEntity.java
+++ b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/LoginFailureEntity.java
@@ -12,10 +12,6 @@ public class LoginFailureEntity {
     private long lastFailure;
     private String lastIPFailure;
 
-    public String getId() {
-        return realm + ":" + username;
-    }
-
     public String getUsername() {
         return username;
     }

diff --git a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/LoginFailureKey.java b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/LoginFailureKey.java
new file mode 100644
index 0000000..f8dbe87
--- /dev/null
+++ b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/LoginFailureKey.java
@@ -0,0 +1,36 @@
+package org.keycloak.models.sessions.infinispan.entities;
+
+/**
+ * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
+ */
+public class LoginFailureKey {
+
+    private final String realm;
+    private final String username;
+
+    public LoginFailureKey(String realm, String username) {
+        this.realm = realm;
+        this.username = username;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        LoginFailureKey key = (LoginFailureKey) o;
+
+        if (realm != null ? !realm.equals(key.realm) : key.realm != null) return false;
+        if (username != null ? !username.equals(key.username) : key.username != null) return false;
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        int result = realm != null ? realm.hashCode() : 0;
+        result = 31 * result + (username != null ? username.hashCode() : 0);
+        return result;
+    }
+
+}

diff --git a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java
index 60bcb59..617f686 100644
--- a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java
+++ b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java
@@ -13,6 +13,7 @@ import org.keycloak.models.UserSessionProvider;
 import org.keycloak.models.UsernameLoginFailureModel;
 import org.keycloak.models.sessions.infinispan.entities.ClientSessionEntity;
 import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity;
+import org.keycloak.models.sessions.infinispan.entities.LoginFailureKey;
 import org.keycloak.models.sessions.infinispan.entities.SessionEntity;
 import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity;
 import org.keycloak.models.sessions.infinispan.mapreduce.ClientSessionMapper;
@@ -39,10 +40,10 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
 
     private final KeycloakSession session;
     private final Cache<String, SessionEntity> sessionCache;
-    private final Cache<String, LoginFailureEntity> loginFailureCache;
+    private final Cache<LoginFailureKey, LoginFailureEntity> loginFailureCache;
     private final InfinispanKeycloakTransaction tx;
 
-    public InfinispanUserSessionProvider(KeycloakSession session, Cache<String, SessionEntity> sessionCache, Cache<String, LoginFailureEntity> loginFailureCache) {
+    public InfinispanUserSessionProvider(KeycloakSession session, Cache<String, SessionEntity> sessionCache, Cache<LoginFailureKey, LoginFailureEntity> loginFailureCache) {
         this.session = session;
         this.sessionCache = sessionCache;
         this.loginFailureCache = loginFailureCache;
@@ -225,16 +226,18 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
 
     @Override
     public UsernameLoginFailureModel getUserLoginFailure(RealmModel realm, String username) {
-        return wrap(loginFailureCache.get(realm.getId() + ":" + username));
+        LoginFailureKey key = new LoginFailureKey(realm.getId(), username);
+        return wrap(key, loginFailureCache.get(key));
     }
 
     @Override
     public UsernameLoginFailureModel addUserLoginFailure(RealmModel realm, String username) {
+        LoginFailureKey key = new LoginFailureKey(realm.getId(), username);
         LoginFailureEntity entity = new LoginFailureEntity();
         entity.setRealm(realm.getId());
         entity.setUsername(username);
-        tx.put(loginFailureCache, entity.getId(), entity);
-        return wrap(entity);
+        tx.put(loginFailureCache, key, entity);
+        return wrap(key, entity);
     }
 
     @Override
@@ -257,6 +260,9 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
     @Override
     public void onUserRemoved(RealmModel realm, UserModel user) {
         removeUserSessions(realm, user);
+
+        loginFailureCache.remove(new LoginFailureKey(realm.getId(), user.getUsername()));
+        loginFailureCache.remove(new LoginFailureKey(realm.getId(), user.getEmail()));
     }
 
     @Override
@@ -321,8 +327,8 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
     }
 
 
-    UsernameLoginFailureModel wrap(LoginFailureEntity entity) {
-        return entity != null ? new UsernameLoginFailureAdapter(this, loginFailureCache, entity) : null;
+    UsernameLoginFailureModel wrap(LoginFailureKey key, LoginFailureEntity entity) {
+        return entity != null ? new UsernameLoginFailureAdapter(this, loginFailureCache, key, entity) : null;
     }
 
     List<ClientSessionModel> wrapClientSessions(RealmModel realm, Collection<ClientSessionEntity> entities) {
@@ -337,7 +343,7 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
 
         private boolean active;
         private boolean rollback;
-        private Map<String, CacheTask> tasks = new HashMap<String, CacheTask>();
+        private Map<Object, CacheTask> tasks = new HashMap<Object, CacheTask>();
 
         @Override
         public void begin() {
@@ -375,7 +381,7 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
             return active;
         }
 
-        public void put(Cache cache, String key, Object value) {
+        public void put(Cache cache, Object key, Object value) {
             if (tasks.containsKey(key)) {
                 throw new IllegalStateException("Can't add session: task in progress for session");
             } else {
@@ -383,7 +389,7 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
             }
         }
 
-        public void replace(Cache cache, String key, Object value) {
+        public void replace(Cache cache, Object key, Object value) {
             CacheTask current = tasks.get(key);
             if (current != null) {
                 switch (current.operation) {
@@ -406,10 +412,10 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
         public class CacheTask {
             private Cache cache;
             private CacheOperation operation;
-            private String key;
+            private Object key;
             private Object value;
 
-            public CacheTask(Cache cache, CacheOperation operation, String key, Object value) {
+            public CacheTask(Cache cache, CacheOperation operation, Object key, Object value) {
                 this.cache = cache;
                 this.operation = operation;
                 this.key = key;

diff --git a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java
index 6098f3f..977eb3d 100755
--- a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java
+++ b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java
@@ -7,6 +7,7 @@ import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.UserSessionProvider;
 import org.keycloak.models.UserSessionProviderFactory;
 import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity;
+import org.keycloak.models.sessions.infinispan.entities.LoginFailureKey;
 import org.keycloak.models.sessions.infinispan.entities.SessionEntity;
 
 /**
@@ -21,7 +22,7 @@ public class InfinispanUserSessionProviderFactory implements UserSessionProvider
     public UserSessionProvider create(KeycloakSession session) {
         InfinispanConnectionProvider connections = session.getProvider(InfinispanConnectionProvider.class);
         Cache<String, SessionEntity> cache = connections.getCache(SESSION_CACHE_NAME);
-        Cache<String, LoginFailureEntity> loginFailures = connections.getCache(LOGIN_FAILURE_CACHE_NAME);
+        Cache<LoginFailureKey, LoginFailureEntity> loginFailures = connections.getCache(LOGIN_FAILURE_CACHE_NAME);
         return new InfinispanUserSessionProvider(session, cache, loginFailures);
     }
 

diff --git a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/UsernameLoginFailureAdapter.java b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/UsernameLoginFailureAdapter.java
index 2de7089..fed8f28 100755
--- a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/UsernameLoginFailureAdapter.java
+++ b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/UsernameLoginFailureAdapter.java
@@ -3,6 +3,7 @@ package org.keycloak.models.sessions.infinispan;
 import org.infinispan.Cache;
 import org.keycloak.models.UsernameLoginFailureModel;
 import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity;
+import org.keycloak.models.sessions.infinispan.entities.LoginFailureKey;
 
 /**
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
@@ -10,12 +11,14 @@ import org.keycloak.models.sessions.infinispan.entities.LoginFailureEntity;
 public class UsernameLoginFailureAdapter implements UsernameLoginFailureModel {
 
     private InfinispanUserSessionProvider provider;
-    private Cache<String, LoginFailureEntity> cache;
+    private Cache<LoginFailureKey, LoginFailureEntity> cache;
+    private LoginFailureKey key;
     private LoginFailureEntity entity;
 
-    public UsernameLoginFailureAdapter(InfinispanUserSessionProvider provider, Cache<String, LoginFailureEntity> cache, LoginFailureEntity entity) {
+    public UsernameLoginFailureAdapter(InfinispanUserSessionProvider provider, Cache<LoginFailureKey, LoginFailureEntity> cache, LoginFailureKey key, LoginFailureEntity entity) {
         this.provider = provider;
         this.cache = cache;
+        this.key = key;
         this.entity = entity;
     }
 
@@ -75,7 +78,7 @@ public class UsernameLoginFailureAdapter implements UsernameLoginFailureModel {
     }
 
     void update() {
-        provider.getTx().replace(cache, entity.getId(), entity);
+        provider.getTx().replace(cache, key, entity);
     }
 
 }

diff --git a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/UsernameLoginFailureEntity.java b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/UsernameLoginFailureEntity.java
index 899d80e..c62e474 100755
--- a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/UsernameLoginFailureEntity.java
+++ b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/UsernameLoginFailureEntity.java
@@ -18,7 +18,7 @@ import java.io.Serializable;
 @NamedQueries({
         @NamedQuery(name="getAllFailures", query="select failure from UsernameLoginFailureEntity failure"),
         @NamedQuery(name = "removeLoginFailuresByRealm", query = "delete from UsernameLoginFailureEntity f where f.realmId = :realmId"),
-        @NamedQuery(name = "removeLoginFailuresByUser", query = "delete from UsernameLoginFailureEntity f where f.realmId = :realmId and f.username = :username")
+        @NamedQuery(name = "removeLoginFailuresByUser", query = "delete from UsernameLoginFailureEntity f where f.realmId = :realmId and (f.username = :username or f.username = :email)")
 })
 @IdClass(UsernameLoginFailureEntity.Key.class)
 public class UsernameLoginFailureEntity {

diff --git a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java
index a54891d..6708a2d 100755
--- a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java
+++ b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java
@@ -242,7 +242,7 @@ public class JpaUserSessionProvider implements UserSessionProvider {
     @Override
     public void onUserRemoved(RealmModel realm, UserModel user) {
         removeUserSessions(realm, user);
-        em.createNamedQuery("removeLoginFailuresByUser").setParameter("username", user.getUsername()).executeUpdate();
+        em.createNamedQuery("removeLoginFailuresByUser").setParameter("realmId", realm.getId()).setParameter("username", user.getUsername()).setParameter("email", user.getEmail()).executeUpdate();
     }
 
     @Override

diff --git a/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/MemUserSessionProvider.java b/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/MemUserSessionProvider.java
index 53cf7f3..a19af58 100755
--- a/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/MemUserSessionProvider.java
+++ b/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/MemUserSessionProvider.java
@@ -224,13 +224,13 @@ public class MemUserSessionProvider implements UserSessionProvider {
 
     @Override
     public UsernameLoginFailureModel getUserLoginFailure(RealmModel realm, String username) {
-        UsernameLoginFailureEntity entity = loginFailures.get(new UsernameLoginFailureKey(username, realm.getId()));
+        UsernameLoginFailureEntity entity = loginFailures.get(new UsernameLoginFailureKey(realm.getId(), username));
         return entity != null ? new UsernameLoginFailureAdapter(entity) : null;
     }
 
     @Override
     public UsernameLoginFailureModel addUserLoginFailure(RealmModel realm, String username) {
-        UsernameLoginFailureKey key = new UsernameLoginFailureKey(username, realm.getId());
+        UsernameLoginFailureKey key = new UsernameLoginFailureKey(realm.getId(), username);
         UsernameLoginFailureEntity entity = new UsernameLoginFailureEntity(username, realm.getId());
         if (loginFailures.putIfAbsent(key, entity) != null) {
             throw new ModelDuplicateException();
@@ -265,6 +265,7 @@ public class MemUserSessionProvider implements UserSessionProvider {
         removeUserSessions(realm, user);
 
         loginFailures.remove(new UsernameLoginFailureKey(realm.getId(), user.getUsername()));
+        loginFailures.remove(new UsernameLoginFailureKey(realm.getId(), user.getEmail()));
     }
 
     @Override

diff --git a/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/MongoUserSessionProvider.java b/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/MongoUserSessionProvider.java
index e2e9478..6ddebb9 100755
--- a/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/MongoUserSessionProvider.java
+++ b/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/MongoUserSessionProvider.java
@@ -250,6 +250,12 @@ public class MongoUserSessionProvider implements UserSessionProvider {
     @Override
     public void onUserRemoved(RealmModel realm, UserModel user) {
         removeUserSessions(realm, user);
+
+        DBObject query = new QueryBuilder()
+                .or(new BasicDBObject("username", user.getUsername()), new BasicDBObject("username", user.getEmail()))
+                .and("realmId").is(realm.getId())
+                .get();
+        mongoStore.removeEntities(MongoUsernameLoginFailureEntity.class, query, invocationContext);
     }
 
     @Override

diff --git a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java
index 2d9983e..c68ae7d 100755
--- a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java
+++ b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java
@@ -12,6 +12,7 @@ import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.UsernameLoginFailureModel;
 import org.keycloak.protocol.oidc.OpenIDConnect;
+import org.keycloak.services.managers.UserManager;
 import org.keycloak.testsuite.rule.KeycloakRule;
 import org.keycloak.util.Time;
 
@@ -38,8 +39,8 @@ public class UserSessionProviderTest {
     public void before() {
         session = kc.startSession();
         realm = session.realms().getRealm("test");
-        session.users().addUser(realm, "user1");
-        session.users().addUser(realm, "user2");
+        session.users().addUser(realm, "user1").setEmail("user1@localhost");
+        session.users().addUser(realm, "user2").setEmail("user2@localhost");
     }
 
     @After
@@ -48,8 +49,10 @@ public class UserSessionProviderTest {
         session.sessions().removeUserSessions(realm);
         UserModel user1 = session.users().getUserByUsername("user1", realm);
         UserModel user2 = session.users().getUserByUsername("user2", realm);
-        session.users().removeUser(realm, user1);
-        session.users().removeUser(realm, user2);
+
+        UserManager um = new UserManager(session);
+        um.removeUser(realm, user1);
+        um.removeUser(realm, user2);
         kc.stopSession(session, true);
     }
 
@@ -360,6 +363,28 @@ public class UserSessionProviderTest {
         assertEquals(0, failure1.getNumFailures());
     }
 
+    @Test
+    public void testOnUserRemoved() {
+        createSessions();
+
+        session.sessions().addUserLoginFailure(realm, "user1");
+        session.sessions().addUserLoginFailure(realm, "user1@localhost");
+        session.sessions().addUserLoginFailure(realm, "user2");
+
+        resetSession();
+
+        session.sessions().onUserRemoved(realm, session.users().getUserByUsername("user1", realm));
+
+        resetSession();
+
+        assertTrue(session.sessions().getUserSessions(realm, session.users().getUserByUsername("user1", realm)).isEmpty());
+        assertFalse(session.sessions().getUserSessions(realm, session.users().getUserByUsername("user2", realm)).isEmpty());
+
+        assertNull(session.sessions().getUserLoginFailure(realm, "user1"));
+        assertNull(session.sessions().getUserLoginFailure(realm, "user1@localhost"));
+        assertNotNull(session.sessions().getUserLoginFailure(realm, "user2"));
+    }
+
     private ClientSessionModel createClientSession(ClientModel client, UserSessionModel userSession, String redirect, String state, Set<String> roles) {
         ClientSessionModel clientSession = session.sessions().createClientSession(realm, client);
         if (userSession != null) clientSession.setUserSession(userSession);