From 1f3815253060bfc45341d3d1ae89f6c4ce58a48d Mon Sep 17 00:00:00 2001
From: md_5 <md_5@live.com.au>
Date: Thu, 1 Aug 2013 13:37:32 +1000
Subject: [PATCH] [URGENT] Add connection throttle.

---
 .../main/java/net/md_5/bungee/BungeeCord.java | 31 +++++++++++++++++++
 .../net/md_5/bungee/config/Configuration.java |  2 ++
 .../bungee/connection/InitialHandler.java     |  1 +
 .../net/md_5/bungee/netty/HandlerBoss.java    |  7 +++++
 4 files changed, 41 insertions(+)

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 853948c1..c175f9e4 100644
--- a/proxy/src/main/java/net/md_5/bungee/BungeeCord.java
+++ b/proxy/src/main/java/net/md_5/bungee/BungeeCord.java
@@ -19,9 +19,11 @@ import java.io.File;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 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;
@@ -192,6 +194,35 @@ public class BungeeCord extends ProxyServer
             }
         }
     }
+    private final Map<SocketAddress, Long> throttle = new HashMap<>();
+
+    public void unThrottle(SocketAddress address)
+    {
+        if ( address != null )
+        {
+            synchronized ( throttle )
+            {
+                throttle.remove( address );
+            }
+        }
+    }
+
+    public boolean throttle(SocketAddress 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
diff --git a/proxy/src/main/java/net/md_5/bungee/config/Configuration.java b/proxy/src/main/java/net/md_5/bungee/config/Configuration.java
index 33855306..5b93a904 100644
--- a/proxy/src/main/java/net/md_5/bungee/config/Configuration.java
+++ b/proxy/src/main/java/net/md_5/bungee/config/Configuration.java
@@ -46,6 +46,7 @@ public class Configuration
     private boolean onlineMode = true;
     private int playerLimit = -1;
     private Collection<String> disabledCommands;
+    private int throttle = 4000;
 
     public void load()
     {
@@ -57,6 +58,7 @@ public class Configuration
         uuid = adapter.getString( "stats", uuid );
         onlineMode = adapter.getBoolean( "online_mode", onlineMode );
         playerLimit = adapter.getInt( "player_limit", playerLimit );
+        throttle = adapter.getInt( "connection_throttle", throttle );
 
         disabledCommands = new CaseInsensitiveSet( (Collection<String>) adapter.getList( "disabled_commands", Arrays.asList( "find" ) ) );
 
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 794243c7..20e115fe 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,6 +153,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
                     + "\00" + response.getMotd()
                     + "\00" + response.getCurrentPlayers()
                     + "\00" + response.getMaxPlayers();
+            BungeeCord.getInstance().unThrottle( getAddress() );
             disconnect( kickMessage );
         } catch ( Throwable t )
         {
diff --git a/proxy/src/main/java/net/md_5/bungee/netty/HandlerBoss.java b/proxy/src/main/java/net/md_5/bungee/netty/HandlerBoss.java
index 7095f18a..f54763a6 100644
--- a/proxy/src/main/java/net/md_5/bungee/netty/HandlerBoss.java
+++ b/proxy/src/main/java/net/md_5/bungee/netty/HandlerBoss.java
@@ -8,6 +8,7 @@ import java.io.IOException;
 import java.util.ArrayDeque;
 import java.util.Queue;
 import java.util.logging.Level;
+import net.md_5.bungee.BungeeCord;
 import net.md_5.bungee.api.ProxyServer;
 import net.md_5.bungee.connection.CancelSendSignal;
 import net.md_5.bungee.connection.InitialHandler;
@@ -34,6 +35,12 @@ public class HandlerBoss extends ChannelInboundHandlerAdapter
     @Override
     public void channelActive(ChannelHandlerContext ctx) throws Exception
     {
+        if ( BungeeCord.getInstance().throttle( ctx.channel().remoteAddress() ) )
+        {
+            ctx.channel().close();
+            return;
+        }
+
         if ( handler != null )
         {
             channel = new ChannelWrapper( ctx );