keycloak-aplcache

diff --git a/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProvider.java b/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProvider.java
index 52d47df..edfd392 100644
--- a/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProvider.java
+++ b/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProvider.java
@@ -19,6 +19,7 @@ package org.keycloak.keys.infinispan;
 
 import java.security.PublicKey;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
@@ -32,6 +33,7 @@ import org.keycloak.common.util.Time;
 import org.keycloak.keys.PublicKeyLoader;
 import org.keycloak.keys.PublicKeyStorageProvider;
 import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakTransaction;
 import org.keycloak.models.cache.infinispan.ClearCacheEvent;
 import org.keycloak.models.cache.infinispan.InfinispanCacheRealmProviderFactory;
 
@@ -51,13 +53,17 @@ public class InfinispanPublicKeyStorageProvider implements PublicKeyStorageProvi
 
     private final int minTimeBetweenRequests ;
 
+    private Set<String> invalidations = new HashSet<>();
+
     public InfinispanPublicKeyStorageProvider(KeycloakSession session, Cache<String, PublicKeysEntry> keys, Map<String, FutureTask<PublicKeysEntry>> tasksInProgress, int minTimeBetweenRequests) {
         this.session = session;
         this.keys = keys;
         this.tasksInProgress = tasksInProgress;
         this.minTimeBetweenRequests = minTimeBetweenRequests;
+        session.getTransactionManager().enlistAfterCompletion(getAfterTransaction());
     }
 
+
     @Override
     public void clearCache() {
         keys.clear();
@@ -65,6 +71,56 @@ public class InfinispanPublicKeyStorageProvider implements PublicKeyStorageProvi
         cluster.notify(InfinispanPublicKeyStorageProviderFactory.KEYS_CLEAR_CACHE_EVENTS, new ClearCacheEvent(), true);
     }
 
+
+    void addInvalidation(String cacheKey) {
+        this.invalidations.add(cacheKey);
+    }
+
+
+    protected KeycloakTransaction getAfterTransaction() {
+        return new KeycloakTransaction() {
+
+            @Override
+            public void begin() {
+            }
+
+            @Override
+            public void commit() {
+                runInvalidations();
+            }
+
+            @Override
+            public void rollback() {
+                runInvalidations();
+            }
+
+            @Override
+            public void setRollbackOnly() {
+            }
+
+            @Override
+            public boolean getRollbackOnly() {
+                return false;
+            }
+
+            @Override
+            public boolean isActive() {
+                return true;
+            }
+        };
+    }
+
+
+    protected void runInvalidations() {
+        ClusterProvider cluster = session.getProvider(ClusterProvider.class);
+
+        for (String cacheKey : invalidations) {
+            keys.remove(cacheKey);
+            cluster.notify(cacheKey, PublicKeyStorageInvalidationEvent.create(cacheKey), true);
+        }
+    }
+
+
     @Override
     public PublicKey getPublicKey(String modelKey, String kid, PublicKeyLoader loader) {
         // Check if key is in cache

diff --git a/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProviderFactory.java
index 8f2e321..42a73fc 100644
--- a/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProviderFactory.java
+++ b/model/infinispan/src/main/java/org/keycloak/keys/infinispan/InfinispanPublicKeyStorageProviderFactory.java
@@ -30,9 +30,14 @@ import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.keys.PublicKeyStorageProvider;
 import org.keycloak.keys.PublicKeyStorageSpi;
 import org.keycloak.keys.PublicKeyStorageProviderFactory;
+import org.keycloak.keys.PublicKeyStorageUtils;
+import org.keycloak.models.ClientModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.RealmModel;
 import org.keycloak.models.cache.infinispan.events.InvalidationEvent;
+import org.keycloak.provider.ProviderEvent;
+import org.keycloak.provider.ProviderEventListener;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -45,7 +50,7 @@ public class InfinispanPublicKeyStorageProviderFactory implements PublicKeyStora
 
     public static final String KEYS_CLEAR_CACHE_EVENTS = "KEYS_CLEAR_CACHE_EVENTS";
 
-    private Cache<String, PublicKeysEntry> keysCache;
+    private volatile Cache<String, PublicKeysEntry> keysCache;
 
     private final Map<String, FutureTask<PublicKeysEntry>> tasksInProgress = new ConcurrentHashMap<>();
 
@@ -64,6 +69,15 @@ public class InfinispanPublicKeyStorageProviderFactory implements PublicKeyStora
                     this.keysCache = session.getProvider(InfinispanConnectionProvider.class).getCache(InfinispanConnectionProvider.KEYS_CACHE_NAME);
 
                     ClusterProvider cluster = session.getProvider(ClusterProvider.class);
+                    cluster.registerListener(ClusterProvider.ALL, (ClusterEvent event) -> {
+
+                        if (event instanceof PublicKeyStorageInvalidationEvent) {
+                            PublicKeyStorageInvalidationEvent invalidationEvent = (PublicKeyStorageInvalidationEvent) event;
+                            keysCache.remove(invalidationEvent.getCacheKey());
+                        }
+
+                    });
+
                     cluster.registerListener(KEYS_CLEAR_CACHE_EVENTS, (ClusterEvent event) -> {
 
                         keysCache.clear();
@@ -82,6 +96,55 @@ public class InfinispanPublicKeyStorageProviderFactory implements PublicKeyStora
 
     @Override
     public void postInit(KeycloakSessionFactory factory) {
+        factory.register(new ProviderEventListener() {
+
+            @Override
+            public void onEvent(ProviderEvent event) {
+                if (keysCache == null) {
+                    return;
+                }
+
+                SessionAndKeyHolder cacheKey = getCacheKeyToInvalidate(event);
+                if (cacheKey != null) {
+                    log.debugf("Invalidating %s from keysCache", cacheKey);
+                    InfinispanPublicKeyStorageProvider provider = (InfinispanPublicKeyStorageProvider) cacheKey.session.getProvider(PublicKeyStorageProvider.class, getId());
+                    provider.addInvalidation(cacheKey.cacheKey);
+                }
+            }
+
+        });
+    }
+
+    private SessionAndKeyHolder getCacheKeyToInvalidate(ProviderEvent event) {
+        if (event instanceof RealmModel.ClientUpdatedEvent) {
+            RealmModel.ClientUpdatedEvent eventt = (RealmModel.ClientUpdatedEvent) event;
+            String cacheKey = PublicKeyStorageUtils.getClientModelCacheKey(eventt.getUpdatedClient().getRealm().getId(), eventt.getUpdatedClient().getId());
+            return new SessionAndKeyHolder(eventt.getKeycloakSession(), cacheKey);
+        } else if (event instanceof RealmModel.ClientRemovedEvent) {
+            RealmModel.ClientRemovedEvent eventt = (RealmModel.ClientRemovedEvent) event;
+            String cacheKey = PublicKeyStorageUtils.getClientModelCacheKey(eventt.getClient().getRealm().getId(), eventt.getClient().getId());
+            return new SessionAndKeyHolder(eventt.getKeycloakSession(), cacheKey);
+        } else if (event instanceof RealmModel.IdentityProviderUpdatedEvent) {
+            RealmModel.IdentityProviderUpdatedEvent eventt = (RealmModel.IdentityProviderUpdatedEvent) event;
+            String cacheKey = PublicKeyStorageUtils.getIdpModelCacheKey(eventt.getRealm().getId(), eventt.getUpdatedIdentityProvider().getInternalId());
+            return new SessionAndKeyHolder(eventt.getKeycloakSession(), cacheKey);
+        } else if (event instanceof RealmModel.IdentityProviderRemovedEvent) {
+            RealmModel.IdentityProviderRemovedEvent eventt = (RealmModel.IdentityProviderRemovedEvent) event;
+            String cacheKey = PublicKeyStorageUtils.getIdpModelCacheKey(eventt.getRealm().getId(), eventt.getRemovedIdentityProvider().getInternalId());
+            return new SessionAndKeyHolder(eventt.getKeycloakSession(), cacheKey);
+        } else {
+            return null;
+        }
+    }
+
+    private class SessionAndKeyHolder {
+        private final KeycloakSession session;
+        private final String cacheKey;
+
+        public SessionAndKeyHolder(KeycloakSession session, String cacheKey) {
+            this.session = session;
+            this.cacheKey = cacheKey;
+        }
 
     }
 

diff --git a/model/infinispan/src/main/java/org/keycloak/keys/infinispan/PublicKeyStorageInvalidationEvent.java b/model/infinispan/src/main/java/org/keycloak/keys/infinispan/PublicKeyStorageInvalidationEvent.java
new file mode 100644
index 0000000..1afcf42
--- /dev/null
+++ b/model/infinispan/src/main/java/org/keycloak/keys/infinispan/PublicKeyStorageInvalidationEvent.java
@@ -0,0 +1,48 @@
+/*
+ * 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.keys.infinispan;
+
+import org.keycloak.models.cache.infinispan.events.InvalidationEvent;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class PublicKeyStorageInvalidationEvent extends InvalidationEvent {
+
+    private String cacheKey;
+
+    public static PublicKeyStorageInvalidationEvent create(String cacheKey) {
+        PublicKeyStorageInvalidationEvent event = new PublicKeyStorageInvalidationEvent();
+        event.cacheKey = cacheKey;
+        return event;
+    }
+
+    @Override
+    public String getId() {
+        return cacheKey;
+    }
+
+    public String getCacheKey() {
+        return cacheKey;
+    }
+
+    @Override
+    public String toString() {
+        return "PublicKeyStorageInvalidationEvent [ " + cacheKey + " ]";
+    }
+}

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java b/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java
index b3f58b4..830da5f 100755
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java
@@ -960,23 +960,7 @@ public class RealmAdapter implements RealmModel, JpaModel<RealmEntity> {
         List<IdentityProviderModel> identityProviders = new ArrayList<IdentityProviderModel>();
 
         for (IdentityProviderEntity entity: entities) {
-            IdentityProviderModel identityProviderModel = new IdentityProviderModel();
-            identityProviderModel.setProviderId(entity.getProviderId());
-            identityProviderModel.setAlias(entity.getAlias());
-            identityProviderModel.setDisplayName(entity.getDisplayName());
-
-            identityProviderModel.setInternalId(entity.getInternalId());
-            Map<String, String> config = entity.getConfig();
-            Map<String, String> copy = new HashMap<>();
-            copy.putAll(config);
-            identityProviderModel.setConfig(copy);
-            identityProviderModel.setEnabled(entity.isEnabled());
-            identityProviderModel.setTrustEmail(entity.isTrustEmail());
-            identityProviderModel.setAuthenticateByDefault(entity.isAuthenticateByDefault());
-            identityProviderModel.setFirstBrokerLoginFlowId(entity.getFirstBrokerLoginFlowId());
-            identityProviderModel.setPostBrokerLoginFlowId(entity.getPostBrokerLoginFlowId());
-            identityProviderModel.setStoreToken(entity.isStoreToken());
-            identityProviderModel.setAddReadTokenRoleOnCreate(entity.isAddReadTokenRoleOnCreate());
+            IdentityProviderModel identityProviderModel = entityToModel(entity);
 
             identityProviders.add(identityProviderModel);
         }
@@ -984,6 +968,27 @@ public class RealmAdapter implements RealmModel, JpaModel<RealmEntity> {
         return Collections.unmodifiableList(identityProviders);
     }
 
+    private IdentityProviderModel entityToModel(IdentityProviderEntity entity) {
+        IdentityProviderModel identityProviderModel = new IdentityProviderModel();
+        identityProviderModel.setProviderId(entity.getProviderId());
+        identityProviderModel.setAlias(entity.getAlias());
+        identityProviderModel.setDisplayName(entity.getDisplayName());
+
+        identityProviderModel.setInternalId(entity.getInternalId());
+        Map<String, String> config = entity.getConfig();
+        Map<String, String> copy = new HashMap<>();
+        copy.putAll(config);
+        identityProviderModel.setConfig(copy);
+        identityProviderModel.setEnabled(entity.isEnabled());
+        identityProviderModel.setTrustEmail(entity.isTrustEmail());
+        identityProviderModel.setAuthenticateByDefault(entity.isAuthenticateByDefault());
+        identityProviderModel.setFirstBrokerLoginFlowId(entity.getFirstBrokerLoginFlowId());
+        identityProviderModel.setPostBrokerLoginFlowId(entity.getPostBrokerLoginFlowId());
+        identityProviderModel.setStoreToken(entity.isStoreToken());
+        identityProviderModel.setAddReadTokenRoleOnCreate(entity.isAddReadTokenRoleOnCreate());
+        return identityProviderModel;
+    }
+
     @Override
     public IdentityProviderModel getIdentityProviderByAlias(String alias) {
         for (IdentityProviderModel identityProviderModel : getIdentityProviders()) {
@@ -1024,8 +1029,28 @@ public class RealmAdapter implements RealmModel, JpaModel<RealmEntity> {
     public void removeIdentityProviderByAlias(String alias) {
         for (IdentityProviderEntity entity : realm.getIdentityProviders()) {
             if (entity.getAlias().equals(alias)) {
+
                 em.remove(entity);
                 em.flush();
+
+                session.getKeycloakSessionFactory().publish(new RealmModel.IdentityProviderRemovedEvent() {
+
+                    @Override
+                    public RealmModel getRealm() {
+                        return RealmAdapter.this;
+                    }
+
+                    @Override
+                    public IdentityProviderModel getRemovedIdentityProvider() {
+                        return entityToModel(entity);
+                    }
+
+                    @Override
+                    public KeycloakSession getKeycloakSession() {
+                        return session;
+                    }
+                });
+
             }
         }
     }
@@ -1048,6 +1073,24 @@ public class RealmAdapter implements RealmModel, JpaModel<RealmEntity> {
         }
 
         em.flush();
+
+        session.getKeycloakSessionFactory().publish(new RealmModel.IdentityProviderUpdatedEvent() {
+
+            @Override
+            public RealmModel getRealm() {
+                return RealmAdapter.this;
+            }
+
+            @Override
+            public IdentityProviderModel getUpdatedIdentityProvider() {
+                return identityProvider;
+            }
+
+            @Override
+            public KeycloakSession getKeycloakSession() {
+                return session;
+            }
+        });
     }
 
     @Override

diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java
index d581710..429d389 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java
@@ -775,23 +775,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         List<IdentityProviderModel> identityProviders = new ArrayList<IdentityProviderModel>();
 
         for (IdentityProviderEntity entity: entities) {
-            IdentityProviderModel identityProviderModel = new IdentityProviderModel();
-
-            identityProviderModel.setProviderId(entity.getProviderId());
-            identityProviderModel.setAlias(entity.getAlias());
-            identityProviderModel.setDisplayName(entity.getDisplayName());
-            identityProviderModel.setInternalId(entity.getInternalId());
-            Map<String, String> config = entity.getConfig();
-            Map<String, String> copy = new HashMap<>();
-            copy.putAll(config);
-            identityProviderModel.setConfig(copy);
-            identityProviderModel.setEnabled(entity.isEnabled());
-            identityProviderModel.setTrustEmail(entity.isTrustEmail());
-            identityProviderModel.setAuthenticateByDefault(entity.isAuthenticateByDefault());
-            identityProviderModel.setFirstBrokerLoginFlowId(entity.getFirstBrokerLoginFlowId());
-            identityProviderModel.setPostBrokerLoginFlowId(entity.getPostBrokerLoginFlowId());
-            identityProviderModel.setStoreToken(entity.isStoreToken());
-            identityProviderModel.setAddReadTokenRoleOnCreate(entity.isAddReadTokenRoleOnCreate());
+            IdentityProviderModel identityProviderModel = entityToModel(entity);
 
             identityProviders.add(identityProviderModel);
         }
@@ -799,6 +783,27 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         return Collections.unmodifiableList(identityProviders);
     }
 
+    private IdentityProviderModel entityToModel(IdentityProviderEntity entity) {
+        IdentityProviderModel identityProviderModel = new IdentityProviderModel();
+
+        identityProviderModel.setProviderId(entity.getProviderId());
+        identityProviderModel.setAlias(entity.getAlias());
+        identityProviderModel.setDisplayName(entity.getDisplayName());
+        identityProviderModel.setInternalId(entity.getInternalId());
+        Map<String, String> config = entity.getConfig();
+        Map<String, String> copy = new HashMap<>();
+        copy.putAll(config);
+        identityProviderModel.setConfig(copy);
+        identityProviderModel.setEnabled(entity.isEnabled());
+        identityProviderModel.setTrustEmail(entity.isTrustEmail());
+        identityProviderModel.setAuthenticateByDefault(entity.isAuthenticateByDefault());
+        identityProviderModel.setFirstBrokerLoginFlowId(entity.getFirstBrokerLoginFlowId());
+        identityProviderModel.setPostBrokerLoginFlowId(entity.getPostBrokerLoginFlowId());
+        identityProviderModel.setStoreToken(entity.isStoreToken());
+        identityProviderModel.setAddReadTokenRoleOnCreate(entity.isAddReadTokenRoleOnCreate());
+        return identityProviderModel;
+    }
+
     @Override
     public IdentityProviderModel getIdentityProviderByAlias(String alias) {
         for (IdentityProviderModel identityProviderModel : getIdentityProviders()) {
@@ -837,6 +842,25 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
             if (entity.getAlias().equals(alias)) {
                 realm.getIdentityProviders().remove(entity);
                 updateRealm();
+
+                session.getKeycloakSessionFactory().publish(new RealmModel.IdentityProviderRemovedEvent() {
+
+                    @Override
+                    public RealmModel getRealm() {
+                        return RealmAdapter.this;
+                    }
+
+                    @Override
+                    public IdentityProviderModel getRemovedIdentityProvider() {
+                        return entityToModel(entity);
+                    }
+
+                    @Override
+                    public KeycloakSession getKeycloakSession() {
+                        return session;
+                    }
+                });
+
                 break;
             }
         }
@@ -860,6 +884,24 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         }
 
         updateRealm();
+
+        session.getKeycloakSessionFactory().publish(new RealmModel.IdentityProviderUpdatedEvent() {
+
+            @Override
+            public RealmModel getRealm() {
+                return RealmAdapter.this;
+            }
+
+            @Override
+            public IdentityProviderModel getUpdatedIdentityProvider() {
+                return identityProvider;
+            }
+
+            @Override
+            public KeycloakSession getKeycloakSession() {
+                return session;
+            }
+        });
     }
 
     @Override

diff --git a/server-spi/src/main/java/org/keycloak/models/RealmModel.java b/server-spi/src/main/java/org/keycloak/models/RealmModel.java
index 4409b9d..3149f50 100755
--- a/server-spi/src/main/java/org/keycloak/models/RealmModel.java
+++ b/server-spi/src/main/java/org/keycloak/models/RealmModel.java
@@ -67,6 +67,18 @@ public interface RealmModel extends RoleContainerModel {
         KeycloakSession getKeycloakSession();
     }
 
+    interface IdentityProviderUpdatedEvent extends ProviderEvent {
+        RealmModel getRealm();
+        IdentityProviderModel getUpdatedIdentityProvider();
+        KeycloakSession getKeycloakSession();
+    }
+
+    interface IdentityProviderRemovedEvent extends ProviderEvent {
+        RealmModel getRealm();
+        IdentityProviderModel getRemovedIdentityProvider();
+        KeycloakSession getKeycloakSession();
+    }
+
     String getId();
 
     String getName();

diff --git a/server-spi-private/src/main/java/org/keycloak/keys/PublicKeyStorageUtils.java b/server-spi-private/src/main/java/org/keycloak/keys/PublicKeyStorageUtils.java
new file mode 100644
index 0000000..52eb7db
--- /dev/null
+++ b/server-spi-private/src/main/java/org/keycloak/keys/PublicKeyStorageUtils.java
@@ -0,0 +1,33 @@
+/*
+ * 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.keys;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class PublicKeyStorageUtils {
+
+    public static String getClientModelCacheKey(String realmId, String clientUuid) {
+        return realmId + "::client::" + clientUuid;
+    }
+
+    public static String getIdpModelCacheKey(String realmId, String idpInternalId) {
+        return realmId + "::idp::" + idpInternalId;
+    }
+
+}

diff --git a/services/src/main/java/org/keycloak/keys/loader/PublicKeyStorageManager.java b/services/src/main/java/org/keycloak/keys/loader/PublicKeyStorageManager.java
index 811dacf..d4507e9 100644
--- a/services/src/main/java/org/keycloak/keys/loader/PublicKeyStorageManager.java
+++ b/services/src/main/java/org/keycloak/keys/loader/PublicKeyStorageManager.java
@@ -22,6 +22,7 @@ import java.security.PublicKey;
 import org.keycloak.broker.oidc.OIDCIdentityProviderConfig;
 import org.keycloak.jose.jws.JWSInput;
 import org.keycloak.keys.PublicKeyStorageProvider;
+import org.keycloak.keys.PublicKeyStorageUtils;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
@@ -36,27 +37,20 @@ public class PublicKeyStorageManager {
 
         PublicKeyStorageProvider keyStorage = session.getProvider(PublicKeyStorageProvider.class);
 
-        String modelKey = getModelKey(client);
+        String modelKey = PublicKeyStorageUtils.getClientModelCacheKey(client.getRealm().getId(), client.getId());
         ClientPublicKeyLoader loader = new ClientPublicKeyLoader(session, client);
         return keyStorage.getPublicKey(modelKey, kid, loader);
     }
 
-    private static String getModelKey(ClientModel client) {
-        return client.getRealm().getId() + "::client::" + client.getId();
-    }
-
 
     public static PublicKey getIdentityProviderPublicKey(KeycloakSession session, RealmModel realm, OIDCIdentityProviderConfig idpConfig, JWSInput input) {
         String kid = input.getHeader().getKeyId();
 
         PublicKeyStorageProvider keyStorage = session.getProvider(PublicKeyStorageProvider.class);
 
-        String modelKey = getModelKey(realm, idpConfig);
+        String modelKey = PublicKeyStorageUtils.getIdpModelCacheKey(realm.getId(), idpConfig.getInternalId());
         OIDCIdentityProviderPublicKeyLoader loader = new OIDCIdentityProviderPublicKeyLoader(session, idpConfig);
         return keyStorage.getPublicKey(modelKey, kid, loader);
     }
 
-    private static String getModelKey(RealmModel realm, OIDCIdentityProviderConfig idpConfig) {
-        return realm.getId() + "::idp::" + idpConfig.getInternalId();
-    }
 }

diff --git a/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/rest/resource/TestingOIDCEndpointsApplicationResource.java b/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/rest/resource/TestingOIDCEndpointsApplicationResource.java
index cf0f55e..fd35caf 100644
--- a/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/rest/resource/TestingOIDCEndpointsApplicationResource.java
+++ b/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/rest/resource/TestingOIDCEndpointsApplicationResource.java
@@ -35,6 +35,8 @@ import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
 import javax.ws.rs.QueryParam;
 import javax.ws.rs.core.MediaType;
+
+import java.security.KeyPair;
 import java.security.KeyPairGenerator;
 import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
@@ -63,13 +65,20 @@ public class TestingOIDCEndpointsApplicationResource {
     @NoCache
     public Map<String, String> generateKeys() {
         try {
-            KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
-            generator.initialize(2048);
-            clientData.setSigningKeyPair(generator.generateKeyPair());
-        } catch (NoSuchAlgorithmException e) {
+            KeyPair keyPair = KeyUtils.generateRsaKeyPair(2048);
+            clientData.setSigningKeyPair(keyPair);
+        } catch (Exception e) {
             throw new BadRequestException("Error generating signing keypair", e);
         }
 
+        return getKeysAsPem();
+    }
+
+
+    @GET
+    @Produces(MediaType.APPLICATION_JSON)
+    @Path("/get-keys-as-pem")
+    public Map<String, String> getKeysAsPem() {
         String privateKeyPem = PemUtils.encodeKey(clientData.getSigningKeyPair().getPrivate());
         String publicKeyPem = PemUtils.encodeKey(clientData.getSigningKeyPair().getPublic());
 

diff --git a/testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/client/resources/TestOIDCEndpointsApplicationResource.java b/testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/client/resources/TestOIDCEndpointsApplicationResource.java
index 9c4f324..f8d8e98 100644
--- a/testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/client/resources/TestOIDCEndpointsApplicationResource.java
+++ b/testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/client/resources/TestOIDCEndpointsApplicationResource.java
@@ -37,6 +37,11 @@ public interface TestOIDCEndpointsApplicationResource {
     @Path("/generate-keys")
     Map<String, String> generateKeys();
 
+    @GET
+    @Produces(MediaType.APPLICATION_JSON)
+    @Path("/get-keys-as-pem")
+    Map<String, String> getKeysAsPem();
+
 
     @GET
     @Produces(MediaType.APPLICATION_JSON)

diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcOIDCBrokerWithSignatureTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcOIDCBrokerWithSignatureTest.java
index 5e6b30d..9bdc40e 100644
--- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcOIDCBrokerWithSignatureTest.java
+++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcOIDCBrokerWithSignatureTest.java
@@ -28,6 +28,8 @@ import org.keycloak.admin.client.resource.RealmResource;
 import org.keycloak.common.util.MultivaluedHashMap;
 import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.keys.KeyProvider;
+import org.keycloak.keys.PublicKeyStorageUtils;
+import org.keycloak.keys.loader.PublicKeyStorageManager;
 import org.keycloak.protocol.oidc.OIDCLoginProtocolService;
 import org.keycloak.representations.idm.ClientRepresentation;
 import org.keycloak.representations.idm.ComponentRepresentation;
@@ -107,15 +109,7 @@ public class KcOIDCBrokerWithSignatureTest extends AbstractBaseBrokerTest {
     @Test
     public void testSignatureVerificationJwksUrl() throws Exception {
         // Configure OIDC identity provider with JWKS URL
-        IdentityProviderRepresentation idpRep = getIdentityProvider();
-        OIDCIdentityProviderConfigRep cfg = new OIDCIdentityProviderConfigRep(idpRep);
-        cfg.setValidateSignature(true);
-        cfg.setUseJwksUrl(true);
-
-        UriBuilder b = OIDCLoginProtocolService.certsUrl(UriBuilder.fromUri(OAuthClient.AUTH_SERVER_ROOT));
-        String jwksUrl = b.build(bc.providerRealmName()).toString();
-        cfg.setJwksUrl(jwksUrl);
-        updateIdentityProvider(idpRep);
+        updateIdentityProviderWithJwksUrl();
 
         // Check that user is able to login
         logInAsUserInIDPForFirstTime();
@@ -139,6 +133,19 @@ public class KcOIDCBrokerWithSignatureTest extends AbstractBaseBrokerTest {
         assertLoggedInAccountManagement();
     }
 
+    // Configure OIDC identity provider with JWKS URL and validateSignature=true
+    private void updateIdentityProviderWithJwksUrl() {
+        IdentityProviderRepresentation idpRep = getIdentityProvider();
+        OIDCIdentityProviderConfigRep cfg = new OIDCIdentityProviderConfigRep(idpRep);
+        cfg.setValidateSignature(true);
+        cfg.setUseJwksUrl(true);
+
+        UriBuilder b = OIDCLoginProtocolService.certsUrl(UriBuilder.fromUri(OAuthClient.AUTH_SERVER_ROOT));
+        String jwksUrl = b.build(bc.providerRealmName()).toString();
+        cfg.setJwksUrl(jwksUrl);
+        updateIdentityProvider(idpRep);
+    }
+
 
     @Test
     public void testSignatureVerificationHardcodedPublicKey() throws Exception {
@@ -178,23 +185,17 @@ public class KcOIDCBrokerWithSignatureTest extends AbstractBaseBrokerTest {
     @Test
     public void testClearKeysCache() throws Exception {
         // Configure OIDC identity provider with JWKS URL
-        IdentityProviderRepresentation idpRep = getIdentityProvider();
-        OIDCIdentityProviderConfigRep cfg = new OIDCIdentityProviderConfigRep(idpRep);
-        cfg.setValidateSignature(true);
-        cfg.setUseJwksUrl(true);
-
-        UriBuilder b = OIDCLoginProtocolService.certsUrl(UriBuilder.fromUri(OAuthClient.AUTH_SERVER_ROOT));
-        String jwksUrl = b.build(bc.providerRealmName()).toString();
-        cfg.setJwksUrl(jwksUrl);
-        updateIdentityProvider(idpRep);
+        updateIdentityProviderWithJwksUrl();
 
         // Check that user is able to login
         logInAsUserInIDPForFirstTime();
         assertLoggedInAccountManagement();
 
+        logoutFromRealm(bc.consumerRealmName());
 
         // Check that key is cached
-        String expectedCacheKey = consumerRealm().toRepresentation().getId() + "::idp::" + idpRep.getInternalId();
+        IdentityProviderRepresentation idpRep = getIdentityProvider();
+        String expectedCacheKey = PublicKeyStorageUtils.getIdpModelCacheKey(consumerRealm().toRepresentation().getId(), idpRep.getInternalId());
         TestingCacheResource cache = testingClient.testing(bc.consumerRealmName()).cache(InfinispanConnectionProvider.KEYS_CACHE_NAME);
         Assert.assertTrue(cache.contains(expectedCacheKey));
 
@@ -205,6 +206,40 @@ public class KcOIDCBrokerWithSignatureTest extends AbstractBaseBrokerTest {
     }
 
 
+    // Test that when I update identityProvier, then the record in publicKey cache is cleared and it's not possible to authenticate with it anymore
+    @Test
+    public void testPublicKeyCacheInvalidatedWhenProviderUpdated() throws Exception {
+        // Configure OIDC identity provider with JWKS URL
+        updateIdentityProviderWithJwksUrl();
+
+        // Check that user is able to login
+        logInAsUserInIDPForFirstTime();
+        assertLoggedInAccountManagement();
+
+        logoutFromRealm(bc.consumerRealmName());
+
+        // Check that key is cached
+        IdentityProviderRepresentation idpRep = getIdentityProvider();
+        String expectedCacheKey = PublicKeyStorageUtils.getIdpModelCacheKey(consumerRealm().toRepresentation().getId(), idpRep.getInternalId());
+        TestingCacheResource cache = testingClient.testing(bc.consumerRealmName()).cache(InfinispanConnectionProvider.KEYS_CACHE_NAME);
+        Assert.assertTrue(cache.contains(expectedCacheKey));
+
+        // Update identityProvider to some bad JWKS_URL
+        OIDCIdentityProviderConfigRep cfg = new OIDCIdentityProviderConfigRep(idpRep);
+        cfg.setJwksUrl("http://localhost:43214/non-existent");
+        updateIdentityProvider(idpRep);
+
+        // Check that key is not cached anymore
+        Assert.assertFalse(cache.contains(expectedCacheKey));
+
+        // Check that user is not able to login with IDP
+        setTimeOffset(20);
+        logInAsUserInIDP();
+        assertErrorPage("Unexpected error when authenticating with identity provider");
+    }
+
+
+
     private void rotateKeys() {
         String activeKid = providerRealm().keys().getKeyMetadata().getActive().get("RSA");
 

diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/AbstractClientRegistrationTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/AbstractClientRegistrationTest.java
index c72bdba..f2f8c6d 100644
--- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/AbstractClientRegistrationTest.java
+++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/AbstractClientRegistrationTest.java
@@ -57,6 +57,7 @@ public abstract class AbstractClientRegistrationTest extends AbstractKeycloakTes
     public void addTestRealms(List<RealmRepresentation> testRealms) {
         RealmRepresentation rep = new RealmRepresentation();
         rep.setEnabled(true);
+        rep.setId(REALM_NAME);
         rep.setRealm(REALM_NAME);
         rep.setUsers(new LinkedList<UserRepresentation>());
 

diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/OIDCJwksClientRegistrationTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/OIDCJwksClientRegistrationTest.java
index 5a86d59..74432a8 100644
--- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/OIDCJwksClientRegistrationTest.java
+++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/OIDCJwksClientRegistrationTest.java
@@ -18,6 +18,7 @@
 package org.keycloak.testsuite.client;
 
 import java.security.KeyPair;
+import java.security.KeyPairGenerator;
 import java.security.PrivateKey;
 import java.security.PublicKey;
 import java.util.Collections;
@@ -40,9 +41,14 @@ import org.keycloak.OAuth2Constants;
 import org.keycloak.adapters.authentication.JWTClientCredentialsProvider;
 import org.keycloak.client.registration.Auth;
 import org.keycloak.common.util.KeycloakUriBuilder;
+import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.constants.ServiceUrlConstants;
 import org.keycloak.jose.jwk.JSONWebKeySet;
+import org.keycloak.jose.jwk.JWK;
+import org.keycloak.jose.jwk.JWKBuilder;
 import org.keycloak.jose.jws.JWSBuilder;
+import org.keycloak.keys.PublicKeyStorageUtils;
+import org.keycloak.keys.loader.PublicKeyStorageManager;
 import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.protocol.oidc.OIDCLoginProtocol;
 import org.keycloak.protocol.oidc.OIDCLoginProtocolService;
@@ -139,6 +145,16 @@ public class OIDCJwksClientRegistrationTest extends AbstractClientRegistrationTe
     // The "kid" is set manually to some custom value
     @Test
     public void createClientWithJWKS_customKid() throws Exception {
+        OIDCClientRepresentation response = createClientWithManuallySetKid("a1");
+
+        Map<String, String> generatedKeys = testingClient.testApp().oidcClientEndpoints().getKeysAsPem();
+
+        // Tries to authenticate client with privateKey JWT
+        assertAuthenticateClientSuccess(generatedKeys, response, "a1");
+    }
+
+
+    private OIDCClientRepresentation createClientWithManuallySetKid(String kid) throws Exception {
         OIDCClientRepresentation clientRep = createRep();
 
         clientRep.setGrantTypes(Collections.singletonList(OAuth2Constants.CLIENT_CREDENTIALS));
@@ -146,20 +162,84 @@ public class OIDCJwksClientRegistrationTest extends AbstractClientRegistrationTe
 
         // Generate keys for client
         TestOIDCEndpointsApplicationResource oidcClientEndpointsResource = testingClient.testApp().oidcClientEndpoints();
-        Map<String, String> generatedKeys = oidcClientEndpointsResource.generateKeys();
+        oidcClientEndpointsResource.generateKeys();
 
         JSONWebKeySet keySet = oidcClientEndpointsResource.getJwks();
 
         // Override kid with custom value
-        keySet.getKeys()[0].setKeyId("a1");
+        keySet.getKeys()[0].setKeyId(kid);
         clientRep.setJwks(keySet);
 
-        OIDCClientRepresentation response = reg.oidc().create(clientRep);
+        return reg.oidc().create(clientRep);
+    }
+
+
+    @Test
+    public void testTwoClientsWithSameKid() throws Exception {
+        // Create client with manually set "kid"
+        OIDCClientRepresentation response = createClientWithManuallySetKid("a1");
+
+
+        // Create client2
+        OIDCClientRepresentation clientRep2 = createRep();
+
+        clientRep2.setGrantTypes(Collections.singletonList(OAuth2Constants.CLIENT_CREDENTIALS));
+        clientRep2.setTokenEndpointAuthMethod(OIDCLoginProtocol.PRIVATE_KEY_JWT);
+
+        // Generate some random keys for client2
+        KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
+        generator.initialize(2048);
+        PublicKey client2PublicKey = generator.generateKeyPair().getPublic();
+
+        // Set client2 with manually set "kid" to be same like kid of client1 (but keys for both clients are different)
+        JSONWebKeySet keySet = new JSONWebKeySet();
+        keySet.setKeys(new JWK[]{JWKBuilder.create().kid("a1").rs256(client2PublicKey)});
+
+        clientRep2.setJwks(keySet);
+        clientRep2 = reg.oidc().create(clientRep2);
+
+
+        // Authenticate client1
+        Map<String, String> generatedKeys = testingClient.testApp().oidcClientEndpoints().getKeysAsPem();
+        assertAuthenticateClientSuccess(generatedKeys, response, "a1");
+
+        // Assert item in publicKey cache for client1
+        String expectedCacheKey = PublicKeyStorageUtils.getClientModelCacheKey(REALM_NAME, response.getClientId());
+        Assert.assertTrue(testingClient.testing().cache(InfinispanConnectionProvider.KEYS_CACHE_NAME).contains(expectedCacheKey));
+
+        // Assert it's not possible to authenticate as client2 with the same "kid" like client1
+        assertAuthenticateClientError(generatedKeys, clientRep2, "a1");
+    }
+
+
+    @Test
+    public void testPublicKeyCacheInvalidatedWhenUpdatingClient() throws Exception {
+        OIDCClientRepresentation response = createClientWithManuallySetKid("a1");
+
+        Map<String, String> generatedKeys = testingClient.testApp().oidcClientEndpoints().getKeysAsPem();
 
         // Tries to authenticate client with privateKey JWT
         assertAuthenticateClientSuccess(generatedKeys, response, "a1");
+
+        // Assert item in publicKey cache for client1
+        String expectedCacheKey = PublicKeyStorageUtils.getClientModelCacheKey(REALM_NAME, response.getClientId());
+        Assert.assertTrue(testingClient.testing().cache(InfinispanConnectionProvider.KEYS_CACHE_NAME).contains(expectedCacheKey));
+
+
+
+        // Update client with some bad JWKS_URI
+        response.setJwksUri("http://localhost:4321/non-existent");
+        reg.auth(Auth.token(response.getRegistrationAccessToken()))
+                .oidc().update(response);
+
+        // Assert item not any longer for client1
+        Assert.assertFalse(testingClient.testing().cache(InfinispanConnectionProvider.KEYS_CACHE_NAME).contains(expectedCacheKey));
+
+        // Assert it's not possible to authenticate as client1
+        assertAuthenticateClientError(generatedKeys, response, "a1");
     }
 
+
     @Test
     public void createClientWithJWKSURI() throws Exception {
         OIDCClientRepresentation clientRep = createRep();