Reimplement join throttle.

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

@@ -1,25 +1,29 @@
 package net.md_5.bungee;
 
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import java.net.InetAddress;
-import java.util.HashMap;
+import java.util.concurrent.TimeUnit;
-import java.util.Map;
-import lombok.RequiredArgsConstructor;
 
-@RequiredArgsConstructor
 public class ConnectionThrottle
 {
 
-    private final Map<InetAddress, Long> throttle = new HashMap<>();
     private final int throttleTime;
+    private final Cache<InetAddress, Long> throttle;
 
-    public void unthrottle(InetAddress address)
+    public ConnectionThrottle(int throttleTime)
     {
-        throttle.remove( address );
+        this.throttleTime = throttleTime;
+        this.throttle = CacheBuilder.newBuilder()
+                .concurrencyLevel( Runtime.getRuntime().availableProcessors() )
+                .initialCapacity( 100 )
+                .expireAfterAccess( throttleTime, TimeUnit.MILLISECONDS )
+                .build();
     }
 
     public boolean throttle(InetAddress address)
     {
-        Long value = throttle.get( address );
+        Long value = throttle.getIfPresent( address );
         long currentTime = System.currentTimeMillis();
 
         throttle.put( address, currentTime );

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

@@ -40,6 +40,7 @@ import net.md_5.bungee.netty.PipelineUtils;
 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.PacketWrapper;
 import net.md_5.bungee.protocol.packet.Handshake;
 import net.md_5.bungee.protocol.packet.PluginMessage;
 import net.md_5.bungee.protocol.packet.EncryptionResponse;
@@ -64,7 +65,7 @@ import net.md_5.bungee.util.BoundedArrayList;
 public class InitialHandler extends PacketHandler implements PendingConnection
 {
 
-    private final ProxyServer bungee;
+    private final BungeeCord bungee;
     private ChannelWrapper ch;
     @Getter
     private final ListenerInfo listener;
@@ -98,6 +99,13 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     private boolean legacy;
     @Getter
     private String extraDataInHandshake = "";
+    private boolean disconnecting;
+
+    @Override
+    public boolean shouldHandle(PacketWrapper packet) throws Exception
+    {
+        return !disconnecting;
+    }
 
     private enum State
     {
@@ -213,7 +221,6 @@ public class InitialHandler extends PacketHandler implements PendingConnection
                     @Override
                     public void done(ProxyPingEvent pingResult, Throwable error)
                     {
-                        BungeeCord.getInstance().getConnectionThrottle().unthrottle( getAddress().getAddress() );
                         Gson gson = BungeeCord.getInstance().gson;
                         unsafe.sendPacket( new StatusResponse( gson.toJson( pingResult.getResponse() ) ) );
                     }
@@ -285,9 +292,14 @@ public class InitialHandler extends PacketHandler implements PendingConnection
                 ch.setProtocol( Protocol.STATUS );
                 break;
             case 2:
+                // Login
                 thisState = State.USERNAME;
                 ch.setProtocol( Protocol.LOGIN );
-                // Login
+
+                if ( bungee.getConnectionThrottle().throttle( ( (InetSocketAddress) ch.getHandle().remoteAddress() ).getAddress() ) )
+                {
+                    disconnect( bungee.getTranslation( "join_throttle_kick", TimeUnit.MILLISECONDS.toSeconds( bungee.getConfig().getThrottle() ) ) );
+                }
                 break;
             default:
                 throw new IllegalArgumentException( "Cannot request protocol " + handshake.getRequestedProtocol() );
@@ -519,8 +531,9 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     @Override
     public void disconnect(final BaseComponent... reason)
     {
-        if ( !ch.isClosed() )
+        if ( !disconnecting || !ch.isClosed() )
         {
+            disconnecting = true;
             // Why do we have to delay this you might ask? Well the simple reason is MOJANG.
             // Despite many a bug report posted, ever since the 1.7 protocol rewrite, the client STILL has a race condition upon switching protocols.
             // As such, despite the protocol switch packets already having been sent, there is the possibility of a client side exception

proxy/src/main/java/net/md_5/bungee/netty/PipelineUtils.java

@@ -49,19 +49,12 @@ public class PipelineUtils
         @Override
         protected void initChannel(Channel ch) throws Exception
         {
-            if ( BungeeCord.getInstance().getConnectionThrottle().throttle( ( (InetSocketAddress) ch.remoteAddress() ).getAddress() ) )
-            {
-                // TODO: Better throttle - we can't throttle this way if we want to maintain 1.7 compat!
-                // ch.close();
-                // return;
-            }
-
             BASE.initChannel( ch );
             ch.pipeline().addBefore( FRAME_DECODER, LEGACY_DECODER, new LegacyDecoder() );
             ch.pipeline().addAfter( FRAME_DECODER, PACKET_DECODER, new MinecraftDecoder( Protocol.HANDSHAKE, true, ProxyServer.getInstance().getProtocolVersion() ) );
             ch.pipeline().addAfter( FRAME_PREPENDER, PACKET_ENCODER, new MinecraftEncoder( Protocol.HANDSHAKE, true, ProxyServer.getInstance().getProtocolVersion() ) );
             ch.pipeline().addBefore( FRAME_PREPENDER, LEGACY_KICKER, new KickStringWriter() );
-            ch.pipeline().get( HandlerBoss.class ).setHandler( new InitialHandler( ProxyServer.getInstance(), ch.attr( LISTENER ).get() ) );
+            ch.pipeline().get( HandlerBoss.class ).setHandler( new InitialHandler( BungeeCord.getInstance(), ch.attr( LISTENER ).get() ) );
         }
     };
     public static final Base BASE = new Base();

proxy/src/main/resources/messages.properties

@@ -24,3 +24,4 @@ total_players=Total players online: {0}
 name_too_long=Cannot have username longer than 16 characters
 name_invalid=Username contains invalid characters.
 ping_cannot_connect=\u00a7c[Bungee] Can't connect to server.
+join_throttle_kick=You have connected too fast. You must wait at least {0} seconds between connections.

proxy/src/test/java/net/md_5/bungee/ThrottleTest.java

@@ -25,9 +25,6 @@ public class ThrottleTest
         Assert.assertFalse( "Address should not be throttled", throttle.throttle( address ) );
         Assert.assertTrue( "Address should be throttled", throttle.throttle( address ) );
 
-        throttle.unthrottle( address );
-        Assert.assertFalse( "Address should not be throttled", throttle.throttle( address ) );
-
         Thread.sleep( 15 );
         Assert.assertFalse( "Address should not be throttled", throttle.throttle( address ) );
     }