IGNITE-28963 Added support for EnumSet collection type during message serialization

modules/codegen/src/main/java/org/apache/ignite/internal/MessageSerializerGenerator.java

@@ -23,6 +23,7 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Collection;
+import java.util.EnumSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -86,6 +87,9 @@ public class MessageSerializerGenerator extends MessageCompanionGenerator {
     /** */
     private static final String MESSAGE_COLLECTION_TYPE_CLS = "org.apache.ignite.plugin.extensions.communication.MessageCollectionType";
 
+    /** */
+    private static final String COLLECTION_IMPL_TYPE_CLS = "org.apache.ignite.plugin.extensions.communication.CollectionImplementationType";
+
     /** */
     private static final String IGNITE_UUID_CLS = "org.apache.ignite.lang.IgniteUuid";
 
@@ -552,9 +556,17 @@ else if (assignableFrom(erasedType(type), type(Collection.class.getName()))) {
 
             assert typeArgs.size() == 1 : type.toString();
 
-            return "new MessageCollectionType(" +
-                messageCollectionItemTypeDescriptor(typeArgs.get(0), field) + ", " +
-                assignableFrom(erasedType(type), type(Set.class.getName())) + ")";
+            TypeMirror itemType = typeArgs.get(0);
+
+            if (EnumSet.class.getName().equals(qualifiedClassName(erasedType(type))) && !enumType(env, itemType))
+                throw new IllegalArgumentException("Unexpected Enum Set element type [itemType=" + itemType + ", colType=" + type + ']');
+
+            imports.add(COLLECTION_IMPL_TYPE_CLS);
+
+            String implType = resolveCollectionImplementationType(type);
+
+            return "new MessageCollectionType(" + messageCollectionItemTypeDescriptor(itemType, field) +
+                ", CollectionImplementationType." + implType + ")";
         }
         else if (enumType(env, type)) {
             imports.add("org.apache.ignite.plugin.extensions.communication.MessageEnumType");
@@ -566,7 +578,7 @@ else if (enumType(env, type)) {
             String decoder = custMapper ? prefix + "Mapper::decode" :
                 "b -> DefaultEnumMapper.INSTANCE.decode(" + prefix + "Vals, b)";
 
-            return String.format("new MessageEnumType<>(%s, %s)", encoder, decoder);
+            return String.format("new MessageEnumType<>(%s.class, %s, %s)", simpleClassName(type), encoder, decoder);
         }
         else {
             imports.add(MESSAGE_ITEM_TYPE_CLS);
@@ -575,6 +587,21 @@ else if (enumType(env, type)) {
         }
     }
 
+    /**
+     * @param type Declared collection type.
+     * @return Name of the {@code CollectionImplementationType} that will be used to create collection instance.
+     */
+    private String resolveCollectionImplementationType(TypeMirror type) {
+        TypeMirror declType = erasedType(type);
+
+        if (EnumSet.class.getName().equals(qualifiedClassName(declType)))
+            return "ENUM_SET";
+        else if (assignableFrom(declType, type(Set.class.getName())))
+            return "HASH_SET";
+        else
+            return "ARRAY_LIST";
+    }
+
     /**
      * Find MessageCollectionItemType for a given type.
      * <p>

modules/core/src/main/java/org/apache/ignite/internal/direct/stream/DirectByteBufferStream.java

@@ -22,6 +22,7 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Collection;
+import java.util.EnumSet;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
@@ -1639,11 +1640,11 @@ public <T> T[] readObjectArray(MessageArrayType type, MessageReader reader) {
     }
 
     /**
-     * Reads collection eather as a {@link ArrayList} or a {@link HashSet}.
+     * Reads collection either as an {@link ArrayList}, a {@link HashSet} or an {@link EnumSet}.
      *
      * @param type Item type.
      * @param reader Reader.
-     * @return {@link ArrayList} or a {@link HashSet}.
+     * @return {@link ArrayList}, {@link HashSet} or {@link EnumSet}.
      */
     public <C extends Collection<?>> C readCollection(MessageCollectionType type, MessageReader reader) {
         if (readSize == -1) {
@@ -1657,7 +1658,7 @@ public <C extends Collection<?>> C readCollection(MessageCollectionType type, Me
 
         if (readSize >= 0) {
             if (col == null)
-                col = type.set() ? U.newHashSet(readSize) : new ArrayList<>(readSize);
+                col = newCollection(type);
 
             for (int i = readItems; i < readSize; i++) {
                 Object item = read(type.valueType(), reader);
@@ -1682,6 +1683,16 @@ public <C extends Collection<?>> C readCollection(MessageCollectionType type, Me
         return col0;
     }
 
+    /** */
+    @SuppressWarnings("unchecked")
+    private Collection<Object> newCollection(MessageCollectionType type) {
+        return switch (type.collectionImplementationType()) {
+            case ENUM_SET -> (Collection<Object>)((MessageEnumType<?>)type.valueType()).newEnumSet();
+            case HASH_SET -> U.newHashSet(readSize);
+            case ARRAY_LIST -> new ArrayList<>(readSize);
+        };
+    }
+
     /**
      * @param type Value type.
      * @param reader Reader.

modules/core/src/main/java/org/apache/ignite/internal/processors/security/SecurityUtils.java

@@ -132,42 +132,30 @@ public static void restoreDefaultSerializeVersion() {
     /**
      * @return Allow all service permissions.
      */
-    public static Map<String, Collection<SecurityPermission>> compatibleServicePermissions() {
+    public static Map<String, EnumSet<SecurityPermission>> compatibleServicePermissions() {
         Map<String, EnumSet<SecurityPermission>> srvcPerms = new HashMap<>();
 
         srvcPerms.put("*", EnumSet.of(
             SecurityPermission.SERVICE_CANCEL,
             SecurityPermission.SERVICE_DEPLOY,
             SecurityPermission.SERVICE_INVOKE));
 
-        return upcast(srvcPerms);
+        return srvcPerms;
     }
 
     /** */
-    @SuppressWarnings("rawtypes")
-    public static Map<String, Collection<SecurityPermission>> upcast(Map<String, EnumSet<SecurityPermission>> map) {
-        return (Map<String, Collection<SecurityPermission>>)(Map)map;
-    }
-
-    /** */
-    @SuppressWarnings("rawtypes")
-    public static Map<String, EnumSet<SecurityPermission>> downcast(Map<String, Collection<SecurityPermission>> map) {
-        return (Map<String, EnumSet<SecurityPermission>>)(Map)map;
-    }
-
-    /**
-     * @param permissionsMap Permissions map.
-     * @return Map with enum sets of security permissions.
-     */
-    public static Map<String, Collection<SecurityPermission>> normalizeValueType(
-        Map<String, Collection<SecurityPermission>> permissionsMap
+    public static Map<String, EnumSet<SecurityPermission>> normalizeResourcePermissions(
+        Map<String, ? extends Collection<SecurityPermission>> rsrcPerms
     ) {
-        return permissionsMap.entrySet().stream()
-            .collect(Collectors.toMap(Map.Entry::getKey, e -> copySafe(e.getValue())));
+        if (rsrcPerms == null)
+            return new HashMap<>();
+
+        return rsrcPerms.entrySet().stream()
+            .collect(Collectors.toMap(Map.Entry::getKey, e -> toEnumSet(e.getValue())));
     }
 
     /** */
-    public static EnumSet<SecurityPermission> copySafe(Collection<SecurityPermission> col) {
+    public static EnumSet<SecurityPermission> toEnumSet(Collection<SecurityPermission> col) {
         if (col instanceof EnumSet<SecurityPermission> enumSet)
             return enumSet;
 

modules/core/src/main/java/org/apache/ignite/plugin/security/SecurityBasicPermissionSet.java

@@ -20,57 +20,63 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.io.ObjectStreamField;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import org.apache.ignite.internal.Order;
-import org.apache.ignite.internal.SelfMarshallingMessage;
 import org.apache.ignite.internal.util.tostring.GridToStringInclude;
 import org.apache.ignite.internal.util.typedef.internal.A;
 import org.apache.ignite.internal.util.typedef.internal.S;
 import org.apache.ignite.internal.util.typedef.internal.U;
 import org.jetbrains.annotations.Nullable;
 
 import static org.apache.ignite.internal.processors.security.SecurityUtils.compatibleServicePermissions;
-import static org.apache.ignite.internal.processors.security.SecurityUtils.copySafe;
-import static org.apache.ignite.internal.processors.security.SecurityUtils.downcast;
 import static org.apache.ignite.internal.processors.security.SecurityUtils.isSecurityCompatibilityMode;
-import static org.apache.ignite.internal.processors.security.SecurityUtils.normalizeValueType;
+import static org.apache.ignite.internal.processors.security.SecurityUtils.normalizeResourcePermissions;
 import static org.apache.ignite.internal.processors.security.SecurityUtils.serializeVersion;
-import static org.apache.ignite.internal.processors.security.SecurityUtils.upcast;
+import static org.apache.ignite.internal.processors.security.SecurityUtils.toEnumSet;
 
 /**
  * Simple implementation of {@link SecurityPermissionSet} interface.
  * Provides convenient way to specify permission set in the XML configuration.
  */
-public class SecurityBasicPermissionSet implements SecurityPermissionSet, SelfMarshallingMessage {
+public class SecurityBasicPermissionSet implements SecurityPermissionSet {
     /** Serial version uid. */
     private static final long serialVersionUID = 0L;
 
+    /** */
+    private static final ObjectStreamField[] serialPersistentFields = {
+        new ObjectStreamField("dfltAllowAll", boolean.class),
+        new ObjectStreamField("cachePermissions", Map.class),
+        new ObjectStreamField("sysPermissions", Collection.class),
+        new ObjectStreamField("taskPermissions", Map.class)
+    };
+
     /** Cache permissions. */
     @GridToStringInclude
     @Order(0)
-    Map<String, Collection<SecurityPermission>> cachePermissions = new HashMap<>();
+    Map<String, EnumSet<SecurityPermission>> cachePermissions = new HashMap<>();
 
     /** Task permissions. */
     @GridToStringInclude
     @Order(1)
-    Map<String, Collection<SecurityPermission>> taskPermissions = new HashMap<>();
+    Map<String, EnumSet<SecurityPermission>> taskPermissions = new HashMap<>();
 
     /** Service permissions. */
     @GridToStringInclude
     @Order(2)
-    transient Map<String, Collection<SecurityPermission>> srvcPermissions = isSecurityCompatibilityMode()
-            ? compatibleServicePermissions()
-            : new HashMap<>();
+    transient Map<String, EnumSet<SecurityPermission>> srvcPermissions = isSecurityCompatibilityMode()
+        ? compatibleServicePermissions()
+        : new HashMap<>();
 
     /** System permissions. */
     @GridToStringInclude
     @Order(3)
-    @Nullable Collection<SecurityPermission> sysPermissions;
+    @Nullable EnumSet<SecurityPermission> sysPermissions;
 
     /** Default allow all. */
     @Order(4)
@@ -82,9 +88,7 @@ public class SecurityBasicPermissionSet implements SecurityPermissionSet, SelfMa
      * @param cachePermissions Cache permissions.
      */
     public void setCachePermissions(Map<String, EnumSet<SecurityPermission>> cachePermissions) {
-        A.notNull(cachePermissions, "cachePermissions");
-
-        this.cachePermissions = normalizeValueType(upcast(cachePermissions));
+        this.cachePermissions = checkPermissions(cachePermissions, "cachePermissions");
     }
 
     /**
@@ -93,9 +97,7 @@ public void setCachePermissions(Map<String, EnumSet<SecurityPermission>> cachePe
      * @param taskPermissions Task permissions.
      */
     public void setTaskPermissions(Map<String, EnumSet<SecurityPermission>> taskPermissions) {
-        A.notNull(taskPermissions, "taskPermissions");
-
-        this.taskPermissions = normalizeValueType(upcast(taskPermissions));
+        this.taskPermissions = checkPermissions(taskPermissions, "taskPermissions");
     }
 
     /**
@@ -104,9 +106,7 @@ public void setTaskPermissions(Map<String, EnumSet<SecurityPermission>> taskPerm
      * @param srvcPermissions Service permissions.
      */
     public void setServicePermissions(Map<String, EnumSet<SecurityPermission>> srvcPermissions) {
-        A.notNull(srvcPermissions, "servicePermissions");
-
-        this.srvcPermissions = normalizeValueType(upcast(srvcPermissions));
+        this.srvcPermissions = checkPermissions(srvcPermissions, "servicePermissions");
     }
 
     /**
@@ -129,22 +129,22 @@ public void setDefaultAllowAll(boolean dfltAllowAll) {
 
     /** {@inheritDoc} */
     @Override public Map<String, EnumSet<SecurityPermission>> cachePermissions() {
-        return downcast(cachePermissions);
+        return cachePermissions;
     }
 
     /** {@inheritDoc} */
     @Override public Map<String, EnumSet<SecurityPermission>> taskPermissions() {
-        return downcast(taskPermissions);
+        return taskPermissions;
     }
 
     /** {@inheritDoc} */
     @Override public Map<String, EnumSet<SecurityPermission>> servicePermissions() {
-        return downcast(srvcPermissions);
+        return srvcPermissions;
     }
 
     /** {@inheritDoc} */
     @Nullable @Override public EnumSet<SecurityPermission> systemPermissions() {
-        return (EnumSet<SecurityPermission>)sysPermissions;
+        return sysPermissions;
     }
 
     /** {@inheritDoc} */
@@ -183,52 +183,64 @@ public void setDefaultAllowAll(boolean dfltAllowAll) {
 
     /** */
     private void writeObject(ObjectOutputStream out) throws IOException {
-        out.defaultWriteObject();
+        ObjectOutputStream.PutField fields = out.putFields();
+
+        fields.put("dfltAllowAll", dfltAllowAll);
+        fields.put("cachePermissions", cachePermissions);
+        fields.put("sysPermissions", sysPermissions);
+        fields.put("taskPermissions", taskPermissions);
+
+        out.writeFields();
 
         if (serializeVersion() >= 2)
             U.writeMap(out, srvcPermissions);
     }
 
     /** */
+    @SuppressWarnings("unchecked")
     private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
-        in.defaultReadObject();
+        ObjectInputStream.GetField fields = in.readFields();
 
-        if (serializeVersion() >= 2)
-            srvcPermissions = U.readMap(in);
+        dfltAllowAll = fields.get("dfltAllowAll", false);
+        cachePermissions = readPermissions(fields, "cachePermissions");
+        taskPermissions = readPermissions(fields, "taskPermissions");
 
-        if (srvcPermissions == null) {
+        Collection<SecurityPermission> sysPerms = (Collection<SecurityPermission>)fields.get("sysPermissions", null);
+
+        sysPermissions = sysPerms == null ? null : toEnumSet(sysPerms);
+
+        Map<String, ? extends Collection<SecurityPermission>> srvcPerms = serializeVersion() >= 2 ? U.readMap(in) : null;
+
+        if (srvcPerms == null) {
             // Allow all for compatibility mode
-            if (serializeVersion() < 2)
-                srvcPermissions = compatibleServicePermissions();
-            else
-                srvcPermissions = Collections.emptyMap();
+            srvcPerms = serializeVersion() < 2 ? compatibleServicePermissions() : Collections.emptyMap();
         }
 
-        normalize();
+        srvcPermissions = normalizeResourcePermissions(srvcPerms);
     }
 
     /** {@inheritDoc} */
     @Override public String toString() {
         return S.toString(SecurityBasicPermissionSet.class, this);
     }
 
-    /** {@inheritDoc} */
-    @Override public void selfMarshal() {
-        // No-op.
-    }
-
-    /** {@inheritDoc} */
-    @Override public void selfUnmarshal() {
-        // Message framework uses ArrayList for ordinary collections,
-        // so we need to convert it to appropriate form explicitly.
-        normalize();
+    /** */
+    @SuppressWarnings("unchecked")
+    private static Map<String, EnumSet<SecurityPermission>> readPermissions(
+        ObjectInputStream.GetField fields,
+        String name
+    ) throws IOException {
+        return normalizeResourcePermissions((Map<String, ? extends Collection<SecurityPermission>>)fields.get(name, null));
     }
 
     /** */
-    private void normalize() {
-        cachePermissions = normalizeValueType(cachePermissions);
-        taskPermissions = normalizeValueType(taskPermissions);
-        srvcPermissions = normalizeValueType(srvcPermissions);
-        sysPermissions = sysPermissions == null ? null : copySafe(sysPermissions);
+    private static Map<String, EnumSet<SecurityPermission>> checkPermissions(
+        Map<String, EnumSet<SecurityPermission>> perms,
+        String name
+    ) {
+        A.notNull(perms, name);
+        A.ensure(perms.values().stream().noneMatch(Objects::isNull), name + " must not contain a null permission set");
+
+        return perms;
     }
 }