File store freeze

* File store: Fix ID determination

* Forbid changing ID (other setters)

* Improve handling of null values

* Support convertible keys in maps

* Fix writing empty values

* Fix updated flag

* Proceed if an object has been deleted in the same tx

* Fix condition

Co-authored-by: Michal Hajas <mhajas@redhat.com>

---------

Co-authored-by: Michal Hajas <mhajas@redhat.com>

model/build-processor/src/main/java/org/keycloak/models/map/processor/GenerateEntityImplementationsProcessor.java

@@ -575,10 +575,18 @@ public class GenerateEntityImplementationsProcessor extends AbstractGenerateEnti
                 pw.println("        return entityFieldDelegate.isUpdated();");
                 pw.println("    }");
 
+                pw.println("    @Override public void markUpdatedFlag() {");
+                pw.println("        entityFieldDelegate.markUpdatedFlag();");
+                pw.println("    }");
+
                 pw.println("    @Override public void clearUpdatedFlag() {");
                 pw.println("        entityFieldDelegate.clearUpdatedFlag();");
                 pw.println("    }");
 
+                pw.println("    @Override public String toString() {");
+                pw.println("        return \"%\" + String.valueOf(entityFieldDelegate);");
+                pw.println("    }");
+
                 getAllAbstractMethods(e)
                   .forEach(ee -> {
                       String originalField = m2field.get(ee);
@@ -701,6 +709,10 @@ public class GenerateEntityImplementationsProcessor extends AbstractGenerateEnti
                 pw.println("        return this.delegateProvider;");
                 pw.println("    }");
 
+                pw.println("    @Override public String toString() {");
+                pw.println("        return \"/\" + String.valueOf(this.delegateProvider);");
+                pw.println("    }");
+
                 getAllAbstractMethods(e)
                   .forEach(ee -> {
                       printMethodHeader(pw, ee);

model/map-file/src/main/java/org/keycloak/models/map/storage/file/FileCrudOperations.java

@@ -25,6 +25,7 @@ import org.keycloak.models.map.common.AbstractEntity;
 import org.keycloak.models.map.common.ExpirationUtils;
 import org.keycloak.models.map.common.HasRealmId;
 import org.keycloak.models.map.common.StringKeyConverter;
+import org.keycloak.models.map.common.StringKeyConverter.StringKey;
 import org.keycloak.models.map.common.UpdatableEntity;
 import org.keycloak.models.map.realm.MapRealmEntity;
 import org.keycloak.models.map.storage.ModelEntityUtil;
@@ -136,7 +137,7 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
 
     // Percent sign + Unix (/) and https://learn.microsoft.com/en-us/windows/win32/fileio/naming-a-file reserved characters
     private static final Pattern RESERVED_CHARACTERS = Pattern.compile("[%<:>\"/\\\\|?*=]");
-    private static final String ID_COMPONENT_SEPARATOR = ":";
+    public static final String ID_COMPONENT_SEPARATOR = ":";
     private static final String ESCAPING_CHARACTER = "=";
     private static final Pattern ID_COMPONENT_SEPARATOR_PATTERN = Pattern.compile(Pattern.quote(ID_COMPONENT_SEPARATOR) + "+");
 
@@ -184,11 +185,11 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
             return null;
         }
 
-        String escapedId = determineKeyFromValue(parsedObject, false);
         final String fileNameStr = fileName.getFileName().toString();
         final String idFromFilename = fileNameStr.substring(0, fileNameStr.length() - FILE_SUFFIX.length());
+        String escapedId = determineKeyFromValue(parsedObject, idFromFilename);
         if (escapedId == null) {
-            LOG.debugf("Determined ID from filename: %s%s", idFromFilename);
+            LOG.tracef("Determined ID from filename: %s%s", idFromFilename);
             escapedId = idFromFilename;
         } else if (!escapedId.endsWith(idFromFilename)) {
             LOG.warnf("Id \"%s\" does not conform with filename \"%s\", expected: %s", escapedId, fileNameStr, escapeId(escapedId));
@@ -210,6 +211,23 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
         return value;
     }
 
+    public String determineKeyFromValue(V value, String lastIdComponentIfUnset) {
+        String[] proposedId = suggestedPath.apply(value);
+
+        if (proposedId == null || proposedId.length == 0) {
+            return lastIdComponentIfUnset;
+        } else if (proposedId[proposedId.length - 1] == null) {
+            proposedId[proposedId.length - 1] = lastIdComponentIfUnset;
+        }
+
+        String[] escapedProposedId = escapeId(proposedId);
+        final String res = String.join(ID_COMPONENT_SEPARATOR, escapedProposedId);
+        if (LOG.isTraceEnabled()) {
+            LOG.tracef("determineKeyFromValue: got %s (%s) for %s", res, res == null ? null : String.join(" [/] ", proposedId), value);
+        }
+        return res;
+    }
+
     /**
      * Returns escaped ID - relative file name in the file system with path separator {@link #ID_COMPONENT_SEPARATOR}.
      *
@@ -218,22 +236,16 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
      * @return
      */
     @Override
-    public String determineKeyFromValue(V value, boolean forCreate) {
+    public String determineKeyFromValue(V value) {
         final boolean randomId;
         String[] proposedId = suggestedPath.apply(value);
 
-        if (!forCreate) {
-            String[] escapedProposedId = escapeId(proposedId);
-            final String res = proposedId == null ? null : String.join(ID_COMPONENT_SEPARATOR, escapedProposedId);
-            if (LOG.isDebugEnabled()) {
-                LOG.debugf("determineKeyFromValue: got %s (%s) for %s", res, res == null ? null : String.join(" [/] ", proposedId), value);
-            }
-            return res;
-        }
-
         if (proposedId == null || proposedId.length == 0) {
             randomId = value.getId() == null;
-            proposedId = new String[]{value.getId() == null ? StringKeyConverter.StringKey.INSTANCE.yieldNewUniqueKey() : value.getId()};
+            proposedId = new String[] { value.getId() == null ? StringKey.INSTANCE.yieldNewUniqueKey() : value.getId() };
+        } else if (proposedId[proposedId.length - 1] == null) {
+            randomId = true;
+            proposedId[proposedId.length - 1] = StringKey.INSTANCE.yieldNewUniqueKey();
         } else {
             randomId = false;
         }
@@ -242,7 +254,7 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
         Path sp = getPathForEscapedId(escapedProposedId);   // sp will never be null
 
         final Path parentDir = sp.getParent();
-        if (!Files.isDirectory(parentDir)) {
+        if (! Files.isDirectory(parentDir)) {
             try {
                 Files.createDirectories(parentDir);
             } catch (IOException ex) {
@@ -253,15 +265,15 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
         for (int counter = 0; counter < 100; counter++) {
             LOG.tracef("Attempting to create file %s", sp, StackUtil.getShortStackTrace());
             try {
-                touch(sp);
                 final String res = String.join(ID_COMPONENT_SEPARATOR, escapedProposedId);
-                LOG.debugf("determineKeyFromValue: got %s for created %s", res, value);
+                touch(res, sp);
+                LOG.tracef("determineKeyFromValue: got %s for created %s", res, value);
                 return res;
             } catch (FileAlreadyExistsException ex) {
-                if (!randomId) {
+                if (! randomId) {
                     throw new ModelDuplicateException("File " + sp + " already exists!");
                 }
-                final String lastComponent = StringKeyConverter.StringKey.INSTANCE.yieldNewUniqueKey();
+                final String lastComponent = StringKey.INSTANCE.yieldNewUniqueKey();
                 escapedProposedId[escapedProposedId.length - 1] = lastComponent;
                 sp = getPathForEscapedId(escapedProposedId);
             } catch (IOException ex) {
@@ -299,7 +311,7 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
 
         // We cannot use Files.find since it throws an UncheckedIOException if it lists a file which is removed concurrently
         // before its BasicAttributes can be retrieved for its BiPredicate parameter
-        try (Stream<Path> dirStream = Files.walk(dataDirectory, entityClass == MapRealmEntity.class ? 1 : 2)) {
+        try (Stream<Path> dirStream = Files.walk(dataDirectory, entityClass == MapRealmEntity.class ? 1 : 3)) {
             // The paths list has to be materialized first, otherwise "dirStream" would be closed
             // before the resulting stream would be read and would return empty result
             paths = dirStream.collect(Collectors.toList());
@@ -396,7 +408,7 @@ public abstract class FileCrudOperations<V extends AbstractEntity & UpdatableEnt
         }
     }
 
-    protected abstract void touch(Path sp) throws IOException;
+    protected abstract void touch(String proposedId, Path sp) throws IOException;
 
     protected abstract boolean removeIfExists(Path sp);
 

model/map-file/src/main/java/org/keycloak/models/map/storage/file/FileMapStorage.java

@@ -37,26 +37,30 @@ import java.nio.file.attribute.BasicFileAttributes;
 import java.nio.file.attribute.FileTime;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 
+import static org.keycloak.models.map.storage.file.FileCrudOperations.ID_COMPONENT_SEPARATOR;
+
 /**
  * {@link MapStorage} implementation used with the file map storage.
  *
  * @author <a href="mailto:sguilhen@redhat.com">Stefan Guilhen</a>
  */
 public class FileMapStorage<V extends AbstractEntity & UpdatableEntity, M>
-  extends ConcurrentHashMapStorage<String, V, M> {
+  extends ConcurrentHashMapStorage<String, V, M, FileCrudOperations<V, M>> {
 
     private static final Logger LOG = Logger.getLogger(FileMapStorage.class);
 
     private final List<Path> createdPaths = new LinkedList<>();
     private final List<Path> pathsToDelete = new LinkedList<>();
-    private final Map<Path, Path> renameOnCommit = new HashMap<>();
+    private final Map<Path, Path> renameOnCommit = new LinkedHashMap<>();
     private final Map<Path, FileTime> lastModified = new HashMap<>();
 
     private final String txId = StringKey.INSTANCE.yieldNewUniqueKey();
@@ -132,7 +136,11 @@ public class FileMapStorage<V extends AbstractEntity & UpdatableEntity, M>
         }
     }
 
-    public void touch(Path path) throws IOException {
+    public void touch(String proposedId, Path path) throws IOException {
+        if (Optional.ofNullable(tasks.get(proposedId)).map(MapTaskWithValue::getOperation).orElse(null) == MapOperation.DELETE) {
+            // If deleted in the current transaction before this operation, then do not touch
+            return;
+        }
         Files.createFile(path);
         createdPaths.add(path);
     }
@@ -144,6 +152,7 @@ public class FileMapStorage<V extends AbstractEntity & UpdatableEntity, M>
     }
 
     void registerRenameOnCommit(Path from, Path to) {
+        pathsToDelete.remove(to);
         this.renameOnCommit.put(from, to);
     }
 
@@ -203,15 +212,15 @@ public class FileMapStorage<V extends AbstractEntity & UpdatableEntity, M>
 
     private static class Crud<V extends AbstractEntity & UpdatableEntity, M> extends FileCrudOperations<V, M> {
 
-        private FileMapStorage store;
+        private FileMapStorage<V, M> store;
 
         public Crud(Class<V> entityClass, Function<String, Path> dataDirectoryFunc, Function<V, String[]> suggestedPath, boolean isExpirableEntity) {
             super(entityClass, dataDirectoryFunc, suggestedPath, isExpirableEntity);
         }
 
         @Override
-        protected void touch(Path sp) throws IOException {
+        protected void touch(String proposedId, Path sp) throws IOException {
-            store.touch(sp);
+            store.touch(proposedId, sp);
         }
 
         @Override
@@ -250,7 +259,46 @@ public class FileMapStorage<V extends AbstractEntity & UpdatableEntity, M>
         public <T, EF extends java.lang.Enum<? extends org.keycloak.models.map.common.EntityField<V>> & org.keycloak.models.map.common.EntityField<V>> void set(EF field, T value) {
             String id = entity.getId();
             super.set(field, value);
-            if (! Objects.equals(id, map.determineKeyFromValue(entity, false))) {
+            checkIdMatches(id, field);
+        }
+
+        @Override
+        public <T, EF extends java.lang.Enum<? extends org.keycloak.models.map.common.EntityField<V>> & org.keycloak.models.map.common.EntityField<V>> void collectionAdd(EF field, T value) {
+            String id = entity.getId();
+            super.collectionAdd(field, value);
+            checkIdMatches(id, field);
+        }
+
+        @Override
+        public <T, EF extends java.lang.Enum<? extends org.keycloak.models.map.common.EntityField<V>> & org.keycloak.models.map.common.EntityField<V>> Object collectionRemove(EF field, T value) {
+            String id = entity.getId();
+            final Object res = super.collectionRemove(field, value);
+            checkIdMatches(id, field);
+            return res;
+        }
+
+        @Override
+        public <K, T, EF extends java.lang.Enum<? extends org.keycloak.models.map.common.EntityField<V>> & org.keycloak.models.map.common.EntityField<V>> void mapPut(EF field, K key, T value) {
+            String id = entity.getId();
+            super.mapPut(field, key, value);
+            checkIdMatches(id, field);
+        }
+
+        @Override
+        public <K, EF extends java.lang.Enum<? extends org.keycloak.models.map.common.EntityField<V>> & org.keycloak.models.map.common.EntityField<V>> Object mapRemove(EF field, K key) {
+            String id = entity.getId();
+            final Object res = super.mapRemove(field, key);
+            checkIdMatches(id, field);
+            return res;
+        }
+
+        private <EF extends java.lang.Enum<? extends org.keycloak.models.map.common.EntityField<V>> & org.keycloak.models.map.common.EntityField<V>> void checkIdMatches(String id, EF field) throws ReadOnlyException {
+            final String idNow = map.determineKeyFromValue(entity, "");
+            if (! Objects.equals(id, idNow)) {
+                if (idNow.endsWith(ID_COMPONENT_SEPARATOR) && id.startsWith(idNow)) {
+                    return;
+                }
+
                 throw new ReadOnlyException("Cannot change " + field + " as that would change primary key");
             }
         }

model/map-file/src/main/java/org/keycloak/models/map/storage/file/FileMapStorageProvider.java

@@ -55,7 +55,7 @@ public class FileMapStorageProvider implements MapStorageProvider {
         return (MapStorage<V, M>) SessionAttributesUtils.createMapStorageIfAbsent(session, getClass(), modelType, factoryId, () -> createFileMapStorage(modelType));
     }
 
-    private <V extends AbstractEntity & UpdatableEntity, M> ConcurrentHashMapStorage<?, V, M> createFileMapStorage(Class<M> modelType) {
+    private <V extends AbstractEntity & UpdatableEntity, M> ConcurrentHashMapStorage<?, V, M, FileCrudOperations<V, M>> createFileMapStorage(Class<M> modelType) {
         String areaName = getModelName(modelType, modelType.getSimpleName());
         final Class<V> et = ModelEntityUtil.getEntityType(modelType);
         Function<V, String[]> uniqueHumanReadableField = (Function<V, String[]>) UNIQUE_HUMAN_READABLE_NAME_FIELD.get(et);

model/map-file/src/main/java/org/keycloak/models/map/storage/file/FileMapStorageProviderFactory.java

@@ -78,8 +78,10 @@ public class FileMapStorageProviderFactory implements AmphibianProviderFactory<M
       // authz
       entry(MapResourceServerEntity.class,  ((Function<MapResourceServerEntity, String[]>) v -> new String[] { v.getClientId() })),
       entry(MapPolicyEntity.class,          ((Function<MapPolicyEntity, String[]>) v -> new String[] { v.getResourceServerId(), v.getName() })),
-      entry(MapPermissionTicketEntity.class,((Function<MapPermissionTicketEntity, String[]>) v -> new String[] { v.getResourceServerId(), v.getId()})),
+      entry(MapPermissionTicketEntity.class,((Function<MapPermissionTicketEntity, String[]>) v -> new String[] { v.getResourceServerId(), null })),
-      entry(MapResourceEntity.class,        ((Function<MapResourceEntity, String[]>) v -> new String[] { v.getResourceServerId(), v.getName() })),
+      entry(MapResourceEntity.class,        ((Function<MapResourceEntity, String[]>) v -> Objects.equals(v.getResourceServerId(), v.getOwner())
+                                                                                          ? new String[] { v.getResourceServerId(), v.getName() }
+                                                                                          : new String[] { v.getResourceServerId(), v.getName(), v.getOwner() })),
       entry(MapScopeEntity.class,           ((Function<MapScopeEntity, String[]>) v -> new String[] { v.getResourceServerId(), v.getName() }))
     );
 

model/map-file/src/main/java/org/keycloak/models/map/storage/file/common/BlockContext.java

@@ -135,7 +135,7 @@ public interface BlockContext<V> {
 
         @Override
         public void writeValue(Object value, WritingMechanism mech) {
-            if (UndefinedValuesUtils.isUndefined(value)) return;
+            if (value == null || (value.getClass() != String.class && UndefinedValuesUtils.isUndefined(value))) return;
             mech.writeObject(value);
         }
 

model/map-file/src/main/java/org/keycloak/models/map/storage/file/common/MapEntityContext.java

@@ -243,6 +243,11 @@ public class MapEntityContext<T> implements BlockContext<T> {
         });
     }
 
+    @Override
+    public String toString() {
+        return "MapEntityContext[" + objectClass.getCanonicalName() + ']';
+    }
+
     public static class MapEntitySequenceYamlContext<T> extends DefaultListContext<T> {
 
         public MapEntitySequenceYamlContext(Class<T> itemClass) {

model/map-file/src/main/java/org/keycloak/models/map/storage/file/yaml/YamlParser.java

@@ -16,6 +16,7 @@
  */
 package org.keycloak.models.map.storage.file.yaml;
 
+import org.keycloak.models.map.common.CastUtils;
 import org.keycloak.models.map.storage.file.common.BlockContextStack;
 import org.keycloak.models.map.storage.file.common.BlockContext.DefaultListContext;
 import org.keycloak.models.map.storage.file.common.BlockContext.DefaultMapContext;
@@ -46,6 +47,9 @@ import org.snakeyaml.engine.v2.resolver.ScalarResolver;
 import org.snakeyaml.engine.v2.scanner.StreamReader;
 import org.keycloak.models.map.storage.file.common.BlockContext;
 
+import java.util.Map;
+import org.snakeyaml.engine.v2.constructor.ConstructScalar;
+import org.snakeyaml.engine.v2.nodes.Node;
 import static org.keycloak.common.util.StackUtil.getShortStackTrace;
 
 /**
@@ -73,15 +77,24 @@ public class YamlParser<E> {
         public Object constructStandardJavaInstance(ScalarNode node) {
             return findConstructorFor(node)
                 .map(constructor -> constructor.construct(node))
-                .orElseThrow(() -> new ConstructorException(null, Optional.empty(), "could not determine a constructor for the tag " + node.getTag(), node.getStartMark()));
+                .orElseThrow(() -> new ConstructorException(null, Optional.empty(), "Could not determine a constructor for the tag " + node.getTag(), node.getStartMark()));
         }
 
         public static final MiniConstructor INSTANCE = new MiniConstructor();
     }
 
+    private static final class NullConstructor extends ConstructScalar {
+
+        @Override
+        public Object construct(Node node) {
+            return null;
+        }
+    }
+
     private static final LoadSettings SETTINGS = LoadSettings.builder()
       .setAllowRecursiveKeys(false)
       .setParseComments(false)
+      .setTagConstructors(Map.of(Tag.NULL, new NullConstructor()))
       .build();
 
     public static <E> E parse(Path path, BlockContext<E> initialContext) {
@@ -175,7 +188,11 @@ public class YamlParser<E> {
             Object key = parseNodeInFreshContext();
             LOG.tracef("Parsed mapping key: %s", key);
             if (! (key instanceof String)) {
-                throw new IllegalStateException("Invalid key in map: " + key);
+                try {
+                    key = CastUtils.cast(key, String.class);
+                } catch (IllegalStateException ex) {
+                    throw new IllegalStateException("Invalid key in map: " + key);
+                }
             }
             Object value = parseNodeInFreshContext((String) key);
             LOG.tracef("Parsed mapping value: %s", value);

model/map-file/src/main/java/org/keycloak/models/map/storage/file/yaml/YamlWritingMechanism.java

@@ -141,8 +141,15 @@ public class YamlWritingMechanism implements WritingMechanism, Closeable {
     }
 
     private ScalarStyle determineStyle(Object value) {
-        if (value instanceof String && ((String) value).lastIndexOf('\n') > 0) {
+        if (value instanceof String) {
-            return ScalarStyle.FOLDED;
+            String sValue = (String) value;
+            // TODO: Check numeric values and quote those as well
+            if ("null".equals(sValue)) {
+                return ScalarStyle.DOUBLE_QUOTED;
+            }
+            if (sValue.length() > 120 || sValue.lastIndexOf('\n') > 0) {
+                return ScalarStyle.FOLDED;
+            }
         }
         return ScalarStyle.PLAIN;
     }

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

@@ -88,7 +88,7 @@ public class HotRodMapStorageProvider implements MapStorageProvider {
         }
     }
 
-    private <K, E extends AbstractHotRodEntity, V extends HotRodEntityDelegate<E> & AbstractEntity, M> ConcurrentHashMapStorage<K, V, M> createHotRodMapStorage(KeycloakSession session, Class<M> modelType, MapStorageProviderFactory.Flag... flags) {
+    private <K, E extends AbstractHotRodEntity, V extends HotRodEntityDelegate<E> & AbstractEntity, M> ConcurrentHashMapStorage<K, V, M, ?> createHotRodMapStorage(KeycloakSession session, Class<M> modelType, MapStorageProviderFactory.Flag... flags) {
         HotRodConnectionProvider connectionProvider = session.getProvider(HotRodConnectionProvider.class);
         HotRodEntityDescriptor<E, V> entityDescriptor = (HotRodEntityDescriptor<E, V>) factory.getEntityDescriptor(modelType);
         Map<SearchableModelField<? super M>, MapModelCriteriaBuilder.UpdatePredicatesFunc<String, V, M>> fieldPredicates = MapFieldPredicates.getPredicates((Class<M>) entityDescriptor.getModelTypeClass());

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/transaction/AllAreasHotRodStoresWrapper.java

@@ -30,10 +30,10 @@ import java.util.function.Supplier;
  */
 public class AllAreasHotRodStoresWrapper extends AbstractKeycloakTransaction {
 
-    private final Map<Class<?>, ConcurrentHashMapStorage<?, ?, ?>> MapKeycloakStoresMap = new ConcurrentHashMap<>();
+    private final Map<Class<?>, ConcurrentHashMapStorage<?, ?, ?, ?>> MapKeycloakStoresMap = new ConcurrentHashMap<>();
 
-    public ConcurrentHashMapStorage<?, ?, ?> getOrCreateStoreForModel(Class<?> modelType, Supplier<ConcurrentHashMapStorage<?, ?, ?>> supplier) {
+    public ConcurrentHashMapStorage<?, ?, ?, ?> getOrCreateStoreForModel(Class<?> modelType, Supplier<ConcurrentHashMapStorage<?, ?, ?, ?>> supplier) {
-        ConcurrentHashMapStorage<?, ?, ?> store = MapKeycloakStoresMap.computeIfAbsent(modelType, t -> supplier.get());
+        ConcurrentHashMapStorage<?, ?, ?, ?> store = MapKeycloakStoresMap.computeIfAbsent(modelType, t -> supplier.get());
         if (!store.isActive()) {
             store.begin();
         }

model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/userSession/HotRodUserSessionMapStorage.java

@@ -25,6 +25,7 @@ import org.keycloak.models.map.common.StringKeyConverter;
 import org.keycloak.models.map.common.delegate.SimpleDelegateProvider;
 import org.keycloak.models.map.storage.QueryParameters;
 import org.keycloak.models.map.storage.CrudOperations;
+import org.keycloak.models.map.storage.chm.ConcurrentHashMapCrudOperations;
 import org.keycloak.models.map.storage.chm.ConcurrentHashMapStorage;
 import org.keycloak.models.map.storage.chm.MapModelCriteriaBuilder;
 import org.keycloak.models.map.storage.criteria.DefaultModelCriteria;
@@ -43,15 +44,15 @@ import java.util.stream.Stream;
 
 import static org.keycloak.models.map.storage.ModelCriteriaBuilder.Operator.IN;
 
-public class HotRodUserSessionMapStorage<K> extends ConcurrentHashMapStorage<K, MapUserSessionEntity, UserSessionModel> {
+public class HotRodUserSessionMapStorage<K> extends ConcurrentHashMapStorage<K, MapUserSessionEntity, UserSessionModel, CrudOperations<MapUserSessionEntity, UserSessionModel>> {
 
-    private final ConcurrentHashMapStorage<String, MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel> clientSessionStore;
+    private final ConcurrentHashMapStorage<String, MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel, CrudOperations<MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel>> clientSessionStore;
 
     public HotRodUserSessionMapStorage(CrudOperations<MapUserSessionEntity, UserSessionModel> map,
                                        StringKeyConverter<K> keyConverter,
                                        DeepCloner cloner,
                                        Map<SearchableModelField<? super UserSessionModel>, MapModelCriteriaBuilder.UpdatePredicatesFunc<K, MapUserSessionEntity, UserSessionModel>> fieldPredicates,
-                                       ConcurrentHashMapStorage<String, MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel> clientSessionStore
+                                       ConcurrentHashMapStorage<String, MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel, CrudOperations<MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel>> clientSessionStore
     ) {
         super(map, keyConverter, cloner, fieldPredicates);
         this.clientSessionStore = clientSessionStore;

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

@@ -33,6 +33,11 @@ public interface UpdatableEntity {
         public void clearUpdatedFlag() {
             this.updated = false;
         }
+
+        @Override
+        public void markUpdatedFlag() {
+            this.updated = true;
+        }
     }
 
     /**
@@ -46,4 +51,10 @@ public interface UpdatableEntity {
      * {@link #isUpdated()} would return {@code false}.
      */
     default void clearUpdatedFlag() { }
+
+    /**
+     * An optional operation setting the updated flag. Right after using this method, the
+     * {@link #isUpdated()} would return {@code true}.
+     */
+    default void markUpdatedFlag() { }
 }

model/map/src/main/java/org/keycloak/models/map/common/delegate/EntityFieldDelegate.java

@@ -53,6 +53,16 @@ public interface EntityFieldDelegate<E> extends UpdatableEntity {
             return entity.isUpdated();
         }
 
+        @Override
+        public void markUpdatedFlag() {
+            entity.markUpdatedFlag();
+        }
+
+        @Override
+        public void clearUpdatedFlag() {
+            entity.clearUpdatedFlag();
+        }
+
         @Override
         public String toString() {
             return "&" + String.valueOf(entity);

model/map/src/main/java/org/keycloak/models/map/storage/CrudOperations.java

@@ -134,7 +134,7 @@ public interface CrudOperations<V extends AbstractEntity & UpdatableEntity, M> {
      * @param value
      * @return
      */
-    default String determineKeyFromValue(V value, boolean forCreate) {
+    default String determineKeyFromValue(V value) {
         return value == null ? null : value.getId();
     }
 }

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

@@ -41,15 +41,18 @@ import org.keycloak.models.map.storage.chm.MapModelCriteriaBuilder.UpdatePredica
 import org.keycloak.models.map.storage.criteria.DefaultModelCriteria;
 import org.keycloak.storage.SearchableModelField;
 import java.util.function.Consumer;
+import java.util.Collection;
+import java.util.Set;
+import java.util.function.BiFunction;
 
-public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEntity, M> implements MapStorage<V, M>, KeycloakTransaction, HasRealmId {
+public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEntity, M, CRUD extends CrudOperations<V, M>> implements MapStorage<V, M>, KeycloakTransaction, HasRealmId {
 
     private final static Logger log = Logger.getLogger(ConcurrentHashMapStorage.class);
 
     protected boolean active;
     protected boolean rollback;
-    protected final Map<String, MapTaskWithValue> tasks = new LinkedHashMap<>();
+    protected final TaskMap tasks = new TaskMap();
-    protected final CrudOperations<V, M> map;
+    protected final CRUD map;
     protected final StringKeyConverter<K> keyConverter;
     protected final DeepCloner cloner;
     protected final Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates;
@@ -57,15 +60,99 @@ public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEnt
     private String realmId;
     private final boolean mapHasRealmId;
 
-    enum MapOperation {
+    protected static final class TaskKey {
+        private final String realmId;
+        private final String key;
+
+        public TaskKey(String realmId, String key) {
+            this.realmId = realmId;
+            this.key = key;
+        }
+
+        private Object getRealmId() {
+            return this.realmId;
+        }
+
+        public String getKey() {
+            return key;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(this.key, this.realmId);
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (this == obj) {
+                return true;
+            }
+            if (obj == null) {
+                return false;
+            }
+            if (getClass() != obj.getClass()) {
+                return false;
+            }
+            final TaskKey other = (TaskKey) obj;
+            return Objects.equals(this.key, other.key) && Objects.equals(this.realmId, other.realmId);
+        }
+
+        @Override
+        public String toString() {
+            return key + " / " + realmId;
+        }
+
+        static TaskKey keyFor(String realmId, String id) {
+            return new TaskKey(realmId, id);
+        }
+    }
+
+    protected class TaskMap {
+
+        private final Map<TaskKey, MapTaskWithValue> map = new LinkedHashMap<>();
+
+        public boolean isEmpty() {
+            return map.isEmpty();
+        }
+
+        public boolean containsKey(String key) {
+            return map.containsKey(TaskKey.keyFor(realmId, key));
+        }
+
+        public MapTaskWithValue get(String key) {
+            return map.get(TaskKey.keyFor(realmId, key));
+        }
+
+        public MapTaskWithValue put(String key, MapTaskWithValue value) {
+            return map.put(TaskKey.keyFor(realmId, key), value);
+        }
+
+        public void clear() {
+            map.clear();
+        }
+
+        public Collection<MapTaskWithValue> values() {
+            return map.values();
+        }
+
+        public Set<Entry<TaskKey, MapTaskWithValue>> entrySet() {
+            return map.entrySet();
+        }
+
+        public MapTaskWithValue merge(String key, MapTaskWithValue value, BiFunction<? super MapTaskWithValue, ? super MapTaskWithValue, ? extends MapTaskWithValue> remappingFunction) {
+            return map.merge(TaskKey.keyFor(realmId, key), value, remappingFunction);
+        }
+    }
+
+    protected enum MapOperation {
         CREATE, UPDATE, DELETE,
     }
 
-    public ConcurrentHashMapStorage(CrudOperations<V, M> map, StringKeyConverter<K> keyConverter, DeepCloner cloner, Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates) {
+    public ConcurrentHashMapStorage(CRUD map, StringKeyConverter<K> keyConverter, DeepCloner cloner, Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates) {
         this(map, keyConverter, cloner, fieldPredicates, null);
     }
 
-    public ConcurrentHashMapStorage(CrudOperations<V, M> map, StringKeyConverter<K> keyConverter, DeepCloner cloner, Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates, EntityField<V> realmIdEntityField) {
+    public ConcurrentHashMapStorage(CRUD map, StringKeyConverter<K> keyConverter, DeepCloner cloner, Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates, EntityField<V> realmIdEntityField) {
         this.map = map;
         this.keyConverter = keyConverter;
         this.cloner = cloner;
@@ -248,7 +335,7 @@ public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEnt
 
     @Override
     public V create(V value) {
-        String key = map.determineKeyFromValue(value, true);
+        String key = map.determineKeyFromValue(value);
         if (key == null) {
             K newKey = keyConverter.yieldNewUniqueKey();
             key = keyConverter.keyToString(newKey);
@@ -295,8 +382,8 @@ public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEnt
         final BulkDeleteOperation bdo = new BulkDeleteOperation(queryParameters);
         Predicate<V> filterForNonDeletedObjects = bdo.getFilterForNonDeletedObjects();
         long res = 0;
-        for (Iterator<Entry<String, MapTaskWithValue>> it = tasks.entrySet().iterator(); it.hasNext();) {
+        for (Iterator<Entry<TaskKey, MapTaskWithValue>> it = tasks.entrySet().iterator(); it.hasNext();) {
-            Entry<String, MapTaskWithValue> me = it.next();
+            Entry<TaskKey, MapTaskWithValue> me = it.next();
             if (! filterForNonDeletedObjects.test(me.getValue().getValue())) {
                 log.tracef(" [DELETE_BULK] removing %s", me.getKey());
                 it.remove();
@@ -328,7 +415,7 @@ public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEnt
 
     private Stream<V> createdValuesStream(Predicate<? super K> keyFilter, Predicate<? super V> entityFilter) {
         return this.tasks.entrySet().stream()
-          .filter(me -> keyFilter.test(keyConverter.fromStringSafe(me.getKey())))
+          .filter(me -> Objects.equals(realmId, me.getKey().getRealmId()) && keyFilter.test(keyConverter.fromStringSafe(me.getKey().getKey())))
           .map(Map.Entry::getValue)
           .filter(v -> v.containsCreate() && ! v.isReplace())
           .map(MapTaskWithValue::getValue)

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

@@ -58,7 +58,7 @@ public class ConcurrentHashMapStorageProvider implements MapStorageProvider {
         });
     }
 
-    private <V extends AbstractEntity & UpdatableEntity, M> ConcurrentHashMapStorage getMapStorage(Class<?> modelType, CrudOperations<V, M> crud) {
+    private <K, V extends AbstractEntity & UpdatableEntity, M> ConcurrentHashMapStorage getMapStorage(Class<?> modelType, ConcurrentHashMapCrudOperations<K, V, M> crud) {
         if (modelType == SingleUseObjectValueModel.class) {
             return new SingleUseObjectMapStorage(crud, factory.getKeyConverter(modelType), CLONER, MapFieldPredicates.getPredicates(modelType));
         }

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

@@ -236,14 +236,14 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
     }
 
     @SuppressWarnings("unchecked")
-    private <V extends AbstractEntity & UpdatableEntity, M> CrudOperations<V, M> loadMap(String mapName,
+    private <K, V extends AbstractEntity & UpdatableEntity, M> ConcurrentHashMapCrudOperations<K, V, M> loadMap(String mapName,
                                                                                             Class<M> modelType,
                                                                                             EnumSet<Flag> flags) {
-        final StringKeyConverter kc = keyConverters.getOrDefault(mapName, defaultKeyConverter);
+        final StringKeyConverter<K> kc = keyConverters.getOrDefault(mapName, defaultKeyConverter);
         Class<?> valueType = ModelEntityUtil.getEntityType(modelType);
         LOG.debugf("Initializing new map storage: %s", mapName);
 
-        CrudOperations<V, M> store;
+        ConcurrentHashMapCrudOperations<K, V, M> store;
         if(modelType == SingleUseObjectValueModel.class) {
             store = new SingleUseObjectConcurrentHashMapCrudOperations(kc, CLONER) {
                 @Override
@@ -252,7 +252,7 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
                 }
             };
         } else {
-            store = new ConcurrentHashMapCrudOperations(modelType, kc, CLONER) {
+            store = new ConcurrentHashMapCrudOperations<>(modelType, kc, CLONER) {
                 @Override
                 public String toString() {
                     return "ConcurrentHashMapStorage(" + mapName + suffix + ")";
@@ -288,7 +288,7 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
     }
 
     @SuppressWarnings("unchecked")
-    public <K, V extends AbstractEntity & UpdatableEntity, M> CrudOperations<V, M> getStorage(
+    public <K, V extends AbstractEntity & UpdatableEntity, M> ConcurrentHashMapCrudOperations<K, V, M> getStorage(
       Class<M> modelType, Flag... flags) {
         EnumSet<Flag> f = flags == null || flags.length == 0 ? EnumSet.noneOf(Flag.class) : EnumSet.of(flags[0], flags);
         String name = getModelName(modelType, modelType.getSimpleName());
@@ -297,7 +297,7 @@ public class ConcurrentHashMapStorageProviderFactory implements AmphibianProvide
          *   "... the computation [...] must not attempt to update any other mappings of this map."
          */
 
-        return (CrudOperations<V, M>) storages.computeIfAbsent(name, n -> loadMap(name, modelType, f));
+        return (ConcurrentHashMapCrudOperations<K, V, M>) storages.computeIfAbsent(name, n -> loadMap(name, modelType, f));
     }
 
     public StringKeyConverter<?> getKeyConverter(Class<?> modelType) {

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

@@ -29,7 +29,7 @@ import java.util.Map;
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
  */
-public class SingleUseObjectMapStorage<K> extends ConcurrentHashMapStorage<K, MapSingleUseObjectEntity, SingleUseObjectValueModel> {
+public class SingleUseObjectMapStorage<K> extends ConcurrentHashMapStorage<K, MapSingleUseObjectEntity, SingleUseObjectValueModel, CrudOperations<MapSingleUseObjectEntity, SingleUseObjectValueModel>> {
 
     public SingleUseObjectMapStorage(CrudOperations<MapSingleUseObjectEntity, SingleUseObjectValueModel> map,
                                      StringKeyConverter<K> keyConverter,

model/map/src/main/java/org/keycloak/models/map/user/MapUserEntity.java

@@ -119,7 +119,7 @@ public interface MapUserEntity extends UpdatableEntity, AbstractEntity, EntityWi
             int indexToRemove = toMoveIndex < ourCredentialIndex ? ourCredentialIndex + 1 : ourCredentialIndex;
             credentialsList.remove(indexToRemove);
 
-            this.updated = true;
+            markUpdatedFlag();
             return true;
         }
     }

testsuite/model/src/test/java/org/keycloak/testsuite/model/ConcurrentHashMapStorageTest.java

@@ -86,9 +86,9 @@ public class ConcurrentHashMapStorageTest extends KeycloakModelTest {
         String component2Id = createMapStorageComponent("component2", "keyType", "string");
 
         String[] ids = withRealm(realmId, (session, realm) -> {
-            ConcurrentHashMapStorage<K, MapClientEntity, ClientModel> storageMain = (ConcurrentHashMapStorage<K, MapClientEntity, ClientModel>) (MapStorage) session.getProvider(MapStorageProvider.class, ConcurrentHashMapStorageProviderFactory.PROVIDER_ID).getMapStorage(ClientModel.class);
+            ConcurrentHashMapStorage<K, MapClientEntity, ClientModel, ?> storageMain = (ConcurrentHashMapStorage<K, MapClientEntity, ClientModel, ?>) (MapStorage) session.getProvider(MapStorageProvider.class, ConcurrentHashMapStorageProviderFactory.PROVIDER_ID).getMapStorage(ClientModel.class);
-            ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel> storage1 = (ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component1Id).getMapStorage(ClientModel.class);
+            ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel, ?> storage1 = (ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel, ?>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component1Id).getMapStorage(ClientModel.class);
-            ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel> storage2 = (ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component2Id).getMapStorage(ClientModel.class);
+            ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel, ?> storage2 = (ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel, ?>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component2Id).getMapStorage(ClientModel.class);
 
             // Assert that the map storage can be used both as a standalone store and a component
             assertThat(storageMain, notNullValue());
@@ -157,7 +157,7 @@ public class ConcurrentHashMapStorageTest extends KeycloakModelTest {
         assertClientsPersisted(component1Id, component2Id, idMain, id1, id2);
     }
 
-    private <K,K1> void assertClientDoesNotExist(ConcurrentHashMapStorage<K, MapClientEntity, ClientModel> storage, String id, final StringKeyConverter<K1> kc, final StringKeyConverter<K> kcStorage) {
+    private <K,K1> void assertClientDoesNotExist(ConcurrentHashMapStorage<K, MapClientEntity, ClientModel, ?> storage, String id, final StringKeyConverter<K1> kc, final StringKeyConverter<K> kcStorage) {
         // Assert that the other stores do not contain the to-be-created clients (if they use compatible key format)
         try {
             assertThat(storage.read(id), nullValue());
@@ -170,11 +170,11 @@ public class ConcurrentHashMapStorageTest extends KeycloakModelTest {
         // Check that in the next transaction, the objects are still there
         withRealm(realmId, (session, realm) -> {
             @SuppressWarnings("unchecked")
-            ConcurrentHashMapStorage<K, MapClientEntity, ClientModel> storageMain = (ConcurrentHashMapStorage<K, MapClientEntity, ClientModel>) (MapStorage) session.getProvider(MapStorageProvider.class, ConcurrentHashMapStorageProviderFactory.PROVIDER_ID).getMapStorage(ClientModel.class);
+            ConcurrentHashMapStorage<K, MapClientEntity, ClientModel, ?> storageMain = (ConcurrentHashMapStorage<K, MapClientEntity, ClientModel, ?>) (MapStorage) session.getProvider(MapStorageProvider.class, ConcurrentHashMapStorageProviderFactory.PROVIDER_ID).getMapStorage(ClientModel.class);
             @SuppressWarnings("unchecked")
-            ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel> storage1 = (ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component1Id).getMapStorage(ClientModel.class);
+            ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel, ?> storage1 = (ConcurrentHashMapStorage<K1, MapClientEntity, ClientModel, ?>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component1Id).getMapStorage(ClientModel.class);
             @SuppressWarnings("unchecked")
-            ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel> storage2 = (ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component2Id).getMapStorage(ClientModel.class);
+            ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel, ?> storage2 = (ConcurrentHashMapStorage<K2, MapClientEntity, ClientModel, ?>) (MapStorage) session.getComponentProvider(MapStorageProvider.class, component2Id).getMapStorage(ClientModel.class);
 
             final StringKeyConverter<K> kcMain = (StringKeyConverter<K>) StringKeyConverter.UUIDKey.INSTANCE;
             final StringKeyConverter<K1> kc1 = (StringKeyConverter<K1>) StringKeyConverter.ULongKey.INSTANCE;