Fix encryption adn zookeeper

Author: shishkovilja

modules/core/src/main/java/org/apache/ignite/internal/CoreMessagesProvider.java

@@ -38,6 +38,7 @@
 import org.apache.ignite.internal.managers.encryption.ChangeCacheEncryptionRequest;
 import org.apache.ignite.internal.managers.encryption.GenerateEncryptionKeyRequest;
 import org.apache.ignite.internal.managers.encryption.GenerateEncryptionKeyResponse;
+import org.apache.ignite.internal.managers.encryption.GroupKeyEncrypted;
 import org.apache.ignite.internal.managers.encryption.MasterKeyChangeRequest;
 import org.apache.ignite.internal.managers.encryption.NodeEncryptionKeys;
 import org.apache.ignite.internal.managers.eventstorage.GridEventStorageMessage;
@@ -644,6 +645,7 @@ public CoreMessagesProvider(Marshaller dfltMarsh, Marshaller schemaAwareMarsh, C
         withNoSchema(ChangeCacheEncryptionRequest.class);
         withNoSchema(MasterKeyChangeRequest.class);
         withNoSchema(NodeEncryptionKeys.class);
+        withNoSchema(GroupKeyEncrypted.class);
 
         // [13000 - 13300]: Control, configuration, diagnostincs and other messages.
         msgIdx = 13000;

modules/core/src/main/java/org/apache/ignite/internal/managers/encryption/GroupKeyEncrypted.java

@@ -18,19 +18,26 @@
 package org.apache.ignite.internal.managers.encryption;
 
 import java.io.Serializable;
+import org.apache.ignite.internal.Order;
+import org.apache.ignite.plugin.extensions.communication.Message;
 
 /**
  * Cache group encryption key with identifier. Key is encrypted.
  */
-public class GroupKeyEncrypted implements Serializable {
+public class GroupKeyEncrypted implements Serializable, Message {
     /** Serial version UID. */
     private static final long serialVersionUID = 0L;
 
     /** Encryption key ID. */
-    private final int id;
+    @Order(0)
+    int id;
 
     /** Encryption key. */
-    private final byte[] key;
+    @Order(1)
+    byte[] key;
+
+    /** */
+    public GroupKeyEncrypted() {}
 
     /**
      * @param id Encryption key ID.

modules/core/src/main/java/org/apache/ignite/internal/managers/encryption/NodeEncryptionKeys.java

@@ -44,6 +44,7 @@ public class NodeEncryptionKeys implements Serializable, Message {
     byte[] masterKeyDigest;
 
     /** Known i.e. stored in {@code ReadWriteMetastorage} keys from node. */
+    @Order(3)
     Map<Integer, List<GroupKeyEncrypted>> knownKeysWithIds;
 
     /** */

modules/zookeeper/src/main/java/org/apache/ignite/spi/discovery/zk/ZookeeperDiscoverySpi.java

@@ -143,6 +143,9 @@ public class ZookeeperDiscoverySpi extends IgniteSpiAdapter implements IgniteDis
     /** */
     private final ZookeeperDiscoveryStatistics stats = new ZookeeperDiscoveryStatistics();
 
+    /** */
+    private final ZkMessageFactory factory = new ZkMessageFactory();
+
     /**
      * @return Base path in ZK for znodes created by SPI.
      */
@@ -468,6 +471,9 @@ public DiscoverySpiNodeAuthenticator getAuthenticator() {
 
         registerMBean(igniteInstanceName, new ZookeeperDiscoverySpiMBeanImpl(this), ZookeeperDiscoverySpiMBean.class);
 
+//        factory.init(((IgniteEx)ignite).context().marshallerContext().jdkMarshaller(),
+//            U.resolveClassLoader(ignite.configuration()));
+
         try {
             impl.startJoinAndWait();
         }
@@ -558,7 +564,7 @@ private ZookeeperClusterNode initLocalNode() {
 
     /** {@inheritDoc} */
     @Override public MessageFactoryProvider messageFactoryProvider() {
-        return new ZkMessageFactory();
+        return factory;
     }
 
     /** {@inheritDoc} */

modules/zookeeper/src/main/java/org/apache/ignite/spi/discovery/zk/internal/DiscoveryMessageParser.java

@@ -31,7 +31,6 @@
 import org.apache.ignite.plugin.extensions.communication.MessageFactory;
 import org.apache.ignite.plugin.extensions.communication.MessageSerializer;
 import org.apache.ignite.spi.IgniteSpiException;
-import org.apache.ignite.spi.discovery.DiscoverySpiCustomMessage;
 
 import static org.apache.ignite.spi.communication.tcp.TcpCommunicationSpi.makeMessageType;
 
@@ -51,11 +50,11 @@ public DiscoveryMessageParser(MessageFactory msgFactory) {
     }
 
     /** Marshals discovery message to bytes array. */
-    public byte[] marshalZip(DiscoverySpiCustomMessage msg) {
+    public byte[] marshalZip(Message msg) {
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
 
         try (DeflaterOutputStream out = new DeflaterOutputStream(baos)) {
-            serializeMessage((Message)msg, out);
+            serializeMessage(msg, out);
         }
         catch (Exception e) {
             throw new IgniteSpiException("Failed to serialize message: " + msg, e);
@@ -65,12 +64,12 @@ public byte[] marshalZip(DiscoverySpiCustomMessage msg) {
     }
 
     /** Unmarshals discovery message from bytes array. */
-    public DiscoverySpiCustomMessage unmarshalZip(byte[] bytes) {
+    public <T extends Message> T unmarshalZip(byte[] bytes) {
         try (
             ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
             InflaterInputStream in = new InflaterInputStream(bais)
         ) {
-            return (DiscoverySpiCustomMessage)deserializeMessage(in);
+            return deserializeMessage(in);
         }
         catch (Exception e) {
             throw new IgniteSpiException("Failed to deserialize message.", e);
@@ -99,7 +98,7 @@ private void serializeMessage(Message m, OutputStream out) throws IOException {
     }
 
     /** */
-    private Message deserializeMessage(InputStream in) throws IOException {
+    private <T extends Message> T deserializeMessage(InputStream in) throws IOException {
         DirectMessageReader msgReader = new DirectMessageReader(msgFactory, null);
         ByteBuffer msgBuf = ByteBuffer.allocate(MSG_BUFFER_SIZE);
 
@@ -127,6 +126,6 @@ private Message deserializeMessage(InputStream in) throws IOException {
         }
         while (!finished);
 
-        return msg;
+        return (T)msg;
     }
 }

modules/zookeeper/src/main/java/org/apache/ignite/spi/discovery/zk/internal/ZkCommunicationErrorResolveFinishMessage.java

@@ -27,7 +27,7 @@
 /**
  *
  */
-class ZkCommunicationErrorResolveFinishMessage implements DiscoverySpiCustomMessage, ZkInternalMessage {
+public class ZkCommunicationErrorResolveFinishMessage implements DiscoverySpiCustomMessage, ZkInternalMessage {
     /** */
     @Order(0)
     UUID futId;

modules/zookeeper/src/main/java/org/apache/ignite/spi/discovery/zk/internal/ZkInternalJoinErrorMessage.java

@@ -18,19 +18,25 @@
 package org.apache.ignite.spi.discovery.zk.internal;
 
 import java.io.Serializable;
+import org.apache.ignite.internal.Order;
+import org.apache.ignite.plugin.extensions.communication.Message;
 
 /**
  *
  */
-class ZkInternalJoinErrorMessage implements ZkInternalMessage, Serializable {
+public class ZkInternalJoinErrorMessage implements ZkInternalMessage, Serializable, Message {
     /** */
     private static final long serialVersionUID = 0L;
 
     /** */
     transient boolean notifyNode = true;
 
     /** */
-    final String err;
+    @Order(0)
+    String err;
+
+    /** */
+    public ZkInternalJoinErrorMessage() {}
 
     /**
      * @param nodeInternalId Joining node internal ID.

modules/zookeeper/src/main/java/org/apache/ignite/spi/discovery/zk/internal/ZkJoiningNodeData.java

@@ -19,26 +19,45 @@
 
 import java.io.Serializable;
 import java.util.Map;
+import org.apache.ignite.IgniteCheckedException;
+import org.apache.ignite.internal.Compress;
+import org.apache.ignite.internal.MarshallableMessage;
+import org.apache.ignite.internal.Order;
+import org.apache.ignite.internal.util.tostring.GridToStringExclude;
 import org.apache.ignite.internal.util.tostring.GridToStringInclude;
 import org.apache.ignite.internal.util.typedef.internal.S;
+import org.apache.ignite.internal.util.typedef.internal.U;
+import org.apache.ignite.marshaller.Marshaller;
+import org.apache.ignite.plugin.extensions.communication.Message;
 
 /**
  *
  */
-class ZkJoiningNodeData implements Serializable {
+public class ZkJoiningNodeData implements Serializable, MarshallableMessage {
     /** */
     private static final long serialVersionUID = 0L;
 
     /** */
-    private int partCnt;
+    @Order(0)
+    int partCnt;
 
     /** */
     @GridToStringInclude
     private ZookeeperClusterNode node;
 
     /** */
+    @Order(1)
+    @GridToStringExclude
+    transient byte[] nodeBytes;
+
+    /** */
+    @Order(2)
+    @Compress
     @GridToStringInclude
-    private Map<Integer, Serializable> discoData;
+    Map<Integer, Message> discoData;
+
+    /** */
+    public ZkJoiningNodeData() {}
 
     /**
      * @param partCnt Number of parts in multi-parts message.
@@ -51,7 +70,7 @@ class ZkJoiningNodeData implements Serializable {
      * @param node Node.
      * @param discoData Discovery data.
      */
-    ZkJoiningNodeData(ZookeeperClusterNode node, Map<Integer, Serializable> discoData) {
+    ZkJoiningNodeData(ZookeeperClusterNode node, Map<Integer, Message> discoData) {
         assert node != null && node.id() != null : node;
         assert discoData != null;
 
@@ -76,12 +95,24 @@ ZookeeperClusterNode node() {
     /**
      * @return Discovery data.
      */
-    Map<Integer, Serializable> discoveryData() {
+    Map<Integer, Message> discoveryData() {
         return discoData;
     }
 
     /** {@inheritDoc} */
     @Override public String toString() {
         return S.toString(ZkJoiningNodeData.class, this);
     }
+
+    /** {@inheritDoc} */
+    @Override public void prepareMarshal(Marshaller marsh) throws IgniteCheckedException {
+        if (node != null)
+            nodeBytes = U.marshal(marsh, node);
+    }
+
+    /** {@inheritDoc} */
+    @Override public void finishUnmarshal(Marshaller marsh, ClassLoader clsLdr) throws IgniteCheckedException {
+        if (nodeBytes != null)
+            node = U.unmarshal(marsh, nodeBytes, clsLdr);
+    }
 }

modules/zookeeper/src/main/java/org/apache/ignite/spi/discovery/zk/internal/ZkMessageFactory.java

@@ -17,16 +17,18 @@
 
 package org.apache.ignite.spi.discovery.zk.internal;
 
+import org.apache.ignite.internal.plugin.AbstractMarshallableMessageFactoryProvider;
 import org.apache.ignite.plugin.extensions.communication.MessageFactory;
-import org.apache.ignite.plugin.extensions.communication.MessageFactoryProvider;
 
 /** */
-public class ZkMessageFactory implements MessageFactoryProvider {
+public class ZkMessageFactory extends AbstractMarshallableMessageFactoryProvider {
     /** {@inheritDoc} */
     @Override public void registerAll(MessageFactory factory) {
-        factory.register(400, ZkCommunicationErrorResolveFinishMessage::new, new ZkCommunicationErrorResolveFinishMessageSerializer());
-        factory.register(401, ZkCommunicationErrorResolveStartMessage::new, new ZkCommunicationErrorResolveStartMessageSerializer());
-        factory.register(402, ZkForceNodeFailMessage::new, new ZkForceNodeFailMessageSerializer());
-        factory.register(403, ZkNoServersMessage::new, new ZkNoServersMessageSerializer());
+        register(factory, ZkCommunicationErrorResolveFinishMessage.class, (short)400, dfltMarsh, resolvedClsLdr);
+        register(factory, ZkCommunicationErrorResolveStartMessage.class, (short)401, dfltMarsh, resolvedClsLdr);
+        register(factory, ZkForceNodeFailMessage.class, (short)402, dfltMarsh, resolvedClsLdr);
+        register(factory, ZkNoServersMessage.class, (short)403, dfltMarsh, resolvedClsLdr);
+        register(factory, ZkJoiningNodeData.class, (short)404, dfltMarsh, resolvedClsLdr);
+        register(factory, ZkInternalJoinErrorMessage.class, (short)405, dfltMarsh, resolvedClsLdr);
     }
 }

modules/zookeeper/src/main/java/org/apache/ignite/spi/discovery/zk/internal/ZkNoServersMessage.java

@@ -24,7 +24,7 @@
 /**
  *
  */
-class ZkNoServersMessage implements DiscoverySpiCustomMessage, ZkInternalMessage {
+public class ZkNoServersMessage implements DiscoverySpiCustomMessage, ZkInternalMessage {
     /** {@inheritDoc} */
     @Nullable @Override public DiscoverySpiCustomMessage ackMessage() {
         return null;