Simpler, unit tested throttle to close #613

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

@@ -19,13 +19,11 @@ import net.md_5.bungee.config.Configuration;
 import java.io.File;
 import java.io.IOException;
 import java.io.PrintStream;
-import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.text.SimpleDateFormat;
 import java.util.Calendar;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.MissingResourceException;
@@ -120,6 +118,8 @@ public class BungeeCord extends ProxyServer
     @Getter
     private final Logger logger;
     public final Gson gson = new Gson();
+    @Getter
+    private ConnectionThrottle connectionThrottle;
 
     
     {
@@ -199,35 +199,6 @@ public class BungeeCord extends ProxyServer
             }
         }
     }
-    private final Map<InetAddress, Long> throttle = new HashMap<>();
-
-    public void unThrottle(InetAddress address)
-    {
-        if ( address != null )
-        {
-            synchronized ( throttle )
-            {
-                throttle.remove( address );
-            }
-        }
-    }
-
-    public boolean throttle(InetAddress address)
-    {
-        long currentTime = System.currentTimeMillis();
-        synchronized ( throttle )
-        {
-            Long value = throttle.get( address );
-            if ( value != null && currentTime - value < config.getThrottle() )
-            {
-                throttle.put( address, currentTime );
-                return true;
-            }
-
-            throttle.put( address, currentTime );
-        }
-        return false;
-    }
 
     /**
      * Start this proxy instance by loading the configuration, plugins and
@@ -255,6 +226,7 @@ public class BungeeCord extends ProxyServer
 
         pluginManager.loadAndEnablePlugins();
 
+        connectionThrottle = new ConnectionThrottle( config.getThrottle() );
         startListeners();
 
         saveThread.scheduleAtFixedRate( new TimerTask()
@@ -367,7 +339,7 @@ public class BungeeCord extends ProxyServer
                     try
                     {
                         plugin.onDisable();
-                    } catch (Throwable t)
+                    } catch ( Throwable t )
                     {
                         getLogger().severe( "Exception disabling plugin " + plugin.getDescription().getName() );
                         t.printStackTrace();

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

@@ -0,0 +1,27 @@
+package net.md_5.bungee;
+
+import gnu.trove.map.hash.TObjectLongHashMap;
+import java.net.InetAddress;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+public class ConnectionThrottle
+{
+
+    private final TObjectLongHashMap<InetAddress> throttle = new TObjectLongHashMap<>();
+    private final int throttleTime;
+
+    public void unthrottle(InetAddress address)
+    {
+        throttle.remove( address );
+    }
+
+    public boolean throttle(InetAddress address)
+    {
+        long value = throttle.get( address );
+        long currentTime = System.currentTimeMillis();
+
+        throttle.put( address, currentTime );
+        return value != 0 && currentTime - value < throttleTime;
+    }
+}

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

@@ -153,7 +153,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
                     + "\00" + response.getMotd()
                     + "\00" + response.getCurrentPlayers()
                     + "\00" + response.getMaxPlayers();
-            BungeeCord.getInstance().unThrottle( getAddress().getAddress() );
+            BungeeCord.getInstance().getConnectionThrottle().unthrottle( getAddress().getAddress() );
             disconnect( kickMessage );
         } catch ( Throwable t )
         {

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

@@ -4,7 +4,6 @@ import io.netty.channel.Channel;
 import io.netty.channel.ChannelException;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOption;
-import io.netty.handler.codec.bytes.ByteArrayEncoder;
 import io.netty.handler.timeout.ReadTimeoutHandler;
 import io.netty.util.AttributeKey;
 import java.net.InetSocketAddress;
@@ -29,7 +28,7 @@ public class PipelineUtils
         @Override
         protected void initChannel(Channel ch) throws Exception
         {
-            if ( BungeeCord.getInstance().throttle( ( (InetSocketAddress) ch.remoteAddress() ).getAddress() ) )
+            if ( BungeeCord.getInstance().getConnectionThrottle().throttle( ( (InetSocketAddress) ch.remoteAddress() ).getAddress() ) )
             {
                 ch.close();
                 return;

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

@@ -0,0 +1,26 @@
+package net.md_5.bungee;
+
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class ThrottleTest
+{
+
+    @Test
+    public void testThrottle() throws InterruptedException, UnknownHostException
+    {
+        ConnectionThrottle throttle = new ConnectionThrottle( 5 );
+        InetAddress address = InetAddress.getLocalHost();
+
+        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 ) );
+    }
+}