IDP storage provider Infinispan implementation

Closes #31251

Signed-off-by: vramik <vramik@redhat.com>

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/CachedCount.java

@@ -14,19 +14,19 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.keycloak.models.cache.infinispan.organization;
+package org.keycloak.models.cache.infinispan;
 
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.cache.infinispan.entities.AbstractRevisioned;
 import org.keycloak.models.cache.infinispan.entities.InRealm;
 
-public class CachedOrganizationCount extends AbstractRevisioned implements InRealm {
+public class CachedCount extends AbstractRevisioned implements InRealm {
 
     private final RealmModel realm;
     private final long count;
 
-    public CachedOrganizationCount(Long revision, RealmModel realm, long count) {
-        super(revision, InfinispanOrganizationProvider.cacheKeyOrgCount(realm));
+    public CachedCount(Long revision, RealmModel realm, String cacheKey, long count) {
+        super(revision, cacheKey);
         this.realm = realm;
         this.count = count;
     }

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/idp/CachedIdentityProvider.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2024 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.idp;
+
+import org.keycloak.models.IdentityProviderModel;
+import org.keycloak.models.RealmModel;
+import org.keycloak.models.cache.infinispan.entities.AbstractRevisioned;
+import org.keycloak.models.cache.infinispan.entities.InRealm;
+
+public class CachedIdentityProvider extends AbstractRevisioned implements InRealm {
+
+    private final RealmModel realm;
+    private final IdentityProviderModel idp;
+
+    public CachedIdentityProvider(Long revision, RealmModel realm, String cacheKey, IdentityProviderModel idp) {
+        super(revision, cacheKey);
+        this.realm = realm;
+        this.idp = idp;
+    }
+
+    @Override
+    public String getRealm() {
+        return realm.getId();
+    }
+
+    public IdentityProviderModel getIdentityProvider() {
+        return new IdentityProviderModel(idp);
+    }
+
+}

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/idp/InfinispanIDPProvider.java

@@ -0,0 +1,186 @@
+/*
+ * Copyright 2024 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.idp;
+
+import org.keycloak.models.cache.infinispan.CachedCount;
+import java.util.Map;
+import java.util.stream.Stream;
+import org.keycloak.models.IDPProvider;
+import org.keycloak.models.IdentityProviderModel;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.RealmModel;
+import org.keycloak.models.cache.CacheRealmProvider;
+import org.keycloak.models.cache.infinispan.RealmCacheSession;
+
+public class InfinispanIDPProvider implements IDPProvider {
+
+    private static final String IDP_COUNT_KEY_SUFFIX = ".idp.count";
+    private static final String IDP_ALIAS_KEY_SUFFIX = ".idp.alias";
+
+    private final KeycloakSession session;
+    private final IDPProvider idpDelegate;
+    private final RealmCacheSession realmCache;
+    
+    public InfinispanIDPProvider(KeycloakSession session) {
+        this.session = session;
+        this.idpDelegate = session.getProvider(IDPProvider.class, "jpa");
+        this.realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class);
+    }
+
+    private static String cacheKeyIdpCount(RealmModel realm) {
+        return realm.getId() + IDP_COUNT_KEY_SUFFIX;
+    }
+
+    private static String cacheKeyIdpAlias(RealmModel realm, String alias) {
+        return realm.getId() + "." + alias + IDP_ALIAS_KEY_SUFFIX;
+    }
+
+    @Override
+    public IdentityProviderModel create(IdentityProviderModel model) {
+        registerCountInvalidation();
+        return idpDelegate.create(model);
+    }
+
+    @Override
+    public void update(IdentityProviderModel model) {
+        // for cases the alias is being updated, it is needed to lookup the idp by id to obtain the original alias
+        IdentityProviderModel idpById = getById(model.getInternalId());
+        registerIDPInvalidation(idpById);
+        idpDelegate.update(model);
+    }
+
+    @Override
+    public boolean remove(String alias) {
+        String cacheKey = cacheKeyIdpAlias(getRealm(), alias);
+        if (isInvalid(cacheKey)) {
+            //lookup idp by alias in cache to be able to invalidate its internalId
+            registerIDPInvalidation(idpDelegate.getByAlias(alias));
+        } else {
+            CachedIdentityProvider cached = realmCache.getCache().get(cacheKey, CachedIdentityProvider.class);
+            if (cached != null) {
+                registerIDPInvalidation(cached.getIdentityProvider());
+            }
+        }
+        registerCountInvalidation();
+        return idpDelegate.remove(alias);
+    }
+
+    @Override
+    public void removeAll() {
+        registerCountInvalidation();
+        // no need to invalidate each entry in cache, removeAll() is (currently) called only in case the realm is being deleted
+        idpDelegate.removeAll();
+    }
+
+    @Override
+    public IdentityProviderModel getById(String internalId) {
+        CachedIdentityProvider cached = realmCache.getCache().get(internalId, CachedIdentityProvider.class);
+        String realmId = getRealm().getId();
+        if (cached != null && !cached.getRealm().equals(realmId)) {
+            cached = null;
+        }
+
+        if (cached == null) {
+            Long loaded = realmCache.getCache().getCurrentRevision(internalId);
+            IdentityProviderModel model = idpDelegate.getById(internalId);
+            if (model == null) return null;
+            if (isInvalid(internalId)) return model;
+            cached = new CachedIdentityProvider(loaded, getRealm(), internalId, model);
+            realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
+        } else if (isInvalid(internalId)) {
+            return idpDelegate.getById(internalId);
+        }
+        return cached.getIdentityProvider();
+    }
+
+    @Override
+    public IdentityProviderModel getByAlias(String alias) {
+        String cacheKey = cacheKeyIdpAlias(getRealm(), alias);
+
+        if (isInvalid(cacheKey)) {
+            return idpDelegate.getByAlias(alias);
+        }
+
+        CachedIdentityProvider cached = realmCache.getCache().get(cacheKey, CachedIdentityProvider.class);
+
+        if (cached == null) {
+            Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
+            IdentityProviderModel model = idpDelegate.getByAlias(alias);
+            if (model == null) {
+                return null;
+            }
+            cached = new CachedIdentityProvider(loaded, getRealm(), cacheKey, model);
+            realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
+        }
+
+        return cached.getIdentityProvider();
+    }
+
+    @Override
+    public Stream<IdentityProviderModel> getAllStream(String search, Integer first, Integer max) {
+        return idpDelegate.getAllStream(search, first, max);
+    }
+
+    @Override
+    public Stream<IdentityProviderModel> getAllStream(Map<String, String> attrs, Integer first, Integer max) {
+        return idpDelegate.getAllStream(attrs, first, max);
+    }
+
+    @Override
+    public long count() {
+        String cacheKey = cacheKeyIdpCount(getRealm());
+        CachedCount cached = realmCache.getCache().get(cacheKey, CachedCount.class);
+
+        // cached and not invalidated
+        if (cached != null && !isInvalid(cacheKey)) {
+            return cached.getCount();
+        }
+
+        Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
+        long count = idpDelegate.count();
+        cached = new CachedCount(loaded, getRealm(), cacheKey, count);
+        realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
+
+        return count;
+    }
+
+    @Override
+    public void close() {
+        idpDelegate.close();
+    }
+
+    private void registerIDPInvalidation(IdentityProviderModel idp) {
+        realmCache.registerInvalidation(idp.getInternalId());
+        realmCache.registerInvalidation(cacheKeyIdpAlias(getRealm(), idp.getAlias()));
+    }
+
+    private void registerCountInvalidation() {
+        realmCache.registerInvalidation(cacheKeyIdpCount(getRealm()));
+    }
+
+    private RealmModel getRealm() {
+        RealmModel realm = session.getContext().getRealm();
+        if (realm == null) {
+            throw new IllegalArgumentException("Session not bound to a realm");
+        }
+        return realm;
+    }
+
+    private boolean isInvalid(String cacheKey) {
+        return realmCache.getInvalidations().contains(cacheKey);
+    }
+}

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/idp/InfinispanIDPProviderFactory.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2024 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.idp;
+
+import org.keycloak.Config;
+import org.keycloak.models.IDPProvider;
+import org.keycloak.models.IDPProviderFactory;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+
+public class InfinispanIDPProviderFactory implements IDPProviderFactory<IDPProvider>{
+
+    public static final String PROVIDER_ID = "infinispan";
+
+    @Override
+    public IDPProvider create(KeycloakSession session) {
+        return new InfinispanIDPProvider(session);
+    }
+
+    @Override
+    public void init(Config.Scope config) {
+    }
+
+    @Override
+    public void postInit(KeycloakSessionFactory factory) {
+    }
+
+    @Override
+    public void close() {
+    }
+
+    @Override
+    public String getId() {
+        return PROVIDER_ID;
+    }
+
+    @Override
+    public int order() {
+        return 10;
+    }
+}

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java

@@ -26,6 +26,7 @@ import org.keycloak.models.OrganizationModel;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.cache.CacheRealmProvider;
+import org.keycloak.models.cache.infinispan.CachedCount;
 import org.keycloak.models.cache.infinispan.RealmCacheSession;
 import org.keycloak.organization.OrganizationProvider;
 
@@ -44,7 +45,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         this.realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class);
     }
 
-    static String cacheKeyOrgCount(RealmModel realm) {
+    private static String cacheKeyOrgCount(RealmModel realm) {
         return realm.getId() + ORG_COUNT_KEY_SUFFIX;
     }
 
@@ -251,7 +252,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
     @Override
     public long count() {
         String cacheKey = cacheKeyOrgCount(getRealm());
-        CachedOrganizationCount cached = realmCache.getCache().get(cacheKey, CachedOrganizationCount.class);
+        CachedCount cached = realmCache.getCache().get(cacheKey, CachedCount.class);
 
         // cached and not invalidated
         if (cached != null && !isInvalid(cacheKey)) {
@@ -260,7 +261,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
 
         Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
         long count = orgDelegate.count();
-        cached = new CachedOrganizationCount(loaded, getRealm(), count);
+        cached = new CachedCount(loaded, getRealm(), cacheKey, count);
         realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
 
         return count;

model/infinispan/src/main/resources/META-INF/services/org.keycloak.models.IDPProviderFactory

@@ -0,0 +1,18 @@
+#
+# Copyright 2024 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.
+#
+
+org.keycloak.models.cache.infinispan.idp.InfinispanIDPProviderFactory

model/jpa/src/main/java/org/keycloak/models/jpa/JpaIDPProvider.java

@@ -142,13 +142,16 @@ public class JpaIDPProvider implements IDPProvider {
         IdentityProviderEntity entity = this.getEntityByAlias(alias);
 
         if (entity != null) {
+            //call toModel(entity) now as after em.remove(entity) and the flush it might throw LazyInitializationException 
+            //when accessing the config of the entity (entity.getConfig()) withing the toModel(entity)
+            IdentityProviderModel model = toModel(entity);
+
             em.remove(entity);
             // flush so that constraint violations are flagged and converted into model exception now rather than at the end of the tx.
             em.flush();
 
             // send identity provider removed event.
             RealmModel realm = this.getRealm();
-            IdentityProviderModel model = toModel(entity);
             session.getKeycloakSessionFactory().publish(new RealmModel.IdentityProviderRemovedEvent() {
 
                 @Override