Skip to content

Automated Test: feature-idp-cache-implementation #307

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
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);
}
Comment on lines 97 to 113

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Potential NPE if remove() is called with a non-existent alias.

storedIdp fetched at line 100 could be null if the alias doesn't exist. registerIDPLoginInvalidation (line 111) handles null safely via the Objects::nonNull guard in the login predicate, but registerIDPInvalidation (line 103) will NPE since it directly dereferences idp.getInternalId().

🛡️ Proposed null guard
     IdentityProviderModel storedIdp = idpDelegate.getByAlias(alias);
+    if (storedIdp == null) {
+        return idpDelegate.remove(alias);
+    }
     if (isInvalid(cacheKey)) {
🤖 Prompt for AI Agents
In
`@model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/idp/InfinispanIdentityProviderStorageProvider.java`
around lines 97 - 113, The remove method may NPE when
idpDelegate.getByAlias(alias) returns null; before calling
registerIDPInvalidation(storedIdp) or registerIDPLoginInvalidation(storedIdp)
guard against null: fetch storedIdp via idpDelegate.getByAlias(alias) and only
call registerIDPInvalidation/stub registration if storedIdp != null (and
likewise ensure cached != null before using cached.getIdentityProvider());
update the logic in InfinispanIdentityProviderStorageProvider.remove to check
storedIdp and cached identity provider for nulls before dereferencing or passing
into registerIDPInvalidation.


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
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
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