From 7578aa27a9073b30a56657fb2dd1f1c2d2f2b65f Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Tue, 9 Oct 2018 19:02:15 -0400
Subject: [PATCH] More reliably tear down connections and fire DisconnectEvent.
 Fixes #111

---
 .../client/ClientPlaySessionHandler.java      |  1 -
 .../connection/client/ConnectedPlayer.java    |  2 ++
 .../client/LoginSessionHandler.java           | 20 ++++++++++++++-----
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
index 42c9d226..9f7001bb 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
@@ -202,7 +202,6 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     @Override
     public void disconnected() {
         player.teardown();
-        server.getEventManager().fireAndForget(new DisconnectEvent(player));
     }
 
     @Override
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 b16c6487..4e3afb26 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
@@ -2,6 +2,7 @@ package com.velocitypowered.proxy.connection.client;
 
 import com.google.common.base.Preconditions;
 import com.google.gson.JsonObject;
+import com.velocitypowered.api.event.connection.DisconnectEvent;
 import com.velocitypowered.api.event.player.KickedFromServerEvent;
 import com.velocitypowered.api.event.player.PlayerSettingsChangedEvent;
 import com.velocitypowered.api.event.player.ServerPreConnectEvent;
@@ -410,6 +411,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
             connectedServer.disconnect();
         }
         server.unregisterConnection(this);
+        server.getEventManager().fireAndForget(new DisconnectEvent(this));
     }
 
     @Override
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/LoginSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/LoginSessionHandler.java
index 5f42779a..48a6859f 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/LoginSessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/LoginSessionHandler.java
@@ -50,6 +50,7 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
     private ServerLogin login;
     private byte[] verify;
     private int playerInfoId;
+    private ConnectedPlayer connectedPlayer;
 
     public LoginSessionHandler(VelocityServer server, MinecraftConnection inbound, InboundConnection apiInbound) {
         this.server = Preconditions.checkNotNull(server, "server");
@@ -196,6 +197,7 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
             // Initiate a regular connection and move over to it.
             ConnectedPlayer player = new ConnectedPlayer(server, profileEvent.getGameProfile(), inbound,
                     apiInbound.getVirtualHost().orElse(null));
+            this.connectedPlayer = player;
 
             return server.getEventManager().fire(new PermissionsSetupEvent(player, ConnectedPlayer.DEFAULT_PERMISSIONS))
                         .thenCompose(event -> {
@@ -212,7 +214,7 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
                             }
                             if (!event.getResult().isAllowed()) {
                                 // The component is guaranteed to be provided if the connection was denied.
-                                inbound.closeWith(Disconnect.create(event.getResult().getReason().get()));
+                                player.disconnect(event.getResult().getReason().get());
                                 return;
                             }
 
@@ -229,6 +231,11 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
             return;
         }
 
+        if (!server.registerConnection(player)) {
+            inbound.closeWith(Disconnect.create(TextComponent.of("You are already on this proxy!", TextColor.RED)));
+            return;
+        }
+
         int threshold = server.getConfiguration().getCompressionThreshold();
         if (threshold >= 0) {
             inbound.write(new SetCompression(threshold));
@@ -243,10 +250,6 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
         inbound.setAssociation(player);
         inbound.setState(StateRegistry.PLAY);
 
-        if (!server.registerConnection(player)) {
-            inbound.closeWith(Disconnect.create(TextComponent.of("You are already on this proxy!", TextColor.RED)));
-        }
-
         logger.info("{} has connected", player);
         inbound.setSessionHandler(new InitialConnectSessionHandler(player));
         server.getEventManager().fire(new PostLoginEvent(player)).thenRun(() -> player.createConnectionRequest(toTry.get()).fireAndForget());
@@ -256,4 +259,11 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
     public void handleUnknown(ByteBuf buf) {
         throw new IllegalStateException("Unknown data " + ByteBufUtil.hexDump(buf));
     }
+
+    @Override
+    public void disconnected() {
+        if (connectedPlayer != null) {
+            connectedPlayer.teardown();
+        }
+    }
 }