Synchronize transition from PLAY to CONFIG state (#126)

Author: UserNugget

plugin/src/main/java/net/elytrium/limboapi/injection/login/LoginListener.java

@@ -76,6 +76,8 @@
 import net.elytrium.limboapi.injection.dummy.ClosedChannel;
 import net.elytrium.limboapi.injection.dummy.ClosedMinecraftConnection;
 import net.elytrium.limboapi.injection.dummy.DummyEventPool;
+import net.elytrium.limboapi.injection.login.confirmation.ConfirmHandler;
+import net.elytrium.limboapi.injection.login.confirmation.LoginConfirmHandler;
 import net.elytrium.limboapi.injection.packet.ServerLoginSuccessHook;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.format.NamedTextColor;
@@ -134,7 +136,7 @@ public void hookLoginSession(GameProfileRequestEvent event) throws Throwable {
       MC_CONNECTION_FIELD.set(handler, CLOSED_MINECRAFT_CONNECTION);
 
       if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) >= 0) {
-        connection.setActiveSessionHandler(StateRegistry.LOGIN, new LoginTrackHandler(connection));
+        connection.setActiveSessionHandler(StateRegistry.LOGIN, new LoginConfirmHandler(connection));
       }
 
       // From Velocity.
@@ -151,7 +153,7 @@ public void hookLoginSession(GameProfileRequestEvent event) throws Throwable {
                 inboundConnection.getIdentifiedKey()
             );
             if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) >= 0) {
-              ((LoginTrackHandler) connection.getActiveSessionHandler()).setPlayer(player);
+              ((ConfirmHandler) connection.getActiveSessionHandler()).setPlayer(player);
             }
             if (this.server.canRegisterConnection(player)) {
               if (!connection.isClosed()) {

plugin/src/main/java/net/elytrium/limboapi/injection/login/LoginTasksQueue.java

@@ -74,6 +74,8 @@
 import net.elytrium.limboapi.LimboAPI;
 import net.elytrium.limboapi.api.event.SafeGameProfileRequestEvent;
 import net.elytrium.limboapi.injection.event.EventManagerHook;
+import net.elytrium.limboapi.injection.login.confirmation.ConfirmHandler;
+import net.elytrium.limboapi.injection.login.confirmation.TransitionConfirmHandler;
 import net.kyori.adventure.text.Component;
 import org.slf4j.Logger;
 
@@ -241,21 +243,13 @@ private void initialize(MinecraftConnection connection) throws Throwable {
           this.player.disconnect0(reason.get(), true);
         } else {
           if (this.server.registerConnection(this.player)) {
-            if (connection.getActiveSessionHandler() instanceof LoginTrackHandler) {
-              LoginTrackHandler track = (LoginTrackHandler) connection.getActiveSessionHandler();
-              track.waitForConfirmation(() -> {
-                try {
-                  this.connectToServer(logger, connection);
-                } catch (Throwable e) {
-                  throw new ReflectionException(e);
-                }
+            if (connection.getActiveSessionHandler() instanceof ConfirmHandler) {
+              ConfirmHandler confirm = (ConfirmHandler) connection.getActiveSessionHandler();
+              confirm.waitForConfirmation(() -> {
+                this.connectToServer(logger, this.player, connection);
               });
             } else {
-              try {
-                this.connectToServer(logger, connection);
-              } catch (Throwable e) {
-                throw new ReflectionException(e);
-              }
+              this.connectToServer(logger, this.player, connection);
             }
           } else {
             this.player.disconnect0(Component.translatable("velocity.error.already-connected-proxy"), true);
@@ -268,13 +262,27 @@ private void initialize(MinecraftConnection connection) throws Throwable {
     });
   }
 
-  private void connectToServer(Logger logger, MinecraftConnection connection) throws Throwable {
+  private void connectToServer(Logger logger, ConnectedPlayer player, MinecraftConnection connection) {
     if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) < 0) {
-      connection.setActiveSessionHandler(connection.getState(),
-          (InitialConnectSessionHandler) INITIAL_CONNECT_SESSION_HANDLER_CONSTRUCTOR.invokeExact(this.player, this.server));
+      try {
+        connection.setActiveSessionHandler(connection.getState(),
+            (InitialConnectSessionHandler) INITIAL_CONNECT_SESSION_HANDLER_CONSTRUCTOR.invokeExact(this.player, this.server));
+      } catch (Throwable e) {
+        throw new ReflectionException(e);
+      }
     } else {
       if (connection.getState() == StateRegistry.PLAY) {
         connection.write(new StartUpdate());
+
+        TransitionConfirmHandler confirm = new TransitionConfirmHandler(connection);
+        confirm.setPlayer(player);
+
+        connection.setActiveSessionHandler(StateRegistry.PLAY, confirm);
+        confirm.waitForConfirmation(() -> {
+          this.connectToServer(logger, player, connection);
+        });
+
+        return;
       }
 
       connection.setActiveSessionHandler(StateRegistry.CONFIG,

plugin/src/main/java/net/elytrium/limboapi/injection/login/LoginTrackHandler.java renamed to plugin/src/main/java/net/elytrium/limboapi/injection/login/confirmation/ConfirmHandler.java

@@ -15,15 +15,13 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-package net.elytrium.limboapi.injection.login;
+package net.elytrium.limboapi.injection.login.confirmation;
 
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.StateRegistry;
-import com.velocitypowered.proxy.protocol.packet.LoginAcknowledged;
-import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.util.ReferenceCountUtil;
 import io.netty.util.internal.PlatformDependent;
@@ -35,16 +33,16 @@
 import net.elytrium.commons.utils.reflection.ReflectionException;
 import net.elytrium.limboapi.LimboAPI;
 
-public class LoginTrackHandler implements MinecraftSessionHandler {
+public abstract class ConfirmHandler implements MinecraftSessionHandler {
 
   private static final MethodHandle TEARDOWN_METHOD;
 
-  private final CompletableFuture<Object> confirmation = new CompletableFuture<>();
-  private final Queue<MinecraftPacket> queuedPackets = PlatformDependent.newMpscQueue();
-  private final MinecraftConnection connection;
-  private ConnectedPlayer player;
+  protected final CompletableFuture<Object> confirmation = new CompletableFuture<>();
+  protected final Queue<MinecraftPacket> queuedPackets = PlatformDependent.newMpscQueue();
+  protected final MinecraftConnection connection;
+  protected ConnectedPlayer player;
 
-  public LoginTrackHandler(MinecraftConnection connection) {
+  public ConfirmHandler(MinecraftConnection connection) {
     this.connection = connection;
   }
 
@@ -63,18 +61,6 @@ public void handleGeneric(MinecraftPacket packet) {
     }
   }
 
-  @Override
-  public boolean handle(LoginAcknowledged packet) {
-    this.connection.setState(StateRegistry.CONFIG);
-    this.confirmation.complete(this);
-    return true;
-  }
-
-  @Override
-  public void handleUnknown(ByteBuf buf) {
-    this.connection.close(true);
-  }
-
   @Override
   public void disconnected() {
     try {

plugin/src/main/java/net/elytrium/limboapi/injection/login/confirmation/LoginConfirmHandler.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2021 - 2023 Elytrium
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package net.elytrium.limboapi.injection.login.confirmation;
+
+import com.velocitypowered.proxy.connection.MinecraftConnection;
+import com.velocitypowered.proxy.protocol.StateRegistry;
+import com.velocitypowered.proxy.protocol.packet.LoginAcknowledged;
+import io.netty.buffer.ByteBuf;
+
+public class LoginConfirmHandler extends ConfirmHandler {
+
+  public LoginConfirmHandler(MinecraftConnection connection) {
+    super(connection);
+  }
+
+  @Override
+  public boolean handle(LoginAcknowledged packet) {
+    this.connection.setState(StateRegistry.CONFIG);
+    this.confirmation.complete(this);
+    return true;
+  }
+
+  @Override
+  public void handleUnknown(ByteBuf buf) {
+    this.connection.close(true);
+  }
+}

plugin/src/main/java/net/elytrium/limboapi/injection/login/confirmation/TransitionConfirmHandler.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2021 - 2023 Elytrium
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package net.elytrium.limboapi.injection.login.confirmation;
+
+import com.velocitypowered.proxy.connection.MinecraftConnection;
+import com.velocitypowered.proxy.protocol.StateRegistry;
+import com.velocitypowered.proxy.protocol.packet.config.FinishedUpdate;
+
+public class TransitionConfirmHandler extends ConfirmHandler {
+
+  public TransitionConfirmHandler(MinecraftConnection connection) {
+    super(connection);
+  }
+
+  @Override
+  public boolean handle(FinishedUpdate packet) {
+    if (this.connection.getState() == StateRegistry.PLAY) {
+      this.connection.setState(StateRegistry.CONFIG);
+      this.confirmation.complete(this);
+      return true;
+    }
+
+    return false;
+  }
+}

plugin/src/main/java/net/elytrium/limboapi/server/LimboImpl.java

@@ -90,7 +90,7 @@
 import net.elytrium.limboapi.api.protocol.PacketDirection;
 import net.elytrium.limboapi.api.protocol.PreparedPacket;
 import net.elytrium.limboapi.api.protocol.packets.PacketMapping;
-import net.elytrium.limboapi.injection.login.LoginTrackHandler;
+import net.elytrium.limboapi.injection.login.confirmation.ConfirmHandler;
 import net.elytrium.limboapi.injection.packet.MinecraftLimitedCompressDecoder;
 import net.elytrium.limboapi.material.Biome;
 import net.elytrium.limboapi.protocol.LimboProtocol;
@@ -414,9 +414,9 @@ private void spawnPlayerLocal(ConnectedPlayer player, LimboSessionHandler handle
           () -> this.limboName
       );
 
-      if (connection.getActiveSessionHandler() instanceof LoginTrackHandler) {
-        LoginTrackHandler track = (LoginTrackHandler) connection.getActiveSessionHandler();
-        track.waitForConfirmation(() -> {
+      if (connection.getActiveSessionHandler() instanceof ConfirmHandler) {
+        ConfirmHandler confirm = (ConfirmHandler) connection.getActiveSessionHandler();
+        confirm.waitForConfirmation(() -> {
           this.spawnPlayerLocal(handlerClass, sessionHandler, player, connection);
         });
       } else {