IGNITE-28010 Use MessageSerializer for SecurityAwareCustomMessageWrapper

modules/core/src/main/java/org/apache/ignite/internal/managers/discovery/DiscoveryMessageFactory.java

@@ -126,5 +126,6 @@ public class DiscoveryMessageFactory implements MessageFactoryProvider {
 
         // DiscoveryCustomMessage
         factory.register((short)500, CacheStatisticsModeChangeMessage::new, new CacheStatisticsModeChangeMessageSerializer());
+        factory.register((short)501, SecurityAwareCustomMessageWrapper::new, new SecurityAwareCustomMessageWrapperSerializer());
     }
 }

modules/core/src/main/java/org/apache/ignite/internal/managers/discovery/SecurityAwareCustomMessageWrapper.java

@@ -18,24 +18,52 @@
 package org.apache.ignite.internal.managers.discovery;
 
 import java.util.UUID;
+import org.apache.ignite.IgniteCheckedException;
+import org.apache.ignite.IgniteException;
+import org.apache.ignite.internal.Order;
+import org.apache.ignite.internal.util.typedef.F;
+import org.apache.ignite.internal.util.typedef.internal.U;
+import org.apache.ignite.plugin.extensions.communication.Message;
+import org.apache.ignite.plugin.extensions.communication.MessageFactory;
 import org.apache.ignite.spi.discovery.DiscoverySpiCustomMessage;
 import org.jetbrains.annotations.Nullable;
 
+import static org.apache.ignite.marshaller.Marshallers.jdk;
+
 /** Custom message wrapper with ID of security subject that initiated the current message. */
-public class SecurityAwareCustomMessageWrapper extends DiscoverySpiCustomMessage {
+public class SecurityAwareCustomMessageWrapper extends DiscoverySpiCustomMessage implements Message {
     /** */
     private static final long serialVersionUID = 0L;
 
     /** Security subject ID. */
-    private final UUID secSubjId;
+    @Order(0)
+    UUID secSubjId;
 
     /** Original message. */
-    private final DiscoveryCustomMessage delegate;
+    private DiscoveryCustomMessage delegate;
+
+    /** */
+    @Order(1)
+    // TODO: Should be removed in https://issues.apache.org/jira/browse/IGNITE-27627
+    Message msg;
+
+    /** Serialized message bytes. */
+    // TODO: Should be removed in https://issues.apache.org/jira/browse/IGNITE-27627
+    @Order(value = 2, method = "messageBytes")
+    byte[] msgBytes;
+
+    /** Default constructor for {@link MessageFactory}. */
+    public SecurityAwareCustomMessageWrapper() {
+        // No-op.
+    }
 
     /** */
     public SecurityAwareCustomMessageWrapper(DiscoveryCustomMessage delegate, UUID secSubjId) {
         this.delegate = delegate;
         this.secSubjId = secSubjId;
+
+        if (delegate instanceof Message)
+            msg = (Message)delegate;
     }
 
     /** Gets security Subject ID. */
@@ -57,7 +85,7 @@ public UUID securitySubjectId() {
      * @return Delegate.
      */
     public DiscoveryCustomMessage delegate() {
-        return delegate;
+        return msg != null ? (DiscoveryCustomMessage)msg : delegate;
     }
 
     /** {@inheritDoc} */
@@ -66,4 +94,38 @@ public DiscoveryCustomMessage delegate() {
 
         return ack == null ? null : new SecurityAwareCustomMessageWrapper(ack, secSubjId);
     }
+
+    /** */
+    public byte[] messageBytes() {
+        if (delegate instanceof Message)
+            return null;
+
+        if (msgBytes != null)
+            return msgBytes;
+
+        try {
+            return msgBytes = U.marshal(jdk(), delegate);
+        }
+        catch (IgniteCheckedException e) {
+            throw new IgniteException(e);
+        }
+    }
+
+    /** */
+    public void messageBytes(byte[] msgBytes) {
+        if (F.isEmpty(msgBytes))
+            return;
+
+        try {
+            delegate = U.unmarshal(jdk(), msgBytes, U.gridClassLoader());
+        }
+        catch (IgniteCheckedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override public short directType() {
+        return 501;
+    }
 }

modules/core/src/test/java/org/apache/ignite/internal/processors/security/NodeSecurityContextPropagationTest.java

@@ -28,7 +28,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
 import org.apache.commons.lang3.reflect.FieldUtils;
-import org.apache.ignite.IgniteCheckedException;
 import org.apache.ignite.configuration.DataRegionConfiguration;
 import org.apache.ignite.configuration.DataStorageConfiguration;
 import org.apache.ignite.configuration.IgniteConfiguration;
@@ -57,6 +56,7 @@
 import static org.apache.ignite.events.EventType.EVT_NODE_LEFT;
 import static org.apache.ignite.internal.IgniteNodeAttributes.ATTR_IGNITE_INSTANCE_NAME;
 import static org.apache.ignite.internal.events.DiscoveryCustomEvent.EVT_DISCOVERY_CUSTOM_EVT;
+import static org.apache.ignite.testframework.GridTestUtils.getFieldValue;
 import static org.apache.ignite.testframework.GridTestUtils.runAsync;
 import static org.apache.ignite.testframework.GridTestUtils.waitForCondition;
 
@@ -186,20 +186,7 @@ private boolean anyReceivedMessageMatch(IgniteEx ignite, Predicate<Object> predi
             Object unwrappedMsg = msg;
 
             if (msg instanceof TcpDiscoveryCustomEventMessage) {
-                DiscoveryCustomMessage customMsg = U.field(msg, "msg");
-
-                if (customMsg == null) {
-                    try {
-                        customMsg = U.unmarshal(
-                            ignite.context().marshallerContext().jdkMarshaller(),
-                            (byte[])U.field(msg, "msgBytes"),
-                            U.resolveClassLoader(ignite.configuration())
-                        );
-                    }
-                    catch (IgniteCheckedException e) {
-                        fail(e.getMessage());
-                    }
-                }
+                DiscoveryCustomMessage customMsg = getFieldValue(msg, "serMsg");
 
                 assert customMsg instanceof SecurityAwareCustomMessageWrapper;