Add cache for IdentityProviderStorageProvider.getForLogin (#32918)

Closes #32573

Signed-off-by: Stefan Guilhen <[email protected]>

Author: sguilhen

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

@@ -18,6 +18,7 @@
 
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -36,11 +37,14 @@
 import org.keycloak.models.cache.infinispan.RealmCacheSession;
 import org.keycloak.organization.OrganizationProvider;
 
+import static org.keycloak.models.IdentityProviderStorageProvider.LoginFilter.getLoginPredicate;
+
 public class InfinispanIdentityProviderStorageProvider implements IdentityProviderStorageProvider {
 
     private static final String IDP_COUNT_KEY_SUFFIX = ".idp.count";
     private static final String IDP_ALIAS_KEY_SUFFIX = ".idp.alias";
     private static final String IDP_ORG_ID_KEY_SUFFIX = ".idp.orgId";
+    private static final String IDP_LOGIN_SUFFIX = ".idp.login";
 
     private final KeycloakSession session;
     private final IdentityProviderStorageProvider idpDelegate;
@@ -70,9 +74,14 @@ public static String cacheKeyOrgId(RealmModel realm, String orgId) {
         return realm.getId() + "." + orgId + IDP_ORG_ID_KEY_SUFFIX;
     }
 
+    public static String cacheKeyForLogin(RealmModel realm, FetchMode fetchMode) {
+        return realm.getId() + IDP_LOGIN_SUFFIX + "." + fetchMode;
+    }
+
     @Override
     public IdentityProviderModel create(IdentityProviderModel model) {
         registerCountInvalidation();
+        registerIDPLoginInvalidation(model);
         return idpDelegate.create(model);
     }
 
@@ -81,22 +90,25 @@ 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);
+        registerIDPLoginInvalidationOnUpdate(idpById, model);
         idpDelegate.update(model);
     }
 
     @Override
     public boolean remove(String alias) {
         String cacheKey = cacheKeyIdpAlias(getRealm(), alias);
+        IdentityProviderModel storedIdp = idpDelegate.getByAlias(alias);
         if (isInvalid(cacheKey)) {
             //lookup idp by alias in cache to be able to invalidate its internalId
-            registerIDPInvalidation(idpDelegate.getByAlias(alias));
+            registerIDPInvalidation(storedIdp);
         } else {
             CachedIdentityProvider cached = realmCache.getCache().get(cacheKey, CachedIdentityProvider.class);
             if (cached != null) {
                 registerIDPInvalidation(cached.getIdentityProvider());
             }
         }
         registerCountInvalidation();
+        registerIDPLoginInvalidation(storedIdp);
         return idpDelegate.remove(alias);
     }
 
@@ -198,6 +210,50 @@ public Stream<IdentityProviderModel> getByOrganization(String orgId, Integer fir
         return identityProviders.stream();
     }
 
+    @Override
+    public Stream<IdentityProviderModel> getForLogin(FetchMode mode, String organizationId) {
+        String cacheKey = cacheKeyForLogin(getRealm(), mode);
+
+        if (isInvalid(cacheKey)) {
+            return idpDelegate.getForLogin(mode, organizationId).map(this::createOrganizationAwareIdentityProviderModel);
+        }
+
+        RealmCacheManager cache = realmCache.getCache();
+        IdentityProviderListQuery query = cache.get(cacheKey, IdentityProviderListQuery.class);
+        String searchKey = organizationId != null ? organizationId : "";
+        Set<String> cached;
+
+        if (query == null) {
+            // not cached yet
+            Long loaded = cache.getCurrentRevision(cacheKey);
+            cached = idpDelegate.getForLogin(mode, organizationId).map(IdentityProviderModel::getInternalId).collect(Collectors.toSet());
+            query = new IdentityProviderListQuery(loaded, cacheKey, getRealm(), searchKey, cached);
+            cache.addRevisioned(query, startupRevision);
+        } else {
+            cached = query.getIDPs(searchKey);
+            if (cached == null) {
+                // there is a cache entry, but the current search is not yet cached
+                cache.invalidateObject(cacheKey);
+                Long loaded = cache.getCurrentRevision(cacheKey);
+                cached = idpDelegate.getForLogin(mode, organizationId).map(IdentityProviderModel::getInternalId).collect(Collectors.toSet());
+                query = new IdentityProviderListQuery(loaded, cacheKey, getRealm(), searchKey, cached, query);
+                cache.addRevisioned(query, cache.getCurrentCounter());
+            }
+        }
+
+        Set<IdentityProviderModel> identityProviders = new HashSet<>();
+        for (String id : cached) {
+            IdentityProviderModel idp = session.identityProviders().getById(id);
+            if (idp == null) {
+                realmCache.registerInvalidation(cacheKey);
+                return idpDelegate.getForLogin(mode, organizationId).map(this::createOrganizationAwareIdentityProviderModel);
+            }
+            identityProviders.add(idp);
+        }
+
+        return identityProviders.stream();
+    }
+
     @Override
     public Stream<String> getByFlow(String flowId, String search, Integer first, Integer max) {
         return idpDelegate.getByFlow(flowId, search, first, max);
@@ -323,6 +379,44 @@ private void registerIDPMapperInvalidation(IdentityProviderMapperModel mapper) {
         realmCache.registerInvalidation(cacheKeyIdpMapperAliasName(getRealm(), mapper.getIdentityProviderAlias(), mapper.getName()));
     }
 
+    private void registerIDPLoginInvalidation(IdentityProviderModel idp) {
+        // only invalidate login caches if the IDP qualifies as a login IDP.
+        if (getLoginPredicate().test(idp)) {
+            for (FetchMode mode : FetchMode.values()) {
+                realmCache.registerInvalidation(cacheKeyForLogin(getRealm(), mode));
+            }
+        }
+    }
+
+    /**
+     * Registers invalidations for the caches that hold the IDPs available for login when an IDP is updated. The caches
+     * are <strong>NOT</strong> invalidated if:
+     * <ul>
+     *     <li>IDP is currently NOT a login IDP, and the update hasn't changed that (i.e. it continues to be unavailable for login);</li>
+     *     <li>IDP is currently a login IDP, and the update hasn't changed that. This includes the organization link not being updated as well</li>
+     * </ul>
+     * In all other scenarios, the caches must be invalidated.
+     *
+     * @param original the identity provider's current model
+     * @param updated the identity provider's updated model
+     */
+    private void registerIDPLoginInvalidationOnUpdate(IdentityProviderModel original, IdentityProviderModel updated) {
+        // IDP isn't currently available for login and update preserves that - no need to invalidate.
+        if (!getLoginPredicate().test(original) && !getLoginPredicate().test(updated)) {
+            return;
+        }
+        // IDP is currently available for login and update preserves that, including organization link - no need to invalidate.
+        if (getLoginPredicate().test(original) && getLoginPredicate().test(updated)
+                && Objects.equals(original.getOrganizationId(), updated.getOrganizationId())) {
+            return;
+        }
+
+        // all other scenarios should invalidate the login caches.
+        for (FetchMode mode : FetchMode.values()) {
+            realmCache.registerInvalidation(cacheKeyForLogin(getRealm(), mode));
+        }
+    }
+
     private RealmModel getRealm() {
         RealmModel realm = session.getContext().getRealm();
         if (realm == null) {

server-spi/src/main/java/org/keycloak/models/IdentityProviderStorageProvider.java

@@ -251,6 +251,7 @@ public static Map<String, String> getLoginSearchOptions() {
 
         public static Predicate<IdentityProviderModel> getLoginPredicate() {
             return ((Predicate<IdentityProviderModel>) Objects::nonNull)
+                    .and(idp -> idp.getOrganizationId() == null || Boolean.parseBoolean(idp.getConfig().get(OrganizationModel.BROKER_PUBLIC)))
                     .and(Stream.of(values()).map(LoginFilter::getFilter).reduce(Predicate::and).get());
         }
     }

services/src/main/java/org/keycloak/organization/forms/login/freemarker/model/OrganizationAwareIdentityProviderBean.java

@@ -72,12 +72,12 @@ protected List<IdentityProvider> searchForIdentityProviders(String existingIDP)
             }
             // we don't have a specific organization - fetch public enabled IDPs linked to any org.
             return session.identityProviders().getForLogin(ORG_ONLY, null)
-                    .filter(idp -> !Objects.equals(existingIDP, idp.getAlias()))
+                    .filter(idp -> idp.isEnabled() && !Objects.equals(existingIDP, idp.getAlias())) // re-check isEnabled as idp might have been wrapped.
                     .map(idp -> createIdentityProvider(this.realm, this.baseURI, idp))
                     .sorted(IDP_COMPARATOR_INSTANCE).toList();
         }
         return session.identityProviders().getForLogin(ALL, this.organization != null ? this.organization.getId() : null)
-                .filter(idp -> !Objects.equals(existingIDP, idp.getAlias()))
+                .filter(idp -> idp.isEnabled() && !Objects.equals(existingIDP, idp.getAlias())) // re-check isEnabled as idp might have been wrapped.
                 .map(idp -> createIdentityProvider(this.realm, this.baseURI, idp))
                 .sorted(IDP_COMPARATOR_INSTANCE).toList();
     }