Fixed threading issues with encryption

src/dorkbox/network/connection/ConnectionImpl.java

@@ -54,6 +54,7 @@ import java.util.Map;
 import java.util.WeakHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 
 
 /**
@@ -100,6 +101,10 @@ class ConnectionImpl extends ChannelInboundHandlerAdapter implements Connection,
 
     private final Map<Integer, RemoteObject> proxyIdCache = Collections.synchronizedMap(new WeakHashMap<Integer, RemoteObject>(8));
 
+    // The IV for AES-GCM must be 12 bytes, since it's 4 (salt) + 8 (external counter) + 4 (GCM counter)
+    // The 12 bytes IV is created during connection registration, and during the AES-GCM crypto, we override the last 8 with this
+    // counter, which is also transmitted as an optimized int. (which is why it starts at 0, so the transmitted bytes are small)
+    private final AtomicLong aes_gcm_iv = new AtomicLong(0);
 
     /**
      * All of the parameters can be null, when metaChannel wants to get the base class type
@@ -143,13 +148,28 @@ class ConnectionImpl extends ChannelInboundHandlerAdapter implements Connection,
 
 
     /**
-     * @return the AES key/IV, etc associated with this connection
+     * @return a threadlocal AES key + IV. key=32 byte, iv=12 bytes (AES-GCM implementation). This is a threadlocal
+     *          because multiple protocols can be performing crypto AT THE SAME TIME, and so we have to make sure that operations don't
+     *          clobber each other
      */
     final
     ParametersWithIV getCryptoParameters() {
         return this.channelWrapper.cryptoParameters();
     }
 
+    /**
+     * This is the per-message sequence number.
+     *
+     *  The IV for AES-GCM must be 12 bytes, since it's 4 (salt) + 8 (external counter) + 4 (GCM counter)
+     *  The 12 bytes IV is created during connection registration, and during the AES-GCM crypto, we override the last 8 with this
+     *  counter, which is also transmitted as an optimized int. (which is why it starts at 0, so the transmitted bytes are small)
+     */
+    final
+    long getNextGcmSequence() {
+        return aes_gcm_iv.getAndIncrement();
+    }
+
+
     /**
      * Has the remote ECC public key changed. This can be useful if specific actions are necessary when the key has changed.
      */

src/dorkbox/network/connection/KryoCryptoSerializationManager.java

@@ -60,7 +60,6 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Threads reading/writing, it messes up a single instance. it is possible to use a single kryo with the use of synchronize, however - that
@@ -86,13 +85,6 @@ class KryoCryptoSerializationManager implements CryptoSerializationManager {
      */
     public static KryoCryptoSerializationManager DEFAULT = DEFAULT();
 
-
-    // The IV for AES-GCM must be 12 bytes, since it's 4 (salt) + 8 (external counter) + 4 (GCM counter)
-    // The 12 bytes IV is created during connection registration, and during the AES-GCM crypto, we override the last 8 with this
-    // counter, which is also transmitted as an optimized int. (which is why it starts at 0, so the transmitted bytes are small)
-    private final AtomicLong aes_gcm_iv = new AtomicLong(0);
-
-
     public static
     KryoCryptoSerializationManager DEFAULT() {
         return DEFAULT(true, true);
@@ -511,7 +503,7 @@ class KryoCryptoSerializationManager implements CryptoSerializationManager {
     void writeWithCrypto(final ConnectionImpl connection, final ByteBuf buffer, final Object message) throws IOException {
         final KryoExtra kryo = kryoPool.take();
         try {
-            kryo.writeCrypto(connection, buffer, message, aes_gcm_iv.getAndIncrement());
+            kryo.writeCrypto(connection, buffer, message);
         } finally {
             kryoPool.put(kryo);
         }

src/dorkbox/network/connection/KryoExtra.java

@@ -80,7 +80,6 @@ class KryoExtra extends Kryo {
     private byte[] decompressOutput;
     private ByteBuf decompressBuf;
 
-
     public
     KryoExtra() {
     }
@@ -123,8 +122,7 @@ class KryoExtra extends Kryo {
 
 
     public synchronized
-    void writeCrypto(final ConnectionImpl connection, final ByteBuf buffer, final Object message, final long gcmIVCounter) throws IOException {
+    void writeCrypto(final ConnectionImpl connection, final ByteBuf buffer, final Object message) throws IOException {
-
         // required by RMI and some serializers to determine which connection wrote (or has info about) this object
         this.connection = connection;
 
@@ -213,10 +211,12 @@ class KryoExtra extends Kryo {
 
 
 
-        /////// encrypting data
+        /////// encrypting data.
+        final long nextGcmSequence = connection.getNextGcmSequence();
 
+        // this is a threadlocal, so that we don't clobber other threads that are performing crypto on the same connection at the same time
         final ParametersWithIV cryptoParameters = connection.getCryptoParameters();
-        BigEndian.Long_.toBytes(gcmIVCounter, cryptoParameters.getIV(), 4); // put our counter into the IV
+        BigEndian.Long_.toBytes(nextGcmSequence, cryptoParameters.getIV(), 4); // put our counter into the IV
 
         final GCMBlockCipher aes = this.aesEngine;
         aes.reset();
@@ -225,7 +225,7 @@ class KryoExtra extends Kryo {
         byte[] cryptoOutput;
 
         // lazy initialize the crypto output buffer
-        int cryptoSize = aes.getOutputSize(length);
+        int cryptoSize = length + 16;   // from:  aes.getOutputSize(length);
 
         // 'output' is the temp byte array
         if (cryptoSize > cryptoOutputLength) {
@@ -249,7 +249,7 @@ class KryoExtra extends Kryo {
         buffer.writeByte(crypto);
 
         // write out our GCM counter
-        OptimizeUtilsByteBuf.writeLong(buffer, gcmIVCounter, true);
+        OptimizeUtilsByteBuf.writeLong(buffer, nextGcmSequence, true);
 
 //        System.err.println("out " + gcmIVCounter);
 
@@ -317,6 +317,7 @@ class KryoExtra extends Kryo {
         // have to make sure to set the position of the buffer, since our conversion to array DOES NOT set the new reader index.
         buffer.readerIndex(buffer.readerIndex() + length);
 
+        // this is a threadlocal, so that we don't clobber other threads that are performing crypto on the same connection at the same time
         final ParametersWithIV cryptoParameters = connection.getCryptoParameters();
         BigEndian.Long_.toBytes(gcmIVCounter, cryptoParameters.getIV(), 4); // put our counter into the IV
 
@@ -324,7 +325,7 @@ class KryoExtra extends Kryo {
         aes.reset();
         aes.init(false, cryptoParameters);
 
-        int cryptoSize = aes.getOutputSize(length);
+        int cryptoSize = length - 16; // from:  aes.getOutputSize(length);
 
         // lazy initialize the decrypt output buffer
         byte[] decryptOutputArray;
@@ -373,11 +374,7 @@ class KryoExtra extends Kryo {
         // read the object from the buffer.
         reader.setBuffer(inputBuf);
 
-        final Object o = readClassAndObject(reader);
+        return readClassAndObject(reader); // this properly sets the readerIndex, but only if it's the correct buffer
-        if (o == null) {
-            System.err.println("what?");
-        }
-        return o; // this properly sets the readerIndex, but only if it's the correct buffer
     }
 
     @Override

src/dorkbox/network/connection/wrapper/ChannelNetworkWrapper.java

@@ -15,7 +15,11 @@
  */
 package dorkbox.network.connection.wrapper;
 
-import dorkbox.network.connection.*;
+import dorkbox.network.connection.Connection;
+import dorkbox.network.connection.ConnectionPointWriter;
+import dorkbox.network.connection.EndPoint;
+import dorkbox.network.connection.ISessionManager;
+import dorkbox.network.connection.UdpServer;
 import dorkbox.network.connection.registration.MetaChannel;
 import io.netty.channel.Channel;
 import io.netty.channel.EventLoop;
@@ -31,14 +35,17 @@ class ChannelNetworkWrapper<C extends Connection> implements ChannelWrapper<C> {
     private final ChannelNetwork udp;
     private final ChannelNetwork udt;
 
-    private final ParametersWithIV cryptoAesKeyAndIV;
-
     // did the remote connection public ECC key change?
     private final boolean remotePublicKeyChanged;
 
     private final String remoteAddress;
     private final EventLoop eventLoop;
 
+    // GCM IV. hacky way to prevent tons of GC and to not clobber the original parameters
+    private final byte[] aesKey; // AES-256 requires 32 bytes
+    private final byte[] aesIV; // AES-GCM requires 12 bytes
+
+    private final ThreadLocal<ParametersWithIV> cryptoParameters;
 
     /**
      * @param udpServer is null when created by the client, non-null when created by the server
@@ -76,9 +83,16 @@ class ChannelNetworkWrapper<C extends Connection> implements ChannelWrapper<C> {
         this.remotePublicKeyChanged = metaChannel.changedRemoteKey;
 
         // AES key & IV (only for networked connections)
-        this.cryptoAesKeyAndIV = new ParametersWithIV(new KeyParameter(metaChannel.aesKey), metaChannel.aesIV);
+        aesKey = metaChannel.aesKey;
-        // TODO: have to be able to get a NEW IV, so we can rotate keys!
+        aesIV = metaChannel.aesIV;
-        // NOTE: IV's must be a NONCE! (ie: one time use only!!!)
+
+        cryptoParameters = new ThreadLocal<ParametersWithIV>() {
+            @Override
+            protected
+            ParametersWithIV initialValue() {
+                return new ParametersWithIV(new KeyParameter(aesKey), aesIV);
+            }
+        };
     }
 
     public final
@@ -136,10 +150,16 @@ class ChannelNetworkWrapper<C extends Connection> implements ChannelWrapper<C> {
         return this.eventLoop;
     }
 
+
+    /**
+     * @return a threadlocal AES key + IV. key=32 byte, iv=12 bytes (AES-GCM implementation). This is a threadlocal
+     *          because multiple protocols can be performing crypto AT THE SAME TIME, and so we have to make sure that operations don't
+     *          clobber each other
+     */
     @Override
     public
     ParametersWithIV cryptoParameters() {
-        return this.cryptoAesKeyAndIV;
+        return this.cryptoParameters.get();
     }
 
     @Override

src/dorkbox/network/connection/wrapper/ChannelWrapper.java

@@ -42,6 +42,11 @@ interface ChannelWrapper<C extends Connection> {
 
     EventLoop getEventLoop();
 
+    /**
+     * @return a threadlocal AES key + IV. key=32 byte, iv=12 bytes (AES-GCM implementation). This is a threadlocal
+     *          because multiple protocols can be performing crypto AT THE SAME TIME, and so we have to make sure that operations don't
+     *          clobber each other
+     */
     ParametersWithIV cryptoParameters();
 
     /**