Volatile encryption (#515)

src/main/java/net/minestom/server/network/packet/client/login/EncryptionResponsePacket.java

@@ -22,6 +22,7 @@ import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.nio.charset.StandardCharsets;
+import java.security.KeyPair;
 import java.util.Arrays;
 import java.util.UUID;
 
@@ -50,8 +51,9 @@ public class EncryptionResponsePacket implements ClientPreplayPacket {
                 MinecraftServer.LOGGER.error("{} tried to login with an invalid nonce!", loginUsername);
                 return;
             }
-
+            final KeyPair keyPair = MojangAuth.getKeyPair();
-            final byte[] digestedData = MojangCrypt.digestData("", MojangAuth.getKeyPair().getPublic(), getSecretKey());
+            final SecretKey secretKey = MojangCrypt.decryptByteToSecretKey(keyPair.getPrivate(), sharedSecret);
+            final byte[] digestedData = MojangCrypt.digestData("", keyPair.getPublic(), secretKey);
             if (digestedData == null) {
                 // Incorrect key, probably because of the client
                 MinecraftServer.LOGGER.error("Connection {} failed initializing encryption.", socketConnection.getRemoteAddress());
@@ -78,7 +80,7 @@ public class EncryptionResponsePacket implements ClientPreplayPacket {
                         // Invalid response
                         return;
                     }
-                    socketConnection.setEncryptionKey(getSecretKey());
+                    socketConnection.setEncryptionKey(secretKey);
                     UUID profileUUID = UUID.fromString(gameProfile.get("id").getAsString()
                             .replaceFirst("(\\w{8})(\\w{4})(\\w{4})(\\w{4})(\\w{12})", "$1-$2-$3-$4-$5"));
                     final String profileName = gameProfile.get("name").getAsString();
@@ -104,10 +106,6 @@ public class EncryptionResponsePacket implements ClientPreplayPacket {
         ByteArrayData.encodeByteArray(writer, verifyToken);
     }
 
-    private SecretKey getSecretKey() {
-        return MojangCrypt.decryptByteToSecretKey(MojangAuth.getKeyPair().getPrivate(), sharedSecret);
-    }
-
     private byte[] getNonce() {
         return MojangAuth.getKeyPair().getPrivate() == null ?
                 this.verifyToken : MojangCrypt.decryptUsingKey(MojangAuth.getKeyPair().getPrivate(), this.verifyToken);

src/main/java/net/minestom/server/network/packet/client/login/LoginStartPacket.java

@@ -69,8 +69,7 @@ public class LoginStartPacket implements ClientPreplayPacket {
             }
             final PlayerSocketConnection socketConnection = (PlayerSocketConnection) connection;
             socketConnection.setConnectionState(ConnectionState.LOGIN);
-            EncryptionRequestPacket encryptionRequestPacket = new EncryptionRequestPacket(socketConnection);
+            socketConnection.sendPacket(new EncryptionRequestPacket(socketConnection));
-            socketConnection.sendPacket(encryptionRequestPacket);
         } else {
             final boolean bungee = BungeeCordProxy.isEnabled();
             // Offline

src/main/java/net/minestom/server/network/player/PlayerSocketConnection.java

@@ -36,6 +36,7 @@ import java.nio.channels.SocketChannel;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.zip.DataFormatException;
 import java.util.zip.Inflater;
 
@@ -47,18 +48,18 @@ import java.util.zip.Inflater;
 @ApiStatus.Internal
 public class PlayerSocketConnection extends PlayerConnection {
     private final static Logger LOGGER = LoggerFactory.getLogger(PlayerSocketConnection.class);
+    private static final AtomicReferenceFieldUpdater<PlayerSocketConnection, EncryptionInfo> ENCRYPTION_UPDATER =
+            AtomicReferenceFieldUpdater.newUpdater(PlayerSocketConnection.class, EncryptionInfo.class, "encryptionInfo");
 
     private final Worker worker;
     private final SocketChannel channel;
     private SocketAddress remoteAddress;
 
-    private boolean encrypted = false;
+    private volatile EncryptionInfo encryptionInfo;
     private boolean compressed = false;
 
     //Could be null. Only used for Mojang Auth
     private byte[] nonce = new byte[4];
-    private Cipher decryptCipher;
-    private Cipher encryptCipher;
 
     // Data from client packets
     private String loginUsername;
@@ -92,11 +93,11 @@ public class PlayerSocketConnection extends PlayerConnection {
     public void processPackets(Worker.Context workerContext, PacketProcessor packetProcessor) {
         final BinaryBuffer readBuffer = workerContext.readBuffer;
         // Decrypt data
-        if (encrypted) {
+        EncryptionInfo encryptionInfo = this.encryptionInfo;
-            final Cipher cipher = decryptCipher;
+        if (encryptionInfo != null) {
             ByteBuffer input = readBuffer.asByteBuffer(0, readBuffer.writerOffset());
             try {
-                cipher.update(input, input.duplicate());
+                encryptionInfo.decryptCipher().update(input, input.duplicate());
             } catch (ShortBufferException e) {
                 MinecraftServer.getExceptionManager().handleException(e);
                 return;
@@ -175,12 +176,10 @@ public class PlayerSocketConnection extends PlayerConnection {
      * @throws IllegalStateException if encryption is already enabled for this connection
      */
     public void setEncryptionKey(@NotNull SecretKey secretKey) {
-        Check.stateCondition(encrypted, "Encryption is already enabled!");
+        final EncryptionInfo prev = ENCRYPTION_UPDATER.getAndSet(this,
-        synchronized (bufferLock) {
+                new EncryptionInfo(MojangCrypt.getCipher(2, secretKey), MojangCrypt.getCipher(1, secretKey)));
-            this.decryptCipher = MojangCrypt.getCipher(2, secretKey);
+        if (prev != null)
-            this.encryptCipher = MojangCrypt.getCipher(1, secretKey);
+            throw new IllegalStateException("Encryption was already enabled for this connection");
-            this.encrypted = true;
-        }
     }
 
     /**
@@ -231,19 +230,20 @@ public class PlayerSocketConnection extends PlayerConnection {
 
     @ApiStatus.Internal
     public void write(@NotNull ByteBuffer buffer, int index, int length) {
-        synchronized (bufferLock) {
+        EncryptionInfo encryptionInfo = this.encryptionInfo;
-            if (encrypted) { // Encryption support
+        if (encryptionInfo != null) {
-                ByteBuffer output = PacketUtils.localBuffer();
+            ByteBuffer output = PacketUtils.localBuffer();
-                try {
+            try {
-                    this.encryptCipher.update(buffer.slice(index,length), output);
+                encryptionInfo.encryptCipher().update(buffer.slice(index, length), output);
-                    buffer = output.flip();
+                buffer = output.flip();
-                    index = 0;
+                index = 0;
-                } catch (ShortBufferException e) {
+            } catch (ShortBufferException e) {
-                    MinecraftServer.getExceptionManager().handleException(e);
+                MinecraftServer.getExceptionManager().handleException(e);
-                    return;
+                return;
-                }
             }
+        }
 
+        synchronized (bufferLock) {
             BinaryBuffer localBuffer = tickBuffer.getPlain();
             final int capacity = localBuffer.capacity();
             if (length <= capacity) {
@@ -472,4 +472,7 @@ public class PlayerSocketConnection extends PlayerConnection {
     public void setNonce(byte[] nonce) {
         this.nonce = nonce;
     }
+
+    public record EncryptionInfo(Cipher decryptCipher, Cipher encryptCipher) {
+    }
 }