From 1bbbfdb0d5ae87444a9f92c4f175846d04678d98 Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Thu, 20 Mar 2014 20:11:56 +1100
Subject: [PATCH] Close #938 - re add our register channel manipulation

---
 .../java/net/md_5/bungee/ServerConnector.java    |  5 +++++
 .../md_5/bungee/connection/InitialHandler.java   | 16 ++++++++++++++++
 .../md_5/bungee/connection/UpstreamBridge.java   |  6 ++++++
 3 files changed, 27 insertions(+)

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 a9dfcabd..7cd5cf14 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
@@ -112,6 +112,11 @@ public class ServerConnector extends PacketHandler
             }
         }
 
+        for ( PluginMessage message : user.getPendingConnection().getRegisterMessages() )
+        {
+            ch.write( message );
+        }
+
         if ( user.getSettings() != null )
         {
             ch.write( user.getSettings() );
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 f4ae9a31..b66cf806 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
@@ -6,6 +6,8 @@ import java.math.BigInteger;
 import java.net.InetSocketAddress;
 import java.net.URLEncoder;
 import java.security.MessageDigest;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.logging.Level;
 import javax.crypto.SecretKey;
 import lombok.Getter;
@@ -34,6 +36,7 @@ import net.md_5.bungee.netty.cipher.CipherDecoder;
 import net.md_5.bungee.netty.cipher.CipherEncoder;
 import net.md_5.bungee.protocol.DefinedPacket;
 import net.md_5.bungee.protocol.packet.Handshake;
+import net.md_5.bungee.protocol.packet.PluginMessage;
 import net.md_5.bungee.protocol.packet.EncryptionResponse;
 import net.md_5.bungee.protocol.packet.EncryptionRequest;
 import net.md_5.bungee.protocol.packet.Kick;
@@ -62,6 +65,8 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     @Getter
     private LoginRequest loginRequest;
     private EncryptionRequest request;
+    @Getter
+    private final List<PluginMessage> registerMessages = new ArrayList<>();
     private State thisState = State.HANDSHAKE;
     private final Unsafe unsafe = new Unsafe()
     {
@@ -96,6 +101,17 @@ public class InitialHandler extends PacketHandler implements PendingConnection
         disconnect( ChatColor.RED + Util.exception( t ) );
     }
 
+    @Override
+    public void handle(PluginMessage pluginMessage) throws Exception
+    {
+        // TODO: Unregister?
+        if ( pluginMessage.getTag().equals( "REGISTER" ) )
+        {
+            Preconditions.checkState( registerMessages.size() < 128, "Too many channels registered" );
+            registerMessages.add( pluginMessage );
+        }
+    }
+
     @Override
     public void handle(LegacyHandshake legacyHandshake) throws Exception
     {
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 bec51bda..e5e402ff 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
@@ -147,6 +147,12 @@ public class UpstreamBridge extends PacketHandler
         {
             throw CancelSendSignal.INSTANCE;
         }
+
+        // TODO: Unregister as well?
+        if ( pluginMessage.getTag().equals( "REGISTER" ) )
+        {
+            con.getPendingConnection().getRegisterMessages().add( pluginMessage );
+        }
     }
 
     @Override