#3207: Rework the plugin message relaying system to allow unregistering channels

protocol/src/main/java/net/md_5/bungee/protocol/packet/PluginMessage.java

@@ -2,7 +2,6 @@ package net.md_5.bungee.protocol.packet;
 
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Predicate;
 import io.netty.buffer.ByteBuf;
 import java.io.ByteArrayInputStream;
 import java.io.DataInput;
@@ -48,14 +47,6 @@ public class PluginMessage extends DefinedPacket
             return "legacy:" + tag.toLowerCase( Locale.ROOT );
         }
     };
-    public static final Predicate<PluginMessage> SHOULD_RELAY = new Predicate<PluginMessage>()
-    {
-        @Override
-        public boolean apply(PluginMessage input)
-        {
-            return ( input.getTag().equals( "REGISTER" ) || input.getTag().equals( "minecraft:register" ) || input.getTag().equals( "MC|Brand" ) || input.getTag().equals( "minecraft:brand" ) ) && input.getData().length < Byte.MAX_VALUE;
-        }
-    };
     //
     private String tag;
     private byte[] data;

proxy/src/main/java/net/md_5/bungee/ServerConnector.java

@@ -1,9 +1,11 @@
 package net.md_5.bungee;
 
+import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
 import java.util.Locale;
 import java.util.Queue;
 import java.util.Set;
@@ -193,9 +195,16 @@ public class ServerConnector extends PacketHandler
             }
         }
 
-        for ( PluginMessage message : user.getPendingConnection().getRelayMessages() )
+        PluginMessage brandMessage = user.getPendingConnection().getBrandMessage();
+        if ( brandMessage != null )
         {
-            ch.write( message );
+            ch.write( brandMessage );
+        }
+
+        Set<String> registeredChannels = user.getPendingConnection().getRegisteredChannels();
+        if ( !registeredChannels.isEmpty() )
+        {
+            ch.write( new PluginMessage( user.getPendingConnection().getVersion() >= ProtocolConstants.MINECRAFT_1_13 ? "minecraft:register" : "REGISTER", Joiner.on( "\0" ).join( registeredChannels ).getBytes( StandardCharsets.UTF_8 ), false ) );
         }
 
         if ( user.getSettings() != null )

proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java

@@ -7,8 +7,10 @@ import java.math.BigInteger;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
 import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
-import java.util.List;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.UUID;
 import java.util.logging.Level;
 import javax.crypto.SecretKey;
@@ -63,7 +65,6 @@ import net.md_5.bungee.protocol.packet.PluginMessage;
 import net.md_5.bungee.protocol.packet.StatusRequest;
 import net.md_5.bungee.protocol.packet.StatusResponse;
 import net.md_5.bungee.util.AllowedCharacters;
-import net.md_5.bungee.util.BoundedArrayList;
 import net.md_5.bungee.util.BufUtil;
 import net.md_5.bungee.util.QuietException;
 
@@ -81,7 +82,9 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     private LoginRequest loginRequest;
     private EncryptionRequest request;
     @Getter
-    private final List<PluginMessage> relayMessages = new BoundedArrayList<>( 128 );
+    private PluginMessage brandMessage;
+    @Getter
+    private final Set<String> registeredChannels = new HashSet<>();
     private State thisState = State.HANDSHAKE;
     private final Unsafe unsafe = new Unsafe()
     {
@@ -154,11 +157,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     @Override
     public void handle(PluginMessage pluginMessage) throws Exception
     {
-        // TODO: Unregister?
+        this.relayMessage( pluginMessage );
-        if ( PluginMessage.SHOULD_RELAY.apply( pluginMessage ) )
-        {
-            relayMessages.add( pluginMessage );
-        }
     }
 
     @Override
@@ -668,4 +667,31 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     {
         return !ch.isClosed();
     }
+
+    public void relayMessage(PluginMessage input) throws Exception
+    {
+        if ( input.getTag().equals( "REGISTER" ) || input.getTag().equals( "minecraft:register" ) )
+        {
+            String content = new String( input.getData(), StandardCharsets.UTF_8 );
+
+            for ( String id : content.split( "\0" ) )
+            {
+                Preconditions.checkState( registeredChannels.size() < 128, "Too many registered channels" );
+                Preconditions.checkArgument( id.length() < 128, "Channel name too long" );
+
+                registeredChannels.add( id );
+            }
+        } else if ( input.getTag().equals( "UNREGISTER" ) || input.getTag().equals( "minecraft:unregister" ) )
+        {
+            String content = new String( input.getData(), StandardCharsets.UTF_8 );
+
+            for ( String id : content.split( "\0" ) )
+            {
+                registeredChannels.remove( id );
+            }
+        } else if ( input.getTag().equals( "MC|Brand" ) || input.getTag().equals( "minecraft:brand" ) )
+        {
+            brandMessage = input;
+        }
+    }
 }

proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java

@@ -265,11 +265,7 @@ public class UpstreamBridge extends PacketHandler
             throw CancelSendSignal.INSTANCE;
         }
 
-        // TODO: Unregister as well?
+        con.getPendingConnection().relayMessage( pluginMessage );
-        if ( PluginMessage.SHOULD_RELAY.apply( pluginMessage ) )
-        {
-            con.getPendingConnection().getRelayMessages().add( pluginMessage );
-        }
     }
 
     @Override

proxy/src/main/java/net/md_5/bungee/util/BoundedArrayList.java

@@ -1,49 +0,0 @@
-package net.md_5.bungee.util;
-
-import com.google.common.base.Preconditions;
-import java.util.ArrayList;
-import java.util.Collection;
-
-public class BoundedArrayList<E> extends ArrayList<E>
-{
-
-    private final int maxSize;
-
-    public BoundedArrayList(int maxSize)
-    {
-        this.maxSize = maxSize;
-    }
-
-    private void checkSize(int increment)
-    {
-        Preconditions.checkState( size() + increment <= maxSize, "Adding %s elements would exceed capacity of %s", increment, maxSize );
-    }
-
-    @Override
-    public boolean add(E e)
-    {
-        checkSize( 1 );
-        return super.add( e );
-    }
-
-    @Override
-    public void add(int index, E element)
-    {
-        checkSize( 1 );
-        super.add( index, element );
-    }
-
-    @Override
-    public boolean addAll(Collection<? extends E> c)
-    {
-        checkSize( c.size() );
-        return super.addAll( c );
-    }
-
-    @Override
-    public boolean addAll(int index, Collection<? extends E> c)
-    {
-        checkSize( c.size() );
-        return super.addAll( index, c );
-    }
-}

proxy/src/test/java/net/md_5/bungee/util/BoundedArrayListTest.java

@@ -1,76 +0,0 @@
-package net.md_5.bungee.util;
-
-import com.google.common.collect.ImmutableList;
-import org.junit.Test;
-
-public class BoundedArrayListTest
-{
-
-    @Test
-    public void testGoodAdd() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 2 );
-        list.add( new Object() );
-        list.add( new Object() );
-    }
-
-    @Test
-    public void testSizeOneAdd() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 1 );
-        list.add( new Object() );
-    }
-
-    @Test(expected = IllegalStateException.class)
-    public void testBadAdd() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 0 );
-        list.add( new Object() );
-    }
-
-    @Test
-    public void testGoodAdd1() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 2 );
-        list.add( new Object() );
-        list.add( 0, new Object() );
-    }
-
-    @Test(expected = IllegalStateException.class)
-    public void testBadAdd1() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 1 );
-        list.add( new Object() );
-        list.add( 0, new Object() );
-    }
-
-    @Test
-    public void testGoodAddAll() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 1 );
-        list.addAll( ImmutableList.of( new Object() ) );
-    }
-
-    @Test
-    public void testGoodAddAll1() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 2 );
-        list.add( new Object() );
-        list.addAll( 0, ImmutableList.of( new Object() ) );
-    }
-
-    @Test(expected = IllegalStateException.class)
-    public void testBadAddAll() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 0 );
-        list.addAll( ImmutableList.of( new Object() ) );
-    }
-
-    @Test(expected = IllegalStateException.class)
-    public void testBadAddAll1() throws Exception
-    {
-        BoundedArrayList<Object> list = new BoundedArrayList<>( 1 );
-        list.add( new Object() );
-        list.addAll( ImmutableList.of( new Object() ) );
-    }
-}