From c948a81d4f6604ea281baa2f9a96ac9366ea9e44 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 3f9efaa8..5d2597ad 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,13 @@ 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)
     {
@@ -69,14 +82,29 @@ public class ServerKickEvent extends Event implements Cancellable
         this( player, kickedFrom, TextComponent.fromArray( kickReasonComponent ), cancelServer, state );
     }
 
-    public ServerKickEvent(ProxiedPlayer player, ServerInfo kickedFrom, BaseComponent reason, ServerInfo cancelServer, State 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[] reason, ServerInfo cancelServer, State state, Cause cause)
+    {
+        this(player, kickedFrom, TextComponent.fromArray(reason), cancelServer, state, cause);
+    }
+
+    public ServerKickEvent(ProxiedPlayer player, ServerInfo kickedFrom, BaseComponent reason, ServerInfo cancelServer, State state, Cause cause)
     {
         this.player = player;
         this.kickedFrom = kickedFrom;
         this.reason = reason;
         this.cancelServer = cancelServer;
         this.state = state;
+        this.cause = cause;
     }
+    // Waterfall end
 
     /**
      * @return the kick reason
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 e670a57a..307a417b 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
@@ -439,7 +439,7 @@ public class ServerConnector extends PacketHandler
         ServerKickEvent event = new ServerKickEvent( user, target, new BaseComponent[]
         {
             kick.getMessage()
-        }, def, ServerKickEvent.State.CONNECTING );
+        }, def, ServerKickEvent.State.CONNECTING, ServerKickEvent.Cause.SERVER );  // Waterfall );
         if ( event.getKickReason().toLowerCase( Locale.ROOT ).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 cd3ed518..d4036c5d 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
@@ -98,16 +98,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.getReason() );
         }
+        // Waterfall end
     }
 
     @Override
@@ -122,7 +125,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.getReason() );
+            }
+            // Waterfall end
         }
 
         ServerDisconnectEvent serverDisconnectEvent = new ServerDisconnectEvent( con, server.getInfo() );
@@ -633,10 +648,14 @@ public class DownstreamBridge extends PacketHandler
     public void handle(Kick kick) throws Exception
     {
         ServerInfo def = con.updateAndGetNextServer( server.getInfo() );
+        if ( java.util.Objects.equals( server.getInfo(), def ) )
+        {
+            def = null;
+        }
         ServerKickEvent event = bungee.getPluginManager().callEvent( new ServerKickEvent( con, server.getInfo(), new BaseComponent[]
         {
             kick.getMessage()
-        }, def, ServerKickEvent.State.CONNECTED ) );
+        }, def, ServerKickEvent.State.CONNECTED, ServerKickEvent.Cause.SERVER ) ); // Waterfall
         if ( event.isCancelled() && event.getCancelServer() != null )
         {
             con.connectNow( event.getCancelServer(), ServerConnectEvent.Reason.KICK_REDIRECT );
-- 
2.44.0