From af58db8a677a61b5b9b6cc00bc34a95bc010dafb Mon Sep 17 00:00:00 2001 From: md_5 Date: Tue, 10 Sep 2013 12:02:29 +1000 Subject: [PATCH] Simpler, unit tested throttle to close #613 --- .../main/java/net/md_5/bungee/BungeeCord.java | 36 +++---------------- .../net/md_5/bungee/ConnectionThrottle.java | 27 ++++++++++++++ .../bungee/connection/InitialHandler.java | 2 +- .../net/md_5/bungee/netty/PipelineUtils.java | 3 +- .../java/net/md_5/bungee/ThrottleTest.java | 26 ++++++++++++++ 5 files changed, 59 insertions(+), 35 deletions(-) create mode 100644 proxy/src/main/java/net/md_5/bungee/ConnectionThrottle.java create mode 100644 proxy/src/test/java/net/md_5/bungee/ThrottleTest.java diff --git a/proxy/src/main/java/net/md_5/bungee/BungeeCord.java b/proxy/src/main/java/net/md_5/bungee/BungeeCord.java index 171ec5ac..d0dccb9e 100644 --- a/proxy/src/main/java/net/md_5/bungee/BungeeCord.java +++ b/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 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(); diff --git a/proxy/src/main/java/net/md_5/bungee/ConnectionThrottle.java b/proxy/src/main/java/net/md_5/bungee/ConnectionThrottle.java new file mode 100644 index 00000000..e3c45fcd --- /dev/null +++ b/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 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; + } +} 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 ae8d3aff..a01d3fd8 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 @@ -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 ) { diff --git a/proxy/src/main/java/net/md_5/bungee/netty/PipelineUtils.java b/proxy/src/main/java/net/md_5/bungee/netty/PipelineUtils.java index b98f4df0..4ea8edd0 100644 --- a/proxy/src/main/java/net/md_5/bungee/netty/PipelineUtils.java +++ b/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; diff --git a/proxy/src/test/java/net/md_5/bungee/ThrottleTest.java b/proxy/src/test/java/net/md_5/bungee/ThrottleTest.java new file mode 100644 index 00000000..931fdb8a --- /dev/null +++ b/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 ) ); + } +}