From fb946123152ed08c2eeb1ad5b9cafe1bf97d6a6d Mon Sep 17 00:00:00 2001
From: md_5 <md_5@live.com.au>
Date: Fri, 2 Aug 2013 19:31:46 +1000
Subject: [PATCH] Fix throttle to 1) Work, 2) Not throttle outbound connections

---
 proxy/src/main/java/net/md_5/bungee/BungeeCord.java        | 7 ++++---
 .../java/net/md_5/bungee/connection/InitialHandler.java    | 2 +-
 proxy/src/main/java/net/md_5/bungee/netty/HandlerBoss.java | 7 +------
 .../src/main/java/net/md_5/bungee/netty/PipelineUtils.java | 7 +++++++
 4 files changed, 13 insertions(+), 10 deletions(-)

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 02bb0a1d..ef68aa08 100644
--- a/proxy/src/main/java/net/md_5/bungee/BungeeCord.java
+++ b/proxy/src/main/java/net/md_5/bungee/BungeeCord.java
@@ -18,6 +18,7 @@ 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.net.SocketAddress;
 import java.util.Calendar;
@@ -194,9 +195,9 @@ public class BungeeCord extends ProxyServer
             }
         }
     }
-    private final Map<SocketAddress, Long> throttle = new HashMap<>();
+    private final Map<InetAddress, Long> throttle = new HashMap<>();
 
-    public void unThrottle(SocketAddress address)
+    public void unThrottle(InetAddress address)
     {
         if ( address != null )
         {
@@ -207,7 +208,7 @@ public class BungeeCord extends ProxyServer
         }
     }
 
-    public boolean throttle(SocketAddress address)
+    public boolean throttle(InetAddress address)
     {
         long currentTime = System.currentTimeMillis();
         synchronized ( throttle )
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 dd6b7f4c..631c2a3c 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() );
+            BungeeCord.getInstance().unThrottle( getAddress().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 67389b4c..a6bd5b23 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
@@ -6,6 +6,7 @@ import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.MessageList;
 import io.netty.handler.timeout.ReadTimeoutException;
 import java.io.IOException;
+import java.net.InetSocketAddress;
 import java.util.logging.Level;
 import net.md_5.bungee.BungeeCord;
 import net.md_5.bungee.api.ProxyServer;
@@ -33,12 +34,6 @@ 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 );
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 ae432017..423e49cb 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
@@ -7,6 +7,7 @@ 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;
 import java.util.concurrent.TimeUnit;
 import net.md_5.bungee.BungeeCord;
 import net.md_5.bungee.BungeeServerInfo;
@@ -28,6 +29,12 @@ public class PipelineUtils
         @Override
         protected void initChannel(Channel ch) throws Exception
         {
+            if ( BungeeCord.getInstance().throttle( ( (InetSocketAddress) ch.remoteAddress() ).getAddress() ) )
+            {
+                ch.close();
+                return;
+            }
+
             BASE.initChannel( ch );
             ch.pipeline().get( HandlerBoss.class ).setHandler( new InitialHandler( ProxyServer.getInstance(), ch.attr( LISTENER ).get() ) );
         }