Ensure sessions are removed when a transfer packet is sent or they exit without the form

src/main/java/org/geysermc/extension/connect/GeyserConnect.java

@@ -109,9 +109,11 @@ public void onPostInitialize(GeyserPostInitializeEvent event) {
             this.logger().warning("Either `passthrough-motd` or `passthrough-player-counts` is enabled in the config, this will likely produce errors");
         }
 
-        // If we are using floodgate then disable the extension
-        if (geyserInstance.getConfig().getRemote().authType() == AuthType.FLOODGATE) {
-            this.logger().error("auth-type set to floodgate in the config, this will break GeyserConnect. Disabling!");
+        // If we are using floodgate then disable the extension.
+        // GeyserConnect also doesn't support the connection sequence that occurs when the default RemoteServer
+        // auth-type is offline (and there is no reason to change it when GeyserConnect is in use).
+        if (geyserInstance.getConfig().getRemote().authType() != AuthType.ONLINE) {
+            this.logger().error("auth-type is not set to 'online' in the Geyser config, this will break GeyserConnect. Disabling!");
             this.disable();
         }
     }

src/main/java/org/geysermc/extension/connect/PacketHandler.java

@@ -39,7 +39,6 @@
 import org.geysermc.extension.connect.utils.ServerManager;
 import org.geysermc.extension.connect.utils.Utils;
 import org.geysermc.geyser.entity.attribute.GeyserAttributeType;
-import org.geysermc.geyser.level.JavaDimension;
 import org.geysermc.geyser.network.UpstreamPacketHandler;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.util.DimensionUtils;
@@ -69,10 +68,10 @@ public PacketHandler(GeyserConnect geyserConnect, GeyserSession session, Bedrock
 
     @Override
     public void onDisconnect(String reason) {
-        if (session.getAuthData() != null) {
-            geyserConnect.logger().info(Utils.displayName(session) + " has disconnected (" + reason + ")");
-            ServerManager.unloadServers(session);
-        }
+        // The user has disconnected without having connected to an actual server. If they have connected to
+        // a server (transfer packet or geyser proxy), then the original packet handler has been restored.
+        ServerManager.unloadServers(session);
+        originalPacketHandler.onDisconnect(reason);
     }
 
     @Override
@@ -156,12 +155,12 @@ public PacketSignal handle(NetworkStackLatencyPacket packet) {
 
     @Override
     public PacketSignal handle(ResourcePackClientResponsePacket packet) {
-        return originalPacketHandler.handle(packet);
+        return originalPacketHandler.handle(packet); // relies on state in the original handler
     }
 
     @Override
     public PacketSignal handle(ResourcePackChunkRequestPacket packet) {
-        return originalPacketHandler.handle(packet);
+        return originalPacketHandler.handle(packet); // relies on state in the original handler
     }
 }
 

src/main/java/org/geysermc/extension/connect/storage/AbstractSQLStorageManager.java

@@ -30,7 +30,6 @@
 import org.geysermc.extension.connect.utils.Server;
 import org.geysermc.extension.connect.utils.ServerManager;
 import org.geysermc.extension.connect.utils.Utils;
-import org.geysermc.geyser.session.GeyserSession;
 
 import java.io.IOException;
 import java.sql.Connection;
@@ -69,23 +68,23 @@ public void closeStorage() {
     }
 
     @Override
-    public void saveServers(GeyserSession session) {
+    public void saveServers(org.geysermc.api.connection.Connection session) {
         // replace into works on MySQL and SQLite
         try (PreparedStatement updatePlayersServers = connection.prepareStatement("REPLACE INTO players(xuid, servers) VALUES(?, ?)")) {
-            updatePlayersServers.setString(1, session.getAuthData().xuid());
+            updatePlayersServers.setString(1, session.xuid());
             updatePlayersServers.setString(2, Utils.OBJECT_MAPPER.writeValueAsString(ServerManager.getServers(session)));
             updatePlayersServers.executeUpdate();
         } catch (IOException | SQLException exception) {
-            GeyserConnect.instance().logger().error("Couldn't save servers for " + session.getAuthData().name(), exception);
+            GeyserConnect.instance().logger().error("Couldn't save servers for " + session.bedrockUsername(), exception);
         }
     }
 
     @Override
-    public List<Server> loadServers(GeyserSession session) {
+    public List<Server> loadServers(org.geysermc.api.connection.Connection session) {
         List<Server> servers = new ArrayList<>();
 
         try (PreparedStatement getPlayersServers = connection.prepareStatement("SELECT servers FROM players WHERE xuid=?")) {
-            getPlayersServers.setString(1, session.getAuthData().xuid());
+            getPlayersServers.setString(1, session.xuid());
             ResultSet rs = getPlayersServers.executeQuery();
 
             while (rs.next()) {
@@ -96,7 +95,7 @@ public List<Server> loadServers(GeyserSession session) {
                 }
             }
         } catch (IOException | SQLException exception) {
-            GeyserConnect.instance().logger().error("Couldn't load servers for " + session.getAuthData().name(), exception);
+            GeyserConnect.instance().logger().error("Couldn't load servers for " + session.bedrockUsername(), exception);
         }
 
         return servers;

src/main/java/org/geysermc/extension/connect/storage/AbstractStorageManager.java

@@ -26,8 +26,8 @@
 package org.geysermc.extension.connect.storage;
 
 import com.fasterxml.jackson.annotation.JsonValue;
+import org.geysermc.api.connection.Connection;
 import org.geysermc.extension.connect.utils.Server;
-import org.geysermc.geyser.session.GeyserSession;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -40,10 +40,10 @@ public void setupStorage() {
     public void closeStorage() {
     }
 
-    public void saveServers(GeyserSession session) {
+    public void saveServers(Connection session) {
     }
 
-    public List<Server> loadServers(GeyserSession session) {
+    public List<Server> loadServers(Connection session) {
         return new ArrayList<>();
     }
 

src/main/java/org/geysermc/extension/connect/storage/DisabledStorageManager.java

@@ -25,8 +25,8 @@
 
 package org.geysermc.extension.connect.storage;
 
+import org.geysermc.api.connection.Connection;
 import org.geysermc.extension.connect.utils.Server;
-import org.geysermc.geyser.session.GeyserSession;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -38,12 +38,12 @@ public void setupStorage() {
     }
 
     @Override
-    public void saveServers(GeyserSession session) {
+    public void saveServers(Connection session) {
 
     }
 
     @Override
-    public List<Server> loadServers(GeyserSession session) {
+    public List<Server> loadServers(Connection session) {
         return new ArrayList<>();
     }
 }

src/main/java/org/geysermc/extension/connect/storage/JsonStorageManager.java

@@ -26,11 +26,11 @@
 package org.geysermc.extension.connect.storage;
 
 import com.fasterxml.jackson.core.type.TypeReference;
+import org.geysermc.api.connection.Connection;
 import org.geysermc.extension.connect.GeyserConnect;
 import org.geysermc.extension.connect.utils.Server;
 import org.geysermc.extension.connect.utils.ServerManager;
 import org.geysermc.extension.connect.utils.Utils;
-import org.geysermc.geyser.session.GeyserSession;
 
 import java.io.IOException;
 import java.nio.file.Path;
@@ -49,19 +49,19 @@ public void setupStorage() {
     }
 
     @Override
-    public void saveServers(GeyserSession session) {
+    public void saveServers(Connection session) {
         try {
-            Utils.OBJECT_MAPPER.writeValue(dataFolder.resolve(session.getAuthData().xuid() + ".json").toFile(), ServerManager.getServers(session));
+            Utils.OBJECT_MAPPER.writeValue(dataFolder.resolve(session.xuid() + ".json").toFile(), ServerManager.getServers(session));
         } catch (IOException ignored) {
         }
     }
 
     @Override
-    public List<Server> loadServers(GeyserSession session) {
+    public List<Server> loadServers(Connection session) {
         List<Server> servers = new ArrayList<>();
 
         try {
-            List<Server> loadedServers = Utils.OBJECT_MAPPER.readValue(dataFolder.resolve(session.getAuthData().xuid() + ".json").toFile(), new TypeReference<>() {
+            List<Server> loadedServers = Utils.OBJECT_MAPPER.readValue(dataFolder.resolve(session.xuid() + ".json").toFile(), new TypeReference<>() {
             });
             if (loadedServers != null) {
                 servers.addAll(loadedServers);

src/main/java/org/geysermc/extension/connect/utils/ServerManager.java

@@ -25,8 +25,8 @@
 
 package org.geysermc.extension.connect.utils;
 
+import org.geysermc.api.connection.Connection;
 import org.geysermc.extension.connect.GeyserConnect;
-import org.geysermc.geyser.session.GeyserSession;
 
 import java.util.HashMap;
 import java.util.List;
@@ -35,35 +35,35 @@
 public class ServerManager {
     private static final Map<String, List<Server>> servers = new HashMap<>();
 
-    public static void loadServers(GeyserSession session) {
+    public static void loadServers(Connection session) {
         GeyserConnect.instance().logger().debug("Loading servers for " + Utils.displayName(session));
         servers.put(session.xuid(), GeyserConnect.instance().storageManager().loadServers(session));
     }
 
-    public static void unloadServers(GeyserSession session) {
+    public static void unloadServers(Connection session) {
         if (getServers(session) == null) return;
         GeyserConnect.instance().logger().debug("Saving and unloading servers for " + Utils.displayName(session));
         GeyserConnect.instance().storageManager().saveServers(session);
         servers.remove(session.xuid());
     }
 
-    public static List<Server> getServers(GeyserSession session) {
+    public static List<Server> getServers(Connection session) {
         return servers.get(session.xuid());
     }
 
-    public static void addServer(GeyserSession session, Server server) {
+    public static void addServer(Connection session, Server server) {
         servers.get(session.xuid()).add(server);
     }
 
-    public static void removeServer(GeyserSession session, Server server) {
+    public static void removeServer(Connection session, Server server) {
         getServers(session).remove(server);
     }
 
-    public static int getServerIndex(GeyserSession session, Server server) {
+    public static int getServerIndex(Connection session, Server server) {
         return getServers(session).indexOf(server);
     }
 
-    public static void updateServer(GeyserSession session, int serverIndex, Server server) {
+    public static void updateServer(Connection session, int serverIndex, Server server) {
         getServers(session).set(serverIndex, server);
     }
 }

src/main/java/org/geysermc/extension/connect/utils/Utils.java

@@ -30,6 +30,7 @@
 import org.cloudburstmc.protocol.bedrock.packet.BedrockPacketHandler;
 import org.cloudburstmc.protocol.bedrock.packet.SetLocalPlayerAsInitializedPacket;
 import org.cloudburstmc.protocol.bedrock.packet.TransferPacket;
+import org.geysermc.api.connection.Connection;
 import org.geysermc.extension.connect.GeyserConnect;
 import org.geysermc.geyser.session.GeyserSession;
 
@@ -75,27 +76,28 @@ public static File fileOrCopiedFromResource(String fileName, String name) throws
         return file;
     }
 
-    public static String displayName(GeyserSession session) {
+    public static String displayName(Connection session) {
         return session.bedrockUsername() + " (" + session.xuid() + ")";
     }
 
     public static void sendToServer(GeyserSession session, BedrockPacketHandler originalPacketHandler, Server server) {
         GeyserConnect.instance().logger().info("Sending " + Utils.displayName(session) + " to " + server.title());
         GeyserConnect.instance().logger().debug(server.toString());
 
+        // Save the player's servers since we are changing packet handlers
+        // (and they are going to disconnect if it is a bedrock server)
+        ServerManager.unloadServers(session);
+
+        // Restore the original packet handler
+        session.getUpstream().getSession().setPacketHandler(originalPacketHandler);
+
         if (server.bedrock()) {
             // Send them to the bedrock server
             TransferPacket transferPacket = new TransferPacket();
             transferPacket.setAddress(server.address());
             transferPacket.setPort(server.port());
             session.sendUpstreamPacket(transferPacket);
         } else {
-            // Save the players servers since we are changing packet handlers
-            ServerManager.unloadServers(session);
-
-            // Restore the original packet handler
-            session.getUpstream().getSession().setPacketHandler(originalPacketHandler);
-
             // Set the remote server and un-initialize the session
             session.remoteServer(server);
             session.getUpstream().setInitialized(false);