From 17e30e9ec10d6243bf1e6758a80c810851ec914b Mon Sep 17 00:00:00 2001 From: Pedro Ruivo Date: Tue, 30 Jul 2024 18:06:54 +0100 Subject: [PATCH] Persist revoke tokens with remote cache feature Stores the revoked tokens into the database and preloads them during startup. Fixes #31760 Signed-off-by: Pedro Ruivo --- ...inispanSingleUseObjectProviderFactory.java | 36 ++----- ...moteInfinispanSingleUseObjectProvider.java | 21 +++- ...inispanSingleUseObjectProviderFactory.java | 101 +++++++++++++++++- .../DefaultDatastoreProviderFactory.java | 35 ++++-- .../SingleUseObjectModelTest.java | 48 ++++++--- 5 files changed, 187 insertions(+), 54 deletions(-) diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanSingleUseObjectProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanSingleUseObjectProviderFactory.java index f3f170af7e..839b817d81 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanSingleUseObjectProviderFactory.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanSingleUseObjectProviderFactory.java @@ -45,12 +45,9 @@ import org.keycloak.models.utils.PostMigrationEvent; import org.keycloak.provider.EnvironmentDependentProviderFactory; import org.keycloak.provider.ProviderConfigProperty; import org.keycloak.provider.ProviderConfigurationBuilder; -import org.keycloak.provider.ProviderEvent; -import org.keycloak.provider.ProviderEventListener; import org.keycloak.provider.ServerInfoAwareProviderFactory; -import org.keycloak.services.scheduled.ClearExpiredRevokedTokens; -import org.keycloak.services.scheduled.ClusterAwareScheduledTaskRunner; -import org.keycloak.timer.TimerProvider; + +import static org.keycloak.storage.datastore.DefaultDatastoreProviderFactory.setupClearExpiredRevokedTokensScheduledTask; /** * @author Marek Posolda @@ -58,7 +55,8 @@ import org.keycloak.timer.TimerProvider; public class InfinispanSingleUseObjectProviderFactory implements SingleUseObjectProviderFactory, EnvironmentDependentProviderFactory, ServerInfoAwareProviderFactory { public static final String CONFIG_PERSIST_REVOKED_TOKENS = "persistRevokedTokens"; - private static final boolean DEFAULT_PERSIST_REVOKED_TOKENS = true; + public static final boolean DEFAULT_PERSIST_REVOKED_TOKENS = true; + public static final String LOADED = "loaded" + SingleUseObjectProvider.REVOKED_KEY; private static final Logger LOG = Logger.getLogger(InfinispanSingleUseObjectProviderFactory.class); @@ -93,8 +91,6 @@ public class InfinispanSingleUseObjectProviderFactory implements SingleUseObject persistRevokedTokens = config.getBoolean(CONFIG_PERSIST_REVOKED_TOKENS, DEFAULT_PERSIST_REVOKED_TOKENS); } - private final static String LOADED = "loaded" + SingleUseObjectProvider.REVOKED_KEY; - private void initialize(KeycloakSession session) { if (persistRevokedTokens && !initialized) { synchronized (this) { @@ -127,25 +123,15 @@ public class InfinispanSingleUseObjectProviderFactory implements SingleUseObject } if (persistRevokedTokens) { - factory.register(new ProviderEventListener() { - public void onEvent(ProviderEvent event) { - if (event instanceof PostMigrationEvent) { - KeycloakSessionFactory sessionFactory = ((PostMigrationEvent) event).getFactory(); - try (KeycloakSession session = sessionFactory.create()) { - TimerProvider timer = session.getProvider(TimerProvider.class); - if (timer != null) { - long interval = Config.scope("scheduled").getLong("interval", 900L) * 1000; - scheduleTask(sessionFactory, timer, interval); - } - // load sessions during startup, not on first request to avoid congestion - initialize(session); - } + factory.register(event -> { + if (event instanceof PostMigrationEvent pme) { + KeycloakSessionFactory sessionFactory = pme.getFactory(); + setupClearExpiredRevokedTokensScheduledTask(sessionFactory); + try (KeycloakSession session = sessionFactory.create()) { + // load sessions during startup, not on first request to avoid congestion + initialize(session); } } - - private void scheduleTask(KeycloakSessionFactory sessionFactory, TimerProvider timer, long interval) { - timer.schedule(new ClusterAwareScheduledTaskRunner(sessionFactory, new ClearExpiredRevokedTokens(), interval), interval); - } }); } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProvider.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProvider.java index b95a51eb54..8cd34dcbc4 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProvider.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProvider.java @@ -18,6 +18,7 @@ package org.keycloak.models.sessions.infinispan.remote; import java.lang.invoke.MethodHandles; +import java.util.Collections; import java.util.Map; import java.util.Objects; import java.util.concurrent.TimeUnit; @@ -33,15 +34,23 @@ import org.keycloak.models.sessions.infinispan.remote.transaction.SingleUseObjec public class RemoteInfinispanSingleUseObjectProvider implements SingleUseObjectProvider { private final static Logger logger = Logger.getLogger(MethodHandles.lookup().lookupClass()); + public static final SingleUseObjectValueEntity REVOKED_TOKEN_VALUE = new SingleUseObjectValueEntity(Collections.emptyMap()); private final SingleUseObjectTransaction transaction; + private final RevokeTokenConsumer revokeTokenConsumer; - public RemoteInfinispanSingleUseObjectProvider(SingleUseObjectTransaction transaction) { + public RemoteInfinispanSingleUseObjectProvider(SingleUseObjectTransaction transaction, RevokeTokenConsumer revokeTokenConsumer) { this.transaction = Objects.requireNonNull(transaction); + this.revokeTokenConsumer = Objects.requireNonNull(revokeTokenConsumer); + } @Override public void put(String key, long lifespanSeconds, Map notes) { + if (key.endsWith(REVOKED_KEY)) { + revokeToken(key, lifespanSeconds); + return; + } transaction.put(key, wrap(notes), lifespanSeconds, TimeUnit.SECONDS); } @@ -93,6 +102,12 @@ public class RemoteInfinispanSingleUseObjectProvider implements SingleUseObjectP return transaction.getCache().withFlags(Flag.FORCE_RETURN_VALUE); } + private void revokeToken(String key, long lifespanSeconds) { + transaction.put(key, REVOKED_TOKEN_VALUE, lifespanSeconds, TimeUnit.SECONDS); + var token = key.substring(0, key.length() - REVOKED_KEY.length()); + revokeTokenConsumer.onTokenRevoke(token, lifespanSeconds); + } + private static Map unwrap(SingleUseObjectValueEntity entity) { return entity == null ? null : entity.getNotes(); } @@ -100,4 +115,8 @@ public class RemoteInfinispanSingleUseObjectProvider implements SingleUseObjectP private static SingleUseObjectValueEntity wrap(Map notes) { return new SingleUseObjectValueEntity(notes); } + + public interface RevokeTokenConsumer { + void onTokenRevoke(String token, long lifespanSeconds); + } } diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProviderFactory.java index 48c5f680a6..46b2373ff6 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProviderFactory.java +++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remote/RemoteInfinispanSingleUseObjectProviderFactory.java @@ -18,41 +18,69 @@ package org.keycloak.models.sessions.infinispan.remote; import java.lang.invoke.MethodHandles; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; import org.infinispan.client.hotrod.RemoteCache; import org.jboss.logging.Logger; import org.keycloak.Config; +import org.keycloak.common.util.Time; import org.keycloak.infinispan.util.InfinispanUtils; import org.keycloak.models.KeycloakSession; import org.keycloak.models.KeycloakSessionFactory; import org.keycloak.models.SingleUseObjectProviderFactory; +import org.keycloak.models.session.RevokedToken; +import org.keycloak.models.session.RevokedTokenPersisterProvider; import org.keycloak.models.sessions.infinispan.entities.SingleUseObjectValueEntity; import org.keycloak.models.sessions.infinispan.remote.transaction.SingleUseObjectTransaction; +import org.keycloak.models.utils.PostMigrationEvent; import org.keycloak.provider.EnvironmentDependentProviderFactory; +import org.keycloak.provider.ProviderConfigProperty; +import org.keycloak.provider.ProviderConfigurationBuilder; +import org.keycloak.provider.ProviderEvent; +import org.keycloak.provider.ProviderEventListener; +import org.keycloak.provider.ServerInfoAwareProviderFactory; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.ACTION_TOKEN_CACHE; import static org.keycloak.connections.infinispan.InfinispanConnectionProvider.getRemoteCache; +import static org.keycloak.models.SingleUseObjectProvider.REVOKED_KEY; +import static org.keycloak.models.sessions.infinispan.InfinispanSingleUseObjectProviderFactory.CONFIG_PERSIST_REVOKED_TOKENS; +import static org.keycloak.models.sessions.infinispan.InfinispanSingleUseObjectProviderFactory.DEFAULT_PERSIST_REVOKED_TOKENS; +import static org.keycloak.models.sessions.infinispan.InfinispanSingleUseObjectProviderFactory.LOADED; +import static org.keycloak.models.sessions.infinispan.remote.RemoteInfinispanSingleUseObjectProvider.REVOKED_TOKEN_VALUE; +import static org.keycloak.models.sessions.infinispan.remote.RemoteInfinispanSingleUseObjectProvider.RevokeTokenConsumer; +import static org.keycloak.storage.datastore.DefaultDatastoreProviderFactory.setupClearExpiredRevokedTokensScheduledTask; -public class RemoteInfinispanSingleUseObjectProviderFactory implements SingleUseObjectProviderFactory, EnvironmentDependentProviderFactory { +public class RemoteInfinispanSingleUseObjectProviderFactory implements SingleUseObjectProviderFactory, EnvironmentDependentProviderFactory, ProviderEventListener, ServerInfoAwareProviderFactory { private final static Logger logger = Logger.getLogger(MethodHandles.lookup().lookupClass()); + private static final RevokeTokenConsumer VOLATILE_REVOKE_TOKEN = (token, lifespanSeconds) -> { + }; + // max of 16 remote cache puts concurrently. + private static final int REVOKED_TOKENS_IMPORT_CONCURRENCY = 16; private volatile RemoteCache cache; + private volatile boolean persistRevokedTokens; @Override public RemoteInfinispanSingleUseObjectProvider create(KeycloakSession session) { assert cache != null; - return new RemoteInfinispanSingleUseObjectProvider(createAndEnlistTransaction(session)); + return new RemoteInfinispanSingleUseObjectProvider(createAndEnlistTransaction(session), createRevokeTokenConsumer(session)); } @Override public void init(Config.Scope config) { - + persistRevokedTokens = config.getBoolean(CONFIG_PERSIST_REVOKED_TOKENS, DEFAULT_PERSIST_REVOKED_TOKENS); } @Override public void postInit(KeycloakSessionFactory factory) { cache = getRemoteCache(factory, ACTION_TOKEN_CACHE); + factory.register(this); logger.debug("Provided initialized."); } @@ -76,9 +104,76 @@ public class RemoteInfinispanSingleUseObjectProviderFactory implements SingleUse return InfinispanUtils.isRemoteInfinispan(); } + @Override + public Map getOperationalInfo() { + Map info = new HashMap<>(); + info.put(CONFIG_PERSIST_REVOKED_TOKENS, Boolean.toString(persistRevokedTokens)); + return info; + } + + @Override + public List getConfigMetadata() { + ProviderConfigurationBuilder builder = ProviderConfigurationBuilder.create(); + builder.property() + .name(CONFIG_PERSIST_REVOKED_TOKENS) + .type("boolean") + .helpText("If revoked tokens are stored persistently across restarts") + .defaultValue(DEFAULT_PERSIST_REVOKED_TOKENS) + .add(); + + return builder.build(); + } + + @Override + public void onEvent(ProviderEvent event) { + if (!(event instanceof PostMigrationEvent pme)) { + return; + } + if (!persistRevokedTokens) { + //nothing to do + return; + } + + // preload revoked tokens from the database and register cleanup expired tokens task + KeycloakSessionFactory sessionFactory = pme.getFactory(); + setupClearExpiredRevokedTokensScheduledTask(sessionFactory); + try (var session = sessionFactory.create()) { + preloadRevokedTokens(session); + } + } + private SingleUseObjectTransaction createAndEnlistTransaction(KeycloakSession session) { var tx = new SingleUseObjectTransaction(cache); session.getTransactionManager().enlistAfterCompletion(tx); return tx; } + + private RevokedTokenPersisterProvider getRevokedTokenPersisterProvider(KeycloakSession session) { + return session.getProvider(RevokedTokenPersisterProvider.class); + } + + private RevokeTokenConsumer createRevokeTokenConsumer(KeycloakSession session) { + return persistRevokedTokens ? getRevokedTokenPersisterProvider(session)::revokeToken : VOLATILE_REVOKE_TOKEN; + } + + private void preloadRevokedTokens(KeycloakSession session) { + var provider = getRevokedTokenPersisterProvider(session); + if (cache.get(LOADED) == null) { + logger.debug("Preloading revoked tokens from database."); + var currentTime = Time.currentTime(); + Flowable.fromStream(provider.getAllRevokedTokens()) + .filter(revokedToken -> revokedToken.expiry() - currentTime > 0) // skip expired tokens + .flatMapCompletable(token -> preloadToken(token, currentTime), false, REVOKED_TOKENS_IMPORT_CONCURRENCY) + .blockingAwait(); + cache.put(LOADED, REVOKED_TOKEN_VALUE); + logger.debug("Preload completed."); + } + } + + private Completable preloadToken(RevokedToken token, long currentTime) { + var lifespan = token.expiry() - currentTime; + return Completable.fromCompletionStage(cache.putIfAbsentAsync(token.tokenId() + REVOKED_KEY, REVOKED_TOKEN_VALUE, lifespan, TimeUnit.SECONDS)); + } + + } diff --git a/model/storage-private/src/main/java/org/keycloak/storage/datastore/DefaultDatastoreProviderFactory.java b/model/storage-private/src/main/java/org/keycloak/storage/datastore/DefaultDatastoreProviderFactory.java index 87d7870615..63ae57739e 100644 --- a/model/storage-private/src/main/java/org/keycloak/storage/datastore/DefaultDatastoreProviderFactory.java +++ b/model/storage-private/src/main/java/org/keycloak/storage/datastore/DefaultDatastoreProviderFactory.java @@ -17,6 +17,9 @@ package org.keycloak.storage.datastore; +import java.util.Arrays; +import java.util.List; + import org.jboss.logging.Logger; import org.keycloak.Config; import org.keycloak.Config.Scope; @@ -29,6 +32,7 @@ import org.keycloak.provider.ProviderEventListener; import org.keycloak.services.scheduled.ClearExpiredAdminEvents; import org.keycloak.services.scheduled.ClearExpiredClientInitialAccessTokens; import org.keycloak.services.scheduled.ClearExpiredEvents; +import org.keycloak.services.scheduled.ClearExpiredRevokedTokens; import org.keycloak.services.scheduled.ClearExpiredUserSessions; import org.keycloak.services.scheduled.ClusterAwareScheduledTaskRunner; import org.keycloak.storage.DatastoreProvider; @@ -38,8 +42,6 @@ import org.keycloak.storage.StoreSyncEvent; import org.keycloak.storage.managers.UserStorageSyncManager; import org.keycloak.timer.ScheduledTask; import org.keycloak.timer.TimerProvider; -import java.util.Arrays; -import java.util.List; public class DefaultDatastoreProviderFactory implements DatastoreProviderFactory, ProviderEventListener { @@ -79,7 +81,7 @@ public class DefaultDatastoreProviderFactory implements DatastoreProviderFactory public String getId() { return PROVIDER_ID; } - + public long getClientStorageProviderTimeout() { return clientStorageProviderTimeout; } @@ -99,20 +101,18 @@ public class DefaultDatastoreProviderFactory implements DatastoreProviderFactory StoreMigrateRepresentationEvent ev = (StoreMigrateRepresentationEvent) event; MigrationModelManager.migrateImport(ev.getSession(), ev.getRealm(), ev.getRep(), ev.isSkipUserDependent()); } - } - - public void setupScheduledTasks(final KeycloakSessionFactory sessionFactory) { - long interval = Config.scope("scheduled").getLong("interval", 900L) * 1000; + } + public static void setupScheduledTasks(final KeycloakSessionFactory sessionFactory) { try (KeycloakSession session = sessionFactory.create()) { TimerProvider timer = session.getProvider(TimerProvider.class); if (timer != null) { - scheduleTasks(sessionFactory, timer, interval); + scheduleTasks(sessionFactory, timer, getScheduledInterval()); } } } - protected void scheduleTasks(KeycloakSessionFactory sessionFactory, TimerProvider timer, long interval) { + protected static void scheduleTasks(KeycloakSessionFactory sessionFactory, TimerProvider timer, long interval) { for (ScheduledTask task : getScheduledTasks()) { scheduleTask(timer, sessionFactory, task, interval); } @@ -120,13 +120,26 @@ public class DefaultDatastoreProviderFactory implements DatastoreProviderFactory UserStorageSyncManager.bootstrapPeriodic(sessionFactory, timer); } - protected List getScheduledTasks() { + protected static List getScheduledTasks() { return Arrays.asList(new ClearExpiredEvents(), new ClearExpiredAdminEvents(), new ClearExpiredClientInitialAccessTokens(), new ClearExpiredUserSessions()); } - protected void scheduleTask(TimerProvider timer, KeycloakSessionFactory sessionFactory, ScheduledTask task, long interval) { + protected static void scheduleTask(TimerProvider timer, KeycloakSessionFactory sessionFactory, ScheduledTask task, long interval) { timer.schedule(new ClusterAwareScheduledTaskRunner(sessionFactory, task, interval), interval); logger.debugf("Scheduled cluster task %s with interval %s ms", task.getTaskName(), interval); } + public static void setupClearExpiredRevokedTokensScheduledTask(KeycloakSessionFactory sessionFactory) { + try (KeycloakSession session = sessionFactory.create()) { + TimerProvider timer = session.getProvider(TimerProvider.class); + if (timer != null) { + scheduleTask(timer, sessionFactory, new ClearExpiredRevokedTokens(), getScheduledInterval()); + } + } + } + + public static long getScheduledInterval() { + return Config.scope("scheduled").getLong("interval", 900L) * 1000; + } + } diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/singleUseObject/SingleUseObjectModelTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/singleUseObject/SingleUseObjectModelTest.java index 6422f8d609..cb9fb721b0 100644 --- a/testsuite/model/src/test/java/org/keycloak/testsuite/model/singleUseObject/SingleUseObjectModelTest.java +++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/singleUseObject/SingleUseObjectModelTest.java @@ -17,20 +17,6 @@ package org.keycloak.testsuite.model.singleUseObject; -import org.hamcrest.Matchers; -import org.junit.Assert; -import org.junit.Test; -import org.keycloak.models.DefaultActionTokenKey; -import org.keycloak.common.util.Time; -import org.keycloak.models.Constants; -import org.keycloak.models.KeycloakSession; -import org.keycloak.models.RealmModel; -import org.keycloak.models.SingleUseObjectProvider; -import org.keycloak.models.UserModel; -import org.keycloak.services.scheduled.ClearExpiredRevokedTokens; -import org.keycloak.testsuite.model.KeycloakModelTest; -import org.keycloak.testsuite.model.RequireProvider; - import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -39,6 +25,25 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import org.hamcrest.Matchers; +import org.infinispan.client.hotrod.RemoteCache; +import org.junit.Assert; +import org.junit.Test; +import org.keycloak.common.util.Time; +import org.keycloak.connections.infinispan.InfinispanConnectionProvider; +import org.keycloak.infinispan.util.InfinispanUtils; +import org.keycloak.models.Constants; +import org.keycloak.models.DefaultActionTokenKey; +import org.keycloak.models.KeycloakSession; +import org.keycloak.models.RealmModel; +import org.keycloak.models.SingleUseObjectProvider; +import org.keycloak.models.UserModel; +import org.keycloak.models.sessions.infinispan.InfinispanSingleUseObjectProviderFactory; +import org.keycloak.models.sessions.infinispan.entities.SingleUseObjectValueEntity; +import org.keycloak.services.scheduled.ClearExpiredRevokedTokens; +import org.keycloak.testsuite.model.KeycloakModelTest; +import org.keycloak.testsuite.model.RequireProvider; + import static org.hamcrest.MatcherAssert.assertThat; @RequireProvider(SingleUseObjectProvider.class) @@ -178,6 +183,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest { }); // simulate restart + removeRevokedTokenFromRemoteCache(revokedKey); reinitializeKeycloakSessionFactory(); inComittedTransaction(session -> { @@ -188,6 +194,7 @@ public class SingleUseObjectModelTest extends KeycloakModelTest { setTimeOffset(120); // simulate restart + removeRevokedTokenFromRemoteCache(revokedKey); reinitializeKeycloakSessionFactory(); inComittedTransaction(session -> { @@ -285,4 +292,17 @@ public class SingleUseObjectModelTest extends KeycloakModelTest { throw new RuntimeException(e); } } + + private void removeRevokedTokenFromRemoteCache(String revokedKey) { + if (!InfinispanUtils.isRemoteInfinispan()) { + return; + } + inComittedTransaction(session -> { + RemoteCache cache = session.getProvider(InfinispanConnectionProvider.class).getRemoteCache(InfinispanConnectionProvider.ACTION_TOKEN_CACHE); + // remove loaded key to enable preloading from database + cache.remove(InfinispanSingleUseObjectProviderFactory.LOADED); + // remote the token + cache.remove(revokedKey); + }); + } }