From 826eddc754866c326206cd091a0449a0df90881c Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sun, 28 Oct 2018 23:00:13 -0400
Subject: [PATCH] Disconnect obsolete server connections as quickly as
 possible.

---
 .../backend/BackendPlaySessionHandler.java          |  4 ++--
 .../connection/backend/LoginSessionHandler.java     |  6 ++++--
 .../backend/VelocityServerConnection.java           | 13 ++++++++++++-
 .../proxy/connection/client/ConnectedPlayer.java    |  8 ++++++--
 4 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
index c7c1a281f..8f87e593b 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
@@ -45,7 +45,7 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
 
   @Override
   public boolean beforeHandle() {
-    if (!serverConn.getPlayer().isActive()) {
+    if (!serverConn.isActive()) {
       // Obsolete connection
       serverConn.disconnect();
       return true;
@@ -121,7 +121,7 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
         packet.getData());
     server.getEventManager().fire(event)
         .thenAcceptAsync(pme -> {
-          if (pme.getResult().isAllowed()) {
+          if (pme.getResult().isAllowed() && serverConn.isActive()) {
             smc.write(packet);
           }
         }, smc.eventLoop());
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/LoginSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/LoginSessionHandler.java
index f04b4d869..4ab39ecca 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/LoginSessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/LoginSessionHandler.java
@@ -117,11 +117,13 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
       serverConn.getPlayer().getConnection()
           .setSessionHandler(new ClientPlaySessionHandler(server, serverConn.getPlayer()));
     } else {
-      // The previous server connection should become obsolete.
-      // Before we remove it, if the server we are departing is modded, we must always reset the client state.
+      // If the server we are departing is modded, we must always reset the client's handshake.
       if (existingConnection.isLegacyForge()) {
         serverConn.getPlayer().sendLegacyForgeHandshakeResetPacket();
       }
+
+      // Shut down the existing server connection.
+      serverConn.getPlayer().setConnectedServer(null);
       existingConnection.disconnect();
     }
 
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
index 7f275e8f3..7ee3804b6 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
@@ -165,9 +165,9 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
 
   public void disconnect() {
     if (connection != null) {
+      gracefulDisconnect = true;
       connection.close();
       connection = null;
-      gracefulDisconnect = true;
     }
   }
 
@@ -230,4 +230,15 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
   public void resetLastPingId() {
     this.lastPingId = -1;
   }
+
+  /**
+   * Ensures that this server connection remains "active": the connection is established and not
+   * closed, the player is still connected to the server, and the player still remains online.
+   *
+   * @return whether or not the player is online
+   */
+  boolean isActive() {
+    return connection != null && !connection.isClosed() && !gracefulDisconnect
+        && proxyPlayer.isActive();
+  }
 }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
index 2c823a00f..71f13b705 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
@@ -443,8 +443,13 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
         });
   }
 
-  public void setConnectedServer(VelocityServerConnection serverConnection) {
+  public void setConnectedServer(@Nullable VelocityServerConnection serverConnection) {
     VelocityServerConnection oldConnection = this.connectedServer;
+    this.connectedServer = serverConnection;
+
+    if (serverConnection == null) {
+      return;
+    }
     if (oldConnection != null && !serverConnection.getServerInfo()
         .equals(oldConnection.getServerInfo())) {
       this.tryIndex = 0;
@@ -452,7 +457,6 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
     if (serverConnection == connectionInFlight) {
       connectionInFlight = null;
     }
-    this.connectedServer = serverConnection;
   }
 
   public void sendLegacyForgeHandshakeResetPacket() {