Fixed copied packets and compression using velocity forwarding

src/main/java/net/minestom/server/extras/velocity/VelocityProxy.java

@@ -1,10 +1,13 @@
 package net.minestom.server.extras.velocity;
 
+import com.google.common.net.InetAddresses;
+import io.netty.buffer.ByteBuf;
 import net.minestom.server.utils.binary.BinaryReader;
 import org.jetbrains.annotations.NotNull;
 
 import javax.crypto.Mac;
 import javax.crypto.spec.SecretKeySpec;
+import java.net.InetAddress;
 import java.security.InvalidKeyException;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
@@ -17,6 +20,7 @@ import java.security.NoSuchAlgorithmException;
 public final class VelocityProxy {
 
     public static final String PLAYER_INFO_CHANNEL = "velocity:player_info";
+    private static final int SUPPORTED_FORWARDING_VERSION = 1;
 
     private static boolean enabled;
     private static byte[] secret;
@@ -48,7 +52,9 @@ public final class VelocityProxy {
 
         final byte[] signature = reader.readBytes(32);
 
-        final byte[] data = reader.getRemainingBytes();
+        ByteBuf buf = reader.getBuffer();
+        final byte[] data = new byte[buf.readableBytes()];
+        buf.getBytes(buf.readerIndex(), data);
 
         try {
             final Mac mac = Mac.getInstance("HmacSHA256");
@@ -61,12 +67,12 @@ public final class VelocityProxy {
             throw new AssertionError(e);
         }
 
-        /*int version = buf.readVarInt();
+        final int version = reader.readVarInt();
-        if (version != SUPPORTED_FORWARDING_VERSION) {
+        return version == SUPPORTED_FORWARDING_VERSION;
-            throw new IllegalStateException("Unsupported forwarding version " + version + ", wanted " + SUPPORTED_FORWARDING_VERSION);
+    }
-        }*/
 
-        return true;
+    public static InetAddress readAddress(@NotNull BinaryReader reader) {
+        return InetAddresses.forString(reader.readSizedString());
     }
 
 }

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

@@ -55,11 +55,6 @@ public class EncryptionResponsePacket implements ClientPreplayPacket {
                         final String string3 = new BigInteger(digestedData).toString(16);
                         final GameProfile gameProfile = MinecraftServer.getSessionService().hasJoinedServer(new GameProfile(null, loginUsername), string3);
                         nettyConnection.setEncryptionKey(getSecretKey());
-                        final int threshold = MinecraftServer.getCompressionThreshold();
-
-                        if (threshold > 0) {
-                            nettyConnection.enableCompression(threshold);
-                        }
 
                         MinecraftServer.getLOGGER().info("UUID of player {} is {}", loginUsername, gameProfile.getId());
                         CONNECTION_MANAGER.startPlayState(connection, gameProfile.getId(), gameProfile.getName());

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

@@ -12,6 +12,9 @@ import net.minestom.server.network.player.PlayerConnection;
 import net.minestom.server.utils.binary.BinaryReader;
 import org.jetbrains.annotations.NotNull;
 
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.util.UUID;
 
 public class LoginPluginResponsePacket implements ClientPreplayPacket {
@@ -34,16 +37,27 @@ public class LoginPluginResponsePacket implements ClientPreplayPacket {
 
             if (channel != null) {
                 boolean success = false;
+                SocketAddress socketAddress = null;
 
                 // Velocity
                 if (VelocityProxy.isEnabled() && channel.equals(VelocityProxy.PLAYER_INFO_CHANNEL)) {
                     if (data != null) {
                         BinaryReader reader = new BinaryReader(data);
                         success = VelocityProxy.checkIntegrity(reader);
+                        if (success) {
+                            // Get the real connection address
+                            final InetAddress address = VelocityProxy.readAddress(reader);
+                            final int port = ((java.net.InetSocketAddress) connection.getRemoteAddress()).getPort();
+                            socketAddress = new InetSocketAddress(address, port);
+                        }
                     }
                 }
 
                 if (success) {
+                    if (socketAddress != null) {
+                        nettyPlayerConnection.setRemoteAddress(socketAddress);
+                    }
+
                     // Proxy usage always mean that the server is in offline mode
                     final String username = nettyPlayerConnection.getLoginUsername();
                     final UUID playerUuid = CONNECTION_MANAGER.getPlayerConnectionUuid(connection, username);

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

@@ -27,9 +27,16 @@ public class LoginStartPacket implements ClientPreplayPacket {
     @Override
     public void process(@NotNull PlayerConnection connection) {
 
-        // Cache the login username
+        // Cache the login username and start compression if enabled
         if (connection instanceof NettyPlayerConnection) {
-            ((NettyPlayerConnection) connection).UNSAFE_setLoginUsername(username);
+            NettyPlayerConnection nettyPlayerConnection = (NettyPlayerConnection) connection;
+            nettyPlayerConnection.UNSAFE_setLoginUsername(username);
+
+            // Compression
+            final int threshold = MinecraftServer.getCompressionThreshold();
+            if (threshold > 0) {
+                nettyPlayerConnection.enableCompression(threshold);
+            }
         }
 
         // Proxy support (only for netty clients)
@@ -75,12 +82,6 @@ public class LoginStartPacket implements ClientPreplayPacket {
             // Offline
             final UUID playerUuid = CONNECTION_MANAGER.getPlayerConnectionUuid(connection, username);
 
-            final int threshold = MinecraftServer.getCompressionThreshold();
-
-            if (threshold > 0 && connection instanceof NettyPlayerConnection) {
-                ((NettyPlayerConnection) connection).enableCompression(threshold);
-            }
-
             CONNECTION_MANAGER.startPlayState(connection, playerUuid, username);
         }
     }

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

@@ -29,6 +29,8 @@ import java.util.concurrent.ConcurrentHashMap;
 public class NettyPlayerConnection extends PlayerConnection {
 
     private final SocketChannel channel;
+
+    private SocketAddress remoteAddress;
     @Getter
     private boolean encrypted = false;
     @Getter
@@ -50,6 +52,7 @@ public class NettyPlayerConnection extends PlayerConnection {
     public NettyPlayerConnection(@NotNull SocketChannel channel) {
         super();
         this.channel = channel;
+        this.remoteAddress = channel.remoteAddress();
     }
 
     /**
@@ -80,25 +83,25 @@ public class NettyPlayerConnection extends PlayerConnection {
 
     @Override
     public void sendPacket(@NotNull ByteBuf buffer, boolean copy) {
-        if ((encrypted || compressed) && copy) {
+        if (copy) {
             buffer = buffer.copy();
             buffer.retain();
             channel.writeAndFlush(buffer);
             buffer.release();
         } else {
-            getChannel().writeAndFlush(buffer);
+            channel.writeAndFlush(buffer);
         }
     }
 
     @Override
     public void writePacket(@NotNull ByteBuf buffer, boolean copy) {
-        if ((encrypted || compressed) && copy) {
+        if (copy) {
             buffer = buffer.copy();
             buffer.retain();
             channel.write(buffer);
             buffer.release();
         } else {
-            getChannel().write(buffer);
+            channel.write(buffer);
         }
     }
 
@@ -115,7 +118,18 @@ public class NettyPlayerConnection extends PlayerConnection {
     @NotNull
     @Override
     public SocketAddress getRemoteAddress() {
-        return getChannel().remoteAddress();
+        return remoteAddress;
+    }
+
+    /**
+     * Changes the internal remote address field.
+     * <p>
+     * Mostly unsafe, used internally when interacting with a proxy.
+     *
+     * @param remoteAddress the new connection remote address
+     */
+    public void setRemoteAddress(@NotNull SocketAddress remoteAddress) {
+        this.remoteAddress = remoteAddress;
     }
 
     @Override
@@ -141,7 +155,7 @@ public class NettyPlayerConnection extends PlayerConnection {
     }
 
     /**
-     * Sets the internal login username field
+     * Sets the internal login username field.
      *
      * @param loginUsername the new login username field
      */