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 <pruivo@redhat.com>

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 static org.keycloak.storage.datastore.DefaultDatastoreProviderFactory.setupClearExpiredRevokedTokensScheduledTask;
-import org.keycloak.timer.TimerProvider;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -58,7 +55,8 @@ import org.keycloak.timer.TimerProvider;
 public class InfinispanSingleUseObjectProviderFactory implements SingleUseObjectProviderFactory<InfinispanSingleUseObjectProvider>, 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() {
+            factory.register(event -> {
-                public void onEvent(ProviderEvent event) {
+                if (event instanceof PostMigrationEvent pme) {
-                    if (event instanceof PostMigrationEvent) {
+                    KeycloakSessionFactory sessionFactory = pme.getFactory();
-                        KeycloakSessionFactory sessionFactory = ((PostMigrationEvent) event).getFactory();
+                    setupClearExpiredRevokedTokensScheduledTask(sessionFactory);
-                        try (KeycloakSession session = sessionFactory.create()) {
+                    try (KeycloakSession session = sessionFactory.create()) {
-                            TimerProvider timer = session.getProvider(TimerProvider.class);
+                        // load sessions during startup, not on first request to avoid congestion
-                            if (timer != null) {
+                        initialize(session);
-                                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);
-                        }
                     }
                 }
-
-                private void scheduleTask(KeycloakSessionFactory sessionFactory, TimerProvider timer, long interval) {
-                    timer.schedule(new ClusterAwareScheduledTaskRunner(sessionFactory, new ClearExpiredRevokedTokens(), interval), interval);
-                }
             });
         }
     }

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<String, String> 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<String, String> unwrap(SingleUseObjectValueEntity entity) {
         return entity == null ? null : entity.getNotes();
     }
@@ -100,4 +115,8 @@ public class RemoteInfinispanSingleUseObjectProvider implements SingleUseObjectP
     private static SingleUseObjectValueEntity wrap(Map<String, String> notes) {
         return new SingleUseObjectValueEntity(notes);
     }
+
+    public interface RevokeTokenConsumer {
+        void onTokenRevoke(String token, long lifespanSeconds);
+    }
 }

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<RemoteInfinispanSingleUseObjectProvider>, EnvironmentDependentProviderFactory {
+public class RemoteInfinispanSingleUseObjectProviderFactory implements SingleUseObjectProviderFactory<RemoteInfinispanSingleUseObjectProvider>, 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<String, SingleUseObjectValueEntity> 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<String, String> getOperationalInfo() {
+        Map<String, String> info = new HashMap<>();
+        info.put(CONFIG_PERSIST_REVOKED_TOKENS, Boolean.toString(persistRevokedTokens));
+        return info;
+    }
+
+    @Override
+    public List<ProviderConfigProperty> 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));
+    }
+
+
 }

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<ScheduledTask> getScheduledTasks() {
+    protected static List<ScheduledTask> 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;
+    }
+
 }

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<String, SingleUseObjectValueEntity> 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);
+        });
+    }
 }