From dd1d2c326ff05dd227612b43a7dc72f7cafd47bd Mon Sep 17 00:00:00 2001
From: Nathan Poirier <nathan@poirier.io>
Date: Tue, 28 Jun 2016 23:00:49 -0500
Subject: [PATCH] Improve ServerKickEvent

ServerKickEvent traditionally will only fire if the server sends a kick packet. During a server shutdown, or server crash this event would not fire for most players. While ServerDisconnectEvent is fired it does not tell us if it was a kick/shutdown/server crash.  This improvement fires the ServerKickEvent for server caused disconnections, and adds a Cause to the kick event.

diff --git a/api/src/main/java/net/md_5/bungee/api/event/ServerKickEvent.java b/api/src/main/java/net/md_5/bungee/api/event/ServerKickEvent.java
index 0e1ef5c4..ee63732d 100644
--- a/api/src/main/java/net/md_5/bungee/api/event/ServerKickEvent.java
+++ b/api/src/main/java/net/md_5/bungee/api/event/ServerKickEvent.java
@@ -44,6 +44,12 @@ public class ServerKickEvent extends Event implements Cancellable
      * State in which the kick occured.
      */
     private State state;
+    // Waterfall start
+    /**
+     * Circumstances which led to the kick.
+     */
+    private Cause cause;
+    // Waterfall end
 
     public enum State
     {
@@ -51,6 +57,14 @@ public class ServerKickEvent extends Event implements Cancellable
         CONNECTING, CONNECTED, UNKNOWN;
     }
 
+    // Waterfall start
+    public enum Cause
+    {
+
+        SERVER, LOST_CONNECTION, EXCEPTION, UNKNOWN;
+    }
+    // Waterfall end
+
     @Deprecated
     public ServerKickEvent(ProxiedPlayer player, BaseComponent[] kickReasonComponent, ServerInfo cancelServer)
     {
@@ -63,14 +77,23 @@ public class ServerKickEvent extends Event implements Cancellable
         this( player, player.getServer().getInfo(), kickReasonComponent, cancelServer, state );
     }
 
+    // Waterfall start
+    @Deprecated
     public ServerKickEvent(ProxiedPlayer player, ServerInfo kickedFrom, BaseComponent[] kickReasonComponent, ServerInfo cancelServer, State state)
+    {
+        this( player, kickedFrom, kickReasonComponent, cancelServer, state, Cause.UNKNOWN );
+    }
+
+    public ServerKickEvent(ProxiedPlayer player, ServerInfo kickedFrom, BaseComponent[] kickReasonComponent, ServerInfo cancelServer, State state, Cause cause)
     {
         this.player = player;
         this.kickedFrom = kickedFrom;
         this.kickReasonComponent = kickReasonComponent;
         this.cancelServer = cancelServer;
         this.state = state;
+        this.cause = cause;
     }
+    // Waterfall end
 
     @Deprecated
     public String getKickReason()
diff --git a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
index 972c70ec..1213c206 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
@@ -315,7 +315,7 @@ public class ServerConnector extends PacketHandler
     public void handle(Kick kick) throws Exception
     {
         ServerInfo def = user.updateAndGetNextServer( target );
-        ServerKickEvent event = new ServerKickEvent( user, target, ComponentSerializer.parse( kick.getMessage() ), def, ServerKickEvent.State.CONNECTING );
+        ServerKickEvent event = new ServerKickEvent( user, target, ComponentSerializer.parse( kick.getMessage() ), def, ServerKickEvent.State.CONNECTING, ServerKickEvent.Cause.SERVER );  // Waterfall
         if ( event.getKickReason().toLowerCase().contains( "outdated" ) && def != null )
         {
             // Pre cancel the event if we are going to try another server
diff --git a/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java b/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java
index 085a4e60..c6da6c47 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java
@@ -4,6 +4,7 @@ import com.google.common.base.Preconditions;
 import com.google.common.io.ByteArrayDataOutput;
 import com.google.common.io.ByteStreams;
 import java.io.DataInput;
+import java.util.Objects; // Waterfall
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.Unpooled;
@@ -60,16 +61,19 @@ public class DownstreamBridge extends PacketHandler
             return;
         }
 
+        // Waterfall start
         ServerInfo def = con.updateAndGetNextServer( server.getInfo() );
-        if ( def != null )
+        ServerKickEvent event = bungee.getPluginManager().callEvent( new ServerKickEvent( con, server.getInfo(), TextComponent.fromLegacyText( bungee.getTranslation( "server_went_down" ) ), def, ServerKickEvent.State.CONNECTED, ServerKickEvent.Cause.EXCEPTION ) );
+        if ( event.isCancelled() && event.getCancelServer() != null )
         {
             server.setObsolete( true );
-            con.connectNow( def, ServerConnectEvent.Reason.SERVER_DOWN_REDIRECT );
-            con.sendMessage( bungee.getTranslation( "server_went_down" ) );
-        } else
+            con.connectNow( event.getCancelServer(), ServerConnectEvent.Reason.SERVER_DOWN_REDIRECT );
+        }
+        else
         {
-            con.disconnect( Util.exception( t ) );
+            con.disconnect0( event.getKickReasonComponent() );
         }
+        // Waterfall end
     }
 
     @Override
@@ -84,7 +88,19 @@ public class DownstreamBridge extends PacketHandler
 
         if ( !server.isObsolete() )
         {
-            con.disconnect( bungee.getTranslation( "lost_connection" ) );
+            // Waterfall start
+            ServerInfo def = con.updateAndGetNextServer( server.getInfo() );
+            ServerKickEvent event = bungee.getPluginManager().callEvent( new ServerKickEvent( con, server.getInfo(), TextComponent.fromLegacyText( bungee.getTranslation( "lost_connection" ) ), def, ServerKickEvent.State.CONNECTED, ServerKickEvent.Cause.LOST_CONNECTION ) );
+            if ( event.isCancelled() && event.getCancelServer() != null )
+            {
+                server.setObsolete( true );
+                con.connectNow( event.getCancelServer() );
+            }
+            else
+            {
+                con.disconnect0( event.getKickReasonComponent() );
+            }
+            // Waterfall end
         }
 
         ServerDisconnectEvent serverDisconnectEvent = new ServerDisconnectEvent( con, server.getInfo() );
@@ -455,7 +471,11 @@ public class DownstreamBridge extends PacketHandler
     public void handle(Kick kick) throws Exception
     {
         ServerInfo def = con.updateAndGetNextServer( server.getInfo() );
-        ServerKickEvent event = bungee.getPluginManager().callEvent( new ServerKickEvent( con, server.getInfo(), ComponentSerializer.parse( kick.getMessage() ), def, ServerKickEvent.State.CONNECTED ) );
+        if ( Objects.equals( server.getInfo(), def ) )
+        {
+            def = null;
+        }
+        ServerKickEvent event = bungee.getPluginManager().callEvent( new ServerKickEvent( con, server.getInfo(), ComponentSerializer.parse( kick.getMessage() ), def, ServerKickEvent.State.CONNECTED, ServerKickEvent.Cause.SERVER ) ); // Waterfall
         if ( event.isCancelled() && event.getCancelServer() != null )
         {
             con.connectNow( event.getCancelServer(), ServerConnectEvent.Reason.KICK_REDIRECT );
-- 
2.16.3