From 70e047e6eea6c011e9fa69eb9d395374f6b394a6 Mon Sep 17 00:00:00 2001
From: themode <themode@outlook.fr>
Date: Fri, 23 Oct 2020 23:55:42 +0200
Subject: [PATCH] Added buffer length check in DynamicChunk#readChunk

---
 .../client/handshake/HandshakePacket.java     | 12 +++++-
 .../network/player/NettyPlayerConnection.java | 43 +++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/src/main/java/net/minestom/server/network/packet/client/handshake/HandshakePacket.java b/src/main/java/net/minestom/server/network/packet/client/handshake/HandshakePacket.java
index 288e09436..b2443fe11 100644
--- a/src/main/java/net/minestom/server/network/packet/client/handshake/HandshakePacket.java
+++ b/src/main/java/net/minestom/server/network/packet/client/handshake/HandshakePacket.java
@@ -6,6 +6,7 @@ import net.minestom.server.chat.ColoredText;
 import net.minestom.server.network.ConnectionState;
 import net.minestom.server.network.packet.client.ClientPreplayPacket;
 import net.minestom.server.network.packet.server.login.LoginDisconnect;
+import net.minestom.server.network.player.NettyPlayerConnection;
 import net.minestom.server.network.player.PlayerConnection;
 import net.minestom.server.utils.binary.BinaryReader;
 
@@ -36,8 +37,15 @@ public class HandshakePacket implements ClientPreplayPacket {
                 connection.setConnectionState(ConnectionState.STATUS);
                 break;
             case 2:
-                connection.setConnectionState(ConnectionState.LOGIN);
-                if (protocolVersion != MinecraftServer.PROTOCOL_VERSION) {
+                if (protocolVersion == MinecraftServer.PROTOCOL_VERSION) {
+                    connection.setConnectionState(ConnectionState.LOGIN);
+
+                    if (connection instanceof NettyPlayerConnection) {
+                        // Give to the connection the server info that the client used
+                        ((NettyPlayerConnection) connection).refreshServerInformation(serverAddress, serverPort);
+                    }
+                } else {
+                    // Incorrect client version
                     connection.sendPacket(new LoginDisconnect(INVALID_VERSION_TEXT.toString()));
                     connection.disconnect();
                 }
diff --git a/src/main/java/net/minestom/server/network/player/NettyPlayerConnection.java b/src/main/java/net/minestom/server/network/player/NettyPlayerConnection.java
index 217cee4b9..a6ac1b06a 100644
--- a/src/main/java/net/minestom/server/network/player/NettyPlayerConnection.java
+++ b/src/main/java/net/minestom/server/network/player/NettyPlayerConnection.java
@@ -10,6 +10,7 @@ import net.minestom.server.extras.mojangAuth.MojangCrypt;
 import net.minestom.server.network.netty.codec.PacketCompressor;
 import net.minestom.server.network.packet.server.ServerPacket;
 import net.minestom.server.network.packet.server.login.SetCompressionPacket;
+import net.minestom.server.utils.validate.Check;
 
 import javax.crypto.SecretKey;
 import java.net.SocketAddress;
@@ -27,12 +28,22 @@ public class NettyPlayerConnection extends PlayerConnection {
     @Getter
     private boolean compressed = false;
 
+    private String serverAddress;
+    private int serverPort;
+
     public NettyPlayerConnection(SocketChannel channel) {
         super();
         this.channel = channel;
     }
 
+    /**
+     * Sets the encryption key and add the channels to the pipeline.
+     *
+     * @param secretKey the secret key to use in the encryption
+     * @throws IllegalStateException if encryption is already enabled for this connection
+     */
     public void setEncryptionKey(SecretKey secretKey) {
+        Check.stateCondition(encrypted, "Encryption is already enabled!");
         this.encrypted = true;
         getChannel().pipeline().addBefore("framer", "decrypt", new Decrypter(MojangCrypt.getCipher(2, secretKey)));
         getChannel().pipeline().addBefore("framer", "encrypt", new Encrypter(MojangCrypt.getCipher(1, secretKey)));
@@ -93,4 +104,36 @@ public class NettyPlayerConnection extends PlayerConnection {
         return channel;
     }
 
+    /**
+     * Get the server address that the client used to connect.
+     * <p>
+     * WARNING: it is given by the client, it is possible for it to be wrong.
+     *
+     * @return the server address used
+     */
+    public String getServerAddress() {
+        return serverAddress;
+    }
+
+    /**
+     * Get the server port that the client used to connect.
+     * <p>
+     * WARNING: it is given by the client, it is possible for it to be wrong.
+     *
+     * @return the server port used
+     */
+    public int getServerPort() {
+        return serverPort;
+    }
+
+    /**
+     * Used in {@link net.minestom.server.network.packet.client.handshake.HandshakePacket} to change the internal fields.
+     *
+     * @param serverAddress the server address which the client used
+     * @param serverPort    the server port which the client used
+     */
+    public void refreshServerInformation(String serverAddress, int serverPort) {
+        this.serverAddress = serverAddress;
+        this.serverPort = serverPort;
+    }
 }