Skip to content

Add caching support for IdentityProviderStorageProvider.getForLogin operations #73

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akshayutture-augment wants to merge 1 commit into
base: feature-idp-cache-baseline
Choose a base branch
from
Open

Add caching support for IdentityProviderStorageProvider.getForLogin operations #73

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 95 additions & 1 deletion ...a/org/keycloak/models/cache/infinispan/idp/InfinispanIdentityProviderStorageProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}

Expand All @@ -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);
}

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions server-spi/src/main/java/org/keycloak/models/IdentityProviderStorageProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}
Expand Down
4 changes: 2 additions & 2 deletions ...loak/organization/forms/login/freemarker/model/OrganizationAwareIdentityProviderBean.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down
Loading
Loading