From 5c809c24990ec06ecdb0ee6550d4d10595ed9323 Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Fri, 26 Aug 2016 16:11:38 +1000
Subject: [PATCH] #1945: Relay MC|Brand messages

---
 .../net/md_5/bungee/protocol/packet/PluginMessage.java | 10 ++++++++++
 .../src/main/java/net/md_5/bungee/ServerConnector.java |  3 +--
 .../net/md_5/bungee/connection/InitialHandler.java     |  6 +++---
 .../net/md_5/bungee/connection/UpstreamBridge.java     |  5 ++---
 4 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/protocol/src/main/java/net/md_5/bungee/protocol/packet/PluginMessage.java b/protocol/src/main/java/net/md_5/bungee/protocol/packet/PluginMessage.java
index bbaef39a..6dcdece5 100644
--- a/protocol/src/main/java/net/md_5/bungee/protocol/packet/PluginMessage.java
+++ b/protocol/src/main/java/net/md_5/bungee/protocol/packet/PluginMessage.java
@@ -1,6 +1,7 @@
 package net.md_5.bungee.protocol.packet;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Predicate;
 import net.md_5.bungee.protocol.DefinedPacket;
 import io.netty.buffer.ByteBuf;
 import java.io.ByteArrayInputStream;
@@ -20,6 +21,15 @@ import net.md_5.bungee.protocol.ProtocolConstants;
 public class PluginMessage extends DefinedPacket
 {
 
+    public static final Predicate<PluginMessage> SHOULD_RELAY = new Predicate<PluginMessage>()
+    {
+        @Override
+        public boolean apply(PluginMessage input)
+        {
+            return input.getTag().equals( "REGISTER" ) || input.getTag().equals( "MC|Brand" );
+        }
+    };
+    //
     private String tag;
     private byte[] data;
 
diff --git a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
index 976636ac..e599e7e0 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
@@ -30,7 +30,6 @@ import net.md_5.bungee.netty.HandlerBoss;
 import net.md_5.bungee.netty.PacketHandler;
 import net.md_5.bungee.protocol.DefinedPacket;
 import net.md_5.bungee.protocol.Protocol;
-import net.md_5.bungee.protocol.packet.BossBar;
 import net.md_5.bungee.protocol.packet.EncryptionRequest;
 import net.md_5.bungee.protocol.packet.Handshake;
 import net.md_5.bungee.protocol.packet.Kick;
@@ -170,7 +169,7 @@ public class ServerConnector extends PacketHandler
             }
         }
 
-        for ( PluginMessage message : user.getPendingConnection().getRegisterMessages() )
+        for ( PluginMessage message : user.getPendingConnection().getRelayMessages() )
         {
             ch.write( message );
         }
diff --git a/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java b/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
index 3fe1d234..068374a0 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
@@ -72,7 +72,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     private LoginRequest loginRequest;
     private EncryptionRequest request;
     @Getter
-    private final List<PluginMessage> registerMessages = new BoundedArrayList<>( 128 );
+    private final List<PluginMessage> relayMessages = new BoundedArrayList<>( 128 );
     private State thisState = State.HANDSHAKE;
     private final Unsafe unsafe = new Unsafe()
     {
@@ -126,9 +126,9 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     public void handle(PluginMessage pluginMessage) throws Exception
     {
         // TODO: Unregister?
-        if ( pluginMessage.getTag().equals( "REGISTER" ) )
+        if ( PluginMessage.SHOULD_RELAY.apply( pluginMessage ) )
         {
-            registerMessages.add( pluginMessage );
+            relayMessages.add( pluginMessage );
         }
     }
 
diff --git a/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java b/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
index 7e1917fa..88b179d3 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
@@ -13,7 +13,6 @@ import net.md_5.bungee.api.event.TabCompleteEvent;
 import net.md_5.bungee.netty.ChannelWrapper;
 import net.md_5.bungee.netty.PacketHandler;
 import net.md_5.bungee.protocol.PacketWrapper;
-import net.md_5.bungee.protocol.ProtocolConstants;
 import net.md_5.bungee.protocol.packet.KeepAlive;
 import net.md_5.bungee.protocol.packet.Chat;
 import net.md_5.bungee.protocol.packet.PlayerListItem;
@@ -191,9 +190,9 @@ public class UpstreamBridge extends PacketHandler
         }
 
         // TODO: Unregister as well?
-        if ( pluginMessage.getTag().equals( "REGISTER" ) )
+        if ( PluginMessage.SHOULD_RELAY.apply( pluginMessage ) )
         {
-            con.getPendingConnection().getRegisterMessages().add( pluginMessage );
+            con.getPendingConnection().getRelayMessages().add( pluginMessage );
         }
     }