Unify expiration handling for SingleUseObjects

Closes #12205

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/realm/HotRodRealmEntity.java

@@ -19,7 +19,6 @@ package org.keycloak.models.map.storage.hotRod.realm;
 
 import org.infinispan.protostream.annotations.ProtoDoc;
 import org.infinispan.protostream.annotations.ProtoField;
-import org.keycloak.common.util.Time;
 import org.keycloak.models.map.annotations.GenerateHotRodEntityImplementation;
 import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.models.map.realm.MapRealmEntity;
@@ -64,6 +63,8 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.keycloak.models.map.common.ExpirationUtils.isExpired;
+
 @GenerateHotRodEntityImplementation(
         implementInterface = "org.keycloak.models.map.realm.MapRealmEntity",
         inherits = "org.keycloak.models.map.storage.hotRod.realm.HotRodRealmEntity.AbstractHotRodRealmEntityDelegate"
@@ -424,8 +425,7 @@ public class HotRodRealmEntity extends AbstractHotRodEntity {
         }
 
         private boolean checkIfExpired(MapClientInitialAccessEntity cia) {
-            return cia.getRemainingCount() < 1 ||
+            return cia.getRemainingCount() < 1 || isExpired(cia, true);
-                    (cia.getExpiration() != null && cia.getExpiration() < Time.currentTimeMillis());
         }
     }
     @Override

model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionProvider.java

@@ -28,7 +28,6 @@ import org.keycloak.models.map.storage.MapKeycloakTransaction;
 import org.keycloak.models.map.storage.MapStorage;
 import org.keycloak.models.map.storage.ModelCriteriaBuilder.Operator;
 import org.keycloak.models.map.storage.criteria.DefaultModelCriteria;
-import org.keycloak.models.utils.SessionExpiration;
 import org.keycloak.sessions.AuthenticationSessionCompoundId;
 import org.keycloak.sessions.AuthenticationSessionProvider;
 import org.keycloak.sessions.RootAuthenticationSessionModel;
@@ -41,6 +40,7 @@ import java.util.function.Function;
 import java.util.function.Predicate;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
+import static org.keycloak.models.map.common.ExpirationUtils.isExpired;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 import static org.keycloak.models.map.storage.criteria.DefaultModelCriteria.criteria;
 import static org.keycloak.models.utils.SessionExpiration.getAuthSessionLifespan;
@@ -65,13 +65,11 @@ public class MapRootAuthenticationSessionProvider implements AuthenticationSessi
 
     private Function<MapRootAuthenticationSessionEntity, RootAuthenticationSessionModel> entityToAdapterFunc(RealmModel realm) {
         return origEntity -> {
-            //return new MapRootAuthenticationSessionAdapter(session, realm, origEntity);
+            if (isExpired(origEntity, true)) {
-            Long expiration = origEntity.getExpiration();
-            if (expiration == null || Time.currentTimeMillis() < origEntity.getExpiration())  {
-                return new MapRootAuthenticationSessionAdapter(session, realm, origEntity);
-            } else {
                 tx.delete(origEntity.getId());
                 return null;
+            } else {
+                return new MapRootAuthenticationSessionAdapter(session, realm, origEntity);
             }
         };
     }

model/map/src/main/java/org/keycloak/models/map/common/ExpirableEntity.java

@@ -34,7 +34,8 @@ public interface ExpirableEntity extends AbstractEntity {
     /**
      * Returns a point in the time (timestamp in milliseconds since The Epoch) when this entity expires.
      *
-     * @return a timestamp in milliseconds since The Epoch or {@code null} if this entity never expires.
+     * @return a timestamp in milliseconds since The Epoch or {@code null} if this entity never expires
+     *         or expiration is not known.
      */
     Long getExpiration();
 

model/map/src/main/java/org/keycloak/models/map/common/ExpirationUtils.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2022 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.models.map.common;
+
+import org.keycloak.common.util.Time;
+
+public class ExpirationUtils {
+
+    /**
+     * Checks whether the {@code entity} is expired
+     *
+     * @param entity to check
+     * @param allowInfiniteValues sets how null values are interpreted, if true entity with expiration equal
+     *                            to {@code null} is interpreted as never expiring entity, if false entities
+     *                            with {@code null} expiration are interpreted as expired entities
+     * @return true if the {@code entity} is expired (expiration time is in the past or now), false otherwise
+     */
+    public static boolean isExpired(ExpirableEntity entity, boolean allowInfiniteValues) {
+        Long expiration = entity.getExpiration();
+        if (!allowInfiniteValues && expiration == null) return false;
+        return expiration != null && expiration <= Time.currentTimeMillis();
+    }
+}

model/map/src/main/java/org/keycloak/models/map/events/MapEventStoreProvider.java

@@ -38,6 +38,7 @@ import java.util.function.Function;
 import java.util.stream.Stream;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
+import static org.keycloak.models.map.common.ExpirationUtils.isExpired;
 import static org.keycloak.models.map.events.EventUtils.modelToEntity;
 
 public class MapEventStoreProvider implements EventStoreProvider {
@@ -79,9 +80,8 @@ public class MapEventStoreProvider implements EventStoreProvider {
     }
 
     private boolean filterExpired(ExpirableEntity event) {
-        Long expiration = event.getExpiration();
         // Check if entity is expired
-        if (expiration != null && expiration <= Time.currentTimeMillis()) {
+        if (isExpired(event, true)) {
             // Remove entity
             authEventsTX.delete(event.getId());
 

model/map/src/main/java/org/keycloak/models/map/realm/MapRealmEntity.java

@@ -17,7 +17,6 @@
 
 package org.keycloak.models.map.realm;
 
-import org.keycloak.common.util.Time;
 import org.keycloak.models.map.annotations.GenerateEntityImplementations;
 import org.keycloak.models.map.annotations.IgnoreForEntityImplementationGenerator;
 import org.keycloak.models.map.common.AbstractEntity;
@@ -43,6 +42,8 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.keycloak.models.map.common.ExpirationUtils.isExpired;
+
 @GenerateEntityImplementations(
         inherits = "org.keycloak.models.map.realm.MapRealmEntity.AbstractRealmEntity"
 )
@@ -237,8 +238,7 @@ public interface MapRealmEntity extends UpdatableEntity, AbstractEntity, EntityW
         }
 
         private boolean checkIfExpired(MapClientInitialAccessEntity cia) {
-            return cia.getRemainingCount() < 1 ||
+            return cia.getRemainingCount() < 1 || isExpired(cia, true);
-                    (cia.getExpiration() != null && cia.getExpiration() < Time.currentTimeMillis());
         }
     }
 

model/map/src/main/java/org/keycloak/models/map/singleUseObject/MapSingleUseObjectAdapter.java

@@ -46,7 +46,7 @@ public class MapSingleUseObjectAdapter extends AbstractSingleUseObjectModel<MapS
     @Override
     public int getExpiration() {
         Long expiration = entity.getExpiration();
-        return expiration != null ? TimeAdapter.fromLongWithTimeInSecondsToIntegerWithTimeInSeconds(expiration) : 0;
+        return expiration != null ? TimeAdapter.fromLongWithTimeInSecondsToIntegerWithTimeInSeconds(TimeAdapter.fromMilliSecondsToSeconds(expiration)) : 0;
     }
 
     @Override

model/map/src/main/java/org/keycloak/models/map/singleUseObject/MapSingleUseObjectProvider.java

@@ -35,6 +35,7 @@ import java.util.Collections;
 import java.util.Map;
 
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
+import static org.keycloak.models.map.common.ExpirationUtils.isExpired;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 import static org.keycloak.models.map.storage.criteria.DefaultModelCriteria.criteria;
 
@@ -55,12 +56,11 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
     }
 
     private ActionTokenValueModel singleUseEntityToAdapter(MapSingleUseObjectEntity origEntity) {
-        long expiration = origEntity.getExpiration() != null ? origEntity.getExpiration() : 0L;
+        if (isExpired(origEntity, false)) {
-        if (Time.currentTime() < expiration)  {
-            return new MapSingleUseObjectAdapter(session, origEntity);
-        } else {
             actionTokenStoreTx.delete(origEntity.getId());
             return null;
+        } else {
+            return new MapSingleUseObjectAdapter(session, origEntity);
         }
     }
 
@@ -75,7 +75,8 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
         DefaultModelCriteria<ActionTokenValueModel> mcb = criteria();
         mcb = mcb.compare(ActionTokenValueModel.SearchableFields.USER_ID, ModelCriteriaBuilder.Operator.EQ, actionTokenKey.getUserId())
                 .compare(ActionTokenValueModel.SearchableFields.ACTION_ID, ModelCriteriaBuilder.Operator.EQ, actionTokenKey.getActionId())
-                .compare(ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE, ModelCriteriaBuilder.Operator.EQ, actionTokenKey.getActionVerificationNonce().toString());
+                .compare(ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE, ModelCriteriaBuilder.Operator.EQ, actionTokenKey.getActionVerificationNonce().toString())
+                .compare(ActionTokenValueModel.SearchableFields.EXPIRATION, ModelCriteriaBuilder.Operator.GT, Time.currentTimeMillis());
 
         ActionTokenValueModel existing = actionTokenStoreTx.read(withCriteria(mcb))
                 .findFirst().map(this::singleUseEntityToAdapter).orElse(null);
@@ -85,7 +86,7 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
             actionTokenStoreEntity.setUserId(actionTokenKey.getUserId());
             actionTokenStoreEntity.setActionId(actionTokenKey.getActionId());
             actionTokenStoreEntity.setActionVerificationNonce(actionTokenKey.getActionVerificationNonce().toString());
-            actionTokenStoreEntity.setExpiration(TimeAdapter.fromIntegerWithTimeInSecondsToLongWithTimeAsInSeconds(actionTokenKey.getExpiration()));
+            actionTokenStoreEntity.setExpiration(TimeAdapter.fromSecondsToMilliseconds(actionTokenKey.getExpiration()));
             actionTokenStoreEntity.setNotes(notes);
 
             LOG.debugf("Adding used action token to actionTokens cache: %s", actionTokenKey.toString());
@@ -105,7 +106,8 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
         DefaultModelCriteria<ActionTokenValueModel> mcb = criteria();
         mcb = mcb.compare(ActionTokenValueModel.SearchableFields.USER_ID, ModelCriteriaBuilder.Operator.EQ, key.getUserId())
                 .compare(ActionTokenValueModel.SearchableFields.ACTION_ID, ModelCriteriaBuilder.Operator.EQ, key.getActionId())
-                .compare(ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE, ModelCriteriaBuilder.Operator.EQ, key.getActionVerificationNonce().toString());
+                .compare(ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE, ModelCriteriaBuilder.Operator.EQ, key.getActionVerificationNonce().toString())
+                .compare(ActionTokenValueModel.SearchableFields.EXPIRATION, ModelCriteriaBuilder.Operator.GT, Time.currentTimeMillis());
 
         return actionTokenStoreTx.read(withCriteria(mcb))
                 .findFirst().map(this::singleUseEntityToAdapter).orElse(null);
@@ -122,7 +124,8 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
         DefaultModelCriteria<ActionTokenValueModel> mcb = criteria();
         mcb = mcb.compare(ActionTokenValueModel.SearchableFields.USER_ID, ModelCriteriaBuilder.Operator.EQ, key.getUserId())
                 .compare(ActionTokenValueModel.SearchableFields.ACTION_ID, ModelCriteriaBuilder.Operator.EQ, key.getActionId())
-                .compare(ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE, ModelCriteriaBuilder.Operator.EQ, key.getActionVerificationNonce().toString());
+                .compare(ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE, ModelCriteriaBuilder.Operator.EQ, key.getActionVerificationNonce().toString())
+                .compare(ActionTokenValueModel.SearchableFields.EXPIRATION, ModelCriteriaBuilder.Operator.GT, Time.currentTimeMillis());
 
         MapSingleUseObjectEntity mapSingleUseObjectEntity = actionTokenStoreTx.read(withCriteria(mcb)).findFirst().orElse(null);
         if (mapSingleUseObjectEntity != null) {
@@ -148,7 +151,7 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
 
         singleUseEntity = new MapSingleUseObjectEntityImpl();
         singleUseEntity.setId(key);
-        singleUseEntity.setExpiration((long) Time.currentTime() + lifespanSeconds);
+        singleUseEntity.setExpiration(Time.currentTimeMillis() + TimeAdapter.fromSecondsToMilliseconds(lifespanSeconds));
         singleUseEntity.setNotes(notes);
 
         actionTokenStoreTx.create(singleUseEntity);
@@ -206,7 +209,7 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
         } else {
             singleUseEntity = new MapSingleUseObjectEntityImpl();
             singleUseEntity.setId(key);
-            singleUseEntity.setExpiration((long) Time.currentTime() + lifespanInSeconds);
+            singleUseEntity.setExpiration(Time.currentTimeMillis() + TimeAdapter.fromSecondsToMilliseconds(lifespanInSeconds));
 
             actionTokenStoreTx.create(singleUseEntity);
 
@@ -231,11 +234,11 @@ public class MapSingleUseObjectProvider implements ActionTokenStoreProvider, Sin
     private MapSingleUseObjectEntity getWithExpiration(String key) {
         MapSingleUseObjectEntity singleUseEntity = actionTokenStoreTx.read(key);
         if (singleUseEntity != null) {
-            long expiration = singleUseEntity.getExpiration() != null ? singleUseEntity.getExpiration() : 0L;
+            if (isExpired(singleUseEntity, false)) {
-            if (Time.currentTime() < expiration) {
+                actionTokenStoreTx.delete(key);
+            } else {
                 return singleUseEntity;
             }
-            actionTokenStoreTx.delete(key);
         }
         return null;
     }

model/map/src/main/java/org/keycloak/models/map/storage/chm/MapFieldPredicates.java

@@ -235,6 +235,7 @@ public class MapFieldPredicates {
         put(ACTION_TOKEN_PREDICATES, ActionTokenValueModel.SearchableFields.USER_ID,                    MapSingleUseObjectEntity::getUserId);
         put(ACTION_TOKEN_PREDICATES, ActionTokenValueModel.SearchableFields.ACTION_ID,                  MapSingleUseObjectEntity::getActionId);
         put(ACTION_TOKEN_PREDICATES, ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE,  MapSingleUseObjectEntity::getActionVerificationNonce);
+        put(ACTION_TOKEN_PREDICATES, ActionTokenValueModel.SearchableFields.EXPIRATION,                 MapSingleUseObjectEntity::getExpiration);
     }
 
     static {

model/map/src/main/java/org/keycloak/models/map/storage/chm/SingleUseObjectConcurrentHashMapStorage.java

@@ -29,6 +29,8 @@ import org.keycloak.models.map.storage.criteria.DefaultModelCriteria;
 
 import java.util.stream.Stream;
 
+import static org.keycloak.models.map.common.ExpirationUtils.isExpired;
+
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
  */
@@ -72,7 +74,8 @@ public class SingleUseObjectConcurrentHashMapStorage<K, V extends AbstractEntity
         SingleUseObjectModelCriteriaBuilder mcb = criteria.flashToModelCriteriaBuilder(createSingleUseObjectCriteriaBuilder());
         if (mcb.isValid()) {
             MapSingleUseObjectEntity value = read(mcb.getKey());
-            return value != null ? Stream.of(value) : Stream.empty();
+            if (value == null || (mcb.checkExpiration() && isExpired(value, false))) return Stream.empty();
+            return Stream.of(value);
         }
 
         return super.read(queryParameters);

model/map/src/main/java/org/keycloak/models/map/storage/chm/SingleUseObjectModelCriteriaBuilder.java

@@ -17,6 +17,7 @@
 
 package org.keycloak.models.map.storage.chm;
 
+import org.keycloak.models.ActionTokenValueModel;
 import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import org.keycloak.storage.SearchableModelField;
 
@@ -31,13 +32,16 @@ public class SingleUseObjectModelCriteriaBuilder implements ModelCriteriaBuilder
 
     private String actionVerificationNonce;
 
+    private Boolean checkExpiration;
+
     public SingleUseObjectModelCriteriaBuilder() {
     }
 
-    public SingleUseObjectModelCriteriaBuilder(String userId, String actionId, String actionVerificationNonce) {
+    public SingleUseObjectModelCriteriaBuilder(String userId, String actionId, String actionVerificationNonce, Boolean checkExpiration) {
         this.userId = userId;
         this.actionId = actionId;
         this.actionVerificationNonce = actionVerificationNonce;
+        this.checkExpiration = checkExpiration;
     }
 
     @Override
@@ -48,8 +52,11 @@ public class SingleUseObjectModelCriteriaBuilder implements ModelCriteriaBuilder
             actionId = value[0].toString();
         } else if (modelField == org.keycloak.models.ActionTokenValueModel.SearchableFields.ACTION_VERIFICATION_NONCE) {
             actionVerificationNonce = value[0].toString();
+        } else if (modelField == ActionTokenValueModel.SearchableFields.EXPIRATION && op == Operator.GT) {
+            checkExpiration = true;
         }
-        return new SingleUseObjectModelCriteriaBuilder(userId, actionId, actionVerificationNonce);
+
+        return new SingleUseObjectModelCriteriaBuilder(userId, actionId, actionVerificationNonce, checkExpiration);
     }
 
     @Override
@@ -57,6 +64,7 @@ public class SingleUseObjectModelCriteriaBuilder implements ModelCriteriaBuilder
         String userId = null;
         String actionId = null;
         String actionVerificationNonce = null;
+        Boolean checkExpiration = null;
 
         for (ModelCriteriaBuilder builder: builders) {
             SingleUseObjectModelCriteriaBuilder suoMcb = (SingleUseObjectModelCriteriaBuilder) builder;
@@ -69,8 +77,11 @@ public class SingleUseObjectModelCriteriaBuilder implements ModelCriteriaBuilder
             if (suoMcb.actionVerificationNonce != null) {
                 actionVerificationNonce = suoMcb.actionVerificationNonce;
             }
+            if (suoMcb.checkExpiration != null) {
+                checkExpiration = suoMcb.checkExpiration;
+            }
         }
-        return new SingleUseObjectModelCriteriaBuilder(userId, actionId, actionVerificationNonce);
+        return new SingleUseObjectModelCriteriaBuilder(userId, actionId, actionVerificationNonce, checkExpiration);
     }
 
     @Override
@@ -87,6 +98,10 @@ public class SingleUseObjectModelCriteriaBuilder implements ModelCriteriaBuilder
         return userId != null && actionId != null && actionVerificationNonce != null;
     }
 
+    public boolean checkExpiration() {
+        return checkExpiration != null && checkExpiration;
+    }
+
     public String getKey() {
         return userId + ":" + actionId + ":" + actionVerificationNonce;
     }

model/map/src/main/java/org/keycloak/models/map/userSession/MapUserSessionProvider.java

@@ -49,6 +49,7 @@ import java.util.stream.Stream;
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
 import static org.keycloak.models.UserSessionModel.CORRESPONDING_SESSION_ID;
 import static org.keycloak.models.UserSessionModel.SessionPersistenceState.TRANSIENT;
+import static org.keycloak.models.map.common.ExpirationUtils.isExpired;
 import static org.keycloak.models.map.storage.QueryParameters.withCriteria;
 import static org.keycloak.models.map.storage.criteria.DefaultModelCriteria.criteria;
 import static org.keycloak.models.map.userSession.SessionExpiration.setClientSessionExpiration;
@@ -87,8 +88,7 @@ public class MapUserSessionProvider implements UserSessionProvider {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
         return (origEntity) -> {
             if (origEntity == null) return null;
-            long expiration = origEntity.getExpiration() != null ? origEntity.getExpiration() : 0L;
+            if (isExpired(origEntity, false)) {
-            if (expiration <= Time.currentTimeMillis()) {
                 if (TRANSIENT == origEntity.getPersistenceState()) {
                     transientUserSessions.remove(origEntity.getId());
                 } else {
@@ -119,8 +119,7 @@ public class MapUserSessionProvider implements UserSessionProvider {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
         return origEntity -> {
             if (origEntity == null) return null;
-            long expiration = origEntity.getExpiration() != null ? origEntity.getExpiration() : 0L;
+            if (isExpired(origEntity, false)) {
-            if (expiration <= Time.currentTimeMillis()) {
                 userSession.removeAuthenticatedClientSessions(Arrays.asList(origEntity.getClientId()));
                 // if a client session is found among transient ones we can skip call to store
                 if (transientClientSessions.remove(origEntity.getId()) == null) {

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

@@ -31,6 +31,7 @@ public interface ActionTokenValueModel {
         public static final SearchableModelField<ActionTokenValueModel> USER_ID                     = new SearchableModelField<>("userId", String.class);
         public static final SearchableModelField<ActionTokenValueModel> ACTION_ID                   = new SearchableModelField<>("actionId", String.class);
         public static final SearchableModelField<ActionTokenValueModel> ACTION_VERIFICATION_NONCE   = new SearchableModelField<>("actionVerificationNonce", String.class);
+        public static final SearchableModelField<ActionTokenValueModel> EXPIRATION                  = new SearchableModelField<>("expiration", Long.class);
     }
     
     /**