From ab568405dd1394045605005baf3960209bb42a77 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sat, 15 Sep 2018 02:22:52 -0400
Subject: [PATCH] Cleaned up client plugin message logic.

---
 .../client/ClientPlaySessionHandler.java      | 33 +++++++------------
 .../protocol/util/PluginMessageUtil.java      | 17 +++++++---
 2 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
index f550865d2..55b023af3 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
@@ -252,8 +252,8 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
         return serverBossBars;
     }
 
-    public void handleClientPluginMessage(PluginMessage packet) {
-        if (packet.getChannel().equals("REGISTER") || packet.getChannel().equals("minecraft:register")) {
+    private void handleClientPluginMessage(PluginMessage packet) {
+        if (PluginMessageUtil.isMCRegister(packet)) {
             List<String> actuallyRegistered = new ArrayList<>();
             List<String> channels = PluginMessageUtil.getChannels(packet);
             for (String channel : channels) {
@@ -270,21 +270,13 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
                 PluginMessage newRegisterPacket = PluginMessageUtil.constructChannelsPacket(packet.getChannel(), actuallyRegistered);
                 player.getConnectedServer().getMinecraftConnection().write(newRegisterPacket);
             }
-
-            return;
-        }
-
-        if (packet.getChannel().equals("UNREGISTER") || packet.getChannel().equals("minecraft:unregister")) {
+        } else if (PluginMessageUtil.isMCUnregister(packet)) {
             List<String> channels = PluginMessageUtil.getChannels(packet);
             clientPluginMsgChannels.removeAll(channels);
-        }
-
-        if (PluginMessageUtil.isMCBrand(packet)) {
+            player.getConnectedServer().getMinecraftConnection().write(packet);
+        } else if (PluginMessageUtil.isMCBrand(packet)) {
             player.getConnectedServer().getMinecraftConnection().write(PluginMessageUtil.rewriteMCBrand(packet));
-            return;
-        }
-
-        if (player.getConnectedServer().isLegacyForge() && !player.getConnectedServer().hasCompletedJoin()) {
+        } else if (player.getConnectedServer().isLegacyForge() && !player.getConnectedServer().hasCompletedJoin()) {
             if (packet.getChannel().equals(VelocityConstants.FORGE_LEGACY_HANDSHAKE_CHANNEL)) {
                 // Always forward the FML handshake to the remote server.
                 player.getConnectedServer().getMinecraftConnection().write(packet);
@@ -294,13 +286,12 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
                 // be sent once the JoinGame packet has been received by the proxy.
                 loginPluginMessages.add(packet);
             }
-            return;
-        }
-
-        MessageHandler.ForwardStatus status = server.getChannelRegistrar().handlePluginMessage(player,
-                ChannelSide.FROM_CLIENT, packet);
-        if (status == MessageHandler.ForwardStatus.FORWARD) {
-            player.getConnectedServer().writeIfJoined(packet);
+        } else {
+            MessageHandler.ForwardStatus status = server.getChannelRegistrar().handlePluginMessage(player,
+                    ChannelSide.FROM_CLIENT, packet);
+            if (status == MessageHandler.ForwardStatus.FORWARD) {
+                player.getConnectedServer().writeIfJoined(packet);
+            }
         }
     }
 
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/util/PluginMessageUtil.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/util/PluginMessageUtil.java
index 075058486..e3fdeab31 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/util/PluginMessageUtil.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/util/PluginMessageUtil.java
@@ -20,13 +20,20 @@ public enum PluginMessageUtil {
         return message.getChannel().equals("MC|Brand") || message.getChannel().equals("minecraft:brand");
     }
 
+    public static boolean isMCRegister(PluginMessage message) {
+        Preconditions.checkNotNull(message, "message");
+        return message.getChannel().equals("REGISTER") || message.getChannel().equals("minecraft:register");
+    }
+
+    public static boolean isMCUnregister(PluginMessage message) {
+        Preconditions.checkNotNull(message, "message");
+        return message.getChannel().equals("UNREGISTER") || message.getChannel().equals("minecraft:unregister");
+    }
+
     public static List<String> getChannels(PluginMessage message) {
         Preconditions.checkNotNull(message, "message");
-        Preconditions.checkArgument(message.getChannel().equals("REGISTER") ||
-                message.getChannel().equals("UNREGISTER") ||
-                message.getChannel().equals("minecraft:register") ||
-                message.getChannel().equals("minecraft:unregister"),
-                "Unknown channel type " + message.getChannel());
+        Preconditions.checkArgument(isMCRegister(message) || isMCUnregister(message),"Unknown channel type %s",
+                message.getChannel());
         String channels = new String(message.getData(), StandardCharsets.UTF_8);
         return ImmutableList.copyOf(channels.split("\0"));
     }