State switching rework with other bugfixes

gradle.properties

@@ -1,5 +1,5 @@
 org.gradle.jvmargs=-Xmx4096m
-fastPrepareVersion=1.0.8
+fastPrepareVersion=1.0.9
 velocityVersion=3.3.0-SNAPSHOT
 nettyVersion=4.1.86.Final
 fastutilVersion=8.5.11

plugin/src/main/java/net/elytrium/limboapi/LimboAPI.java

@@ -32,6 +32,7 @@
 import com.velocitypowered.natives.util.Natives;
 import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
+import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
 import com.velocitypowered.proxy.event.VelocityEventManager;
 import com.velocitypowered.proxy.network.Connections;
@@ -61,6 +62,7 @@
 import net.elytrium.commons.utils.reflection.ReflectionException;
 import net.elytrium.commons.utils.updates.UpdatesChecker;
 import net.elytrium.fastprepare.PreparedPacketFactory;
+import net.elytrium.fastprepare.handler.PreparedPacketEncoder;
 import net.elytrium.limboapi.api.Limbo;
 import net.elytrium.limboapi.api.LimboFactory;
 import net.elytrium.limboapi.api.chunk.BuiltInBiome;
@@ -414,7 +416,40 @@ public ByteBuf encodeSingleLoginUncompressed(MinecraftPacket packet, ProtocolVer
   }
 
   public void inject3rdParty(Player player, MinecraftConnection connection, ChannelPipeline pipeline) {
-    this.preparedPacketFactory.inject(player, connection, pipeline);
+    StateRegistry state = connection.getState();
+    if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) < 0
+        || (state != StateRegistry.CONFIG && state != StateRegistry.LOGIN)) {
+      this.preparedPacketFactory.inject(player, connection, pipeline);
+    } else {
+      this.configPreparedPacketFactory.inject(player, connection, pipeline);
+    }
+  }
+
+  public void setState(MinecraftConnection connection, StateRegistry stateRegistry) {
+    connection.setState(stateRegistry);
+    this.setEncoderState(connection, stateRegistry);
+  }
+
+  public void setActiveSessionHandler(MinecraftConnection connection, StateRegistry stateRegistry,
+                                      MinecraftSessionHandler sessionHandler) {
+    connection.setActiveSessionHandler(stateRegistry, sessionHandler);
+    this.setEncoderState(connection, stateRegistry);
+  }
+
+  public void setEncoderState(MinecraftConnection connection, StateRegistry state) {
+    // As CONFIG state was added in 1.20.2, no need to track it for lower versions
+    if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) < 0) {
+      return;
+    }
+
+    PreparedPacketEncoder encoder = connection.getChannel().pipeline().get(PreparedPacketEncoder.class);
+    if (encoder != null) {
+      if (state != StateRegistry.CONFIG && state != StateRegistry.LOGIN) {
+        encoder.setFactory(this.preparedPacketFactory);
+      } else {
+        encoder.setFactory(this.configPreparedPacketFactory);
+      }
+    }
   }
 
   public void deject3rdParty(ChannelPipeline pipeline) {

plugin/src/main/java/net/elytrium/limboapi/Settings.java

@@ -115,6 +115,7 @@ public static class MESSAGES {
 
       public String TOO_BIG_PACKET = "{PRFX}{NL}{NL}&cYour client sent too big packet!";
       public String INVALID_PING = "{PRFX}{NL}{NL}&cYour client sent invalid ping packet!";
+      public String INVALID_SWITCH = "{PRFX}{NL}{NL}&cYour client sent an unexpected state switching packet!";
       public String TIME_OUT = "{PRFX}{NL}{NL}Timed out.";
     }
   }

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

@@ -76,7 +76,6 @@
 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;
@@ -136,7 +135,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 LoginConfirmHandler(connection));
+        connection.setActiveSessionHandler(StateRegistry.LOGIN, new LoginConfirmHandler(this.plugin, connection));
       }
 
       // From Velocity.
@@ -153,7 +152,7 @@ public void hookLoginSession(GameProfileRequestEvent event) throws Throwable {
                 inboundConnection.getIdentifiedKey()
             );
             if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) >= 0) {
-              ((ConfirmHandler) connection.getActiveSessionHandler()).setPlayer(player);
+              ((LoginConfirmHandler) connection.getActiveSessionHandler()).setPlayer(player);
             }
             if (this.server.canRegisterConnection(player)) {
               if (!connection.isClosed()) {
@@ -204,19 +203,13 @@ public void hookLoginSession(GameProfileRequestEvent event) throws Throwable {
 
                 this.plugin.setInitialID(player, playerUniqueID);
 
-                if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) < 0) {
+                if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) >= 0) {
+                  ((LoginConfirmHandler) connection.getActiveSessionHandler())
+                      .thenRun(() -> this.fireRegisterEvent(player, connection, inbound, handler));
+                } else {
                   connection.setState(StateRegistry.PLAY);
+                  this.fireRegisterEvent(player, connection, inbound, handler);
                 }
-
-                this.server.getEventManager().fire(new LoginLimboRegisterEvent(player)).thenAcceptAsync(limboRegisterEvent -> {
-                  LoginTasksQueue queue = new LoginTasksQueue(this.plugin, handler, this.server, player, inbound, limboRegisterEvent.getOnJoinCallbacks());
-                  this.plugin.addLoginQueue(player, queue);
-                  this.plugin.setKickCallback(player, limboRegisterEvent.getOnKickCallback());
-                  queue.next();
-                }, connection.eventLoop()).exceptionally(t -> {
-                  LimboAPI.getLogger().error("Exception while registering LimboAPI login handlers for {}.", player, t);
-                  return null;
-                });
               }
             } else {
               player.disconnect0(Component.translatable("velocity.error.already-connected-proxy", NamedTextColor.RED), true);
@@ -229,11 +222,29 @@ public void hookLoginSession(GameProfileRequestEvent event) throws Throwable {
     }
   }
 
+  private void fireRegisterEvent(ConnectedPlayer player, MinecraftConnection connection,
+      InitialInboundConnection inbound, Object handler) {
+    this.server.getEventManager().fire(new LoginLimboRegisterEvent(player)).thenAcceptAsync(limboRegisterEvent -> {
+      LoginTasksQueue queue = new LoginTasksQueue(this.plugin, handler, this.server, player, inbound, limboRegisterEvent.getOnJoinCallbacks());
+      this.plugin.addLoginQueue(player, queue);
+      this.plugin.setKickCallback(player, limboRegisterEvent.getOnKickCallback());
+      queue.next();
+    }, connection.eventLoop()).exceptionally(t -> {
+      LimboAPI.getLogger().error("Exception while registering LimboAPI login handlers for {}.", player, t);
+      return null;
+    });
+  }
+
   @Subscribe
   public void hookPlaySession(ServerConnectedEvent event) {
     ConnectedPlayer player = (ConnectedPlayer) event.getPlayer();
     MinecraftConnection connection = player.getConnection();
 
+    // 1.20.2+ can ignore this, as it should be despawned by default
+    if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) >= 0) {
+      return;
+    }
+
     connection.eventLoop().execute(() -> {
       if (!(connection.getActiveSessionHandler() instanceof ClientPlaySessionHandler)) {
         try {

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

@@ -57,7 +57,6 @@
 import com.velocitypowered.proxy.protocol.packet.LegacyPlayerListItem;
 import com.velocitypowered.proxy.protocol.packet.UpsertPlayerInfo;
 import com.velocitypowered.proxy.protocol.packet.chat.ComponentHolder;
-import com.velocitypowered.proxy.protocol.packet.config.StartUpdate;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import io.netty.channel.ChannelPipeline;
 import io.netty.channel.EventLoop;
@@ -74,8 +73,8 @@
 import net.elytrium.commons.utils.reflection.ReflectionException;
 import net.elytrium.limboapi.LimboAPI;
 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.elytrium.limboapi.injection.login.confirmation.LoginConfirmHandler;
+import net.elytrium.limboapi.server.LimboSessionHandlerImpl;
 import net.kyori.adventure.text.Component;
 import org.slf4j.Logger;
 
@@ -146,7 +145,7 @@ private void finish() {
                           .setProperties(gameProfile.getGameProfile().getProperties())
                   )
               ));
-            } else if (connection.getState() == StateRegistry.PLAY) {
+            } else if (connection.getState() != StateRegistry.CONFIG) {
               UpsertPlayerInfo.Entry playerInfoEntry = new UpsertPlayerInfo.Entry(this.player.getUniqueId());
               playerInfoEntry.setDisplayName(new ComponentHolder(this.player.getProtocolVersion(), Component.text(gameProfile.getUsername())));
               playerInfoEntry.setProfile(gameProfile.getGameProfile());
@@ -200,8 +199,8 @@ private void finish() {
   private void initialize(MinecraftConnection connection) throws Throwable {
     connection.setAssociation(this.player);
     if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) < 0
-        || (connection.getState() != StateRegistry.LOGIN && connection.getState() != StateRegistry.CONFIG)) {
-      connection.setState(StateRegistry.PLAY);
+        || connection.getState() != StateRegistry.CONFIG) {
+      this.plugin.setState(connection, StateRegistry.PLAY);
     }
 
     ChannelPipeline pipeline = connection.getChannel().pipeline();
@@ -235,16 +234,16 @@ private void initialize(MinecraftConnection connection) throws Throwable {
       } else {
         Optional<Component> reason = event.getResult().getReasonComponent();
         if (reason.isPresent()) {
-          this.player.disconnect0(reason.get(), true);
+          this.player.disconnect0(reason.get(), false);
         } else {
           if (this.server.registerConnection(this.player)) {
-            if (connection.getActiveSessionHandler() instanceof ConfirmHandler confirm) {
+            if (connection.getActiveSessionHandler() instanceof LoginConfirmHandler confirm) {
               confirm.waitForConfirmation(() -> this.connectToServer(logger, this.player, connection));
             } else {
               this.connectToServer(logger, this.player, connection);
             }
           } else {
-            this.player.disconnect0(Component.translatable("velocity.error.already-connected-proxy"), true);
+            this.player.disconnect0(Component.translatable("velocity.error.already-connected-proxy"), false);
           }
         }
       }
@@ -254,6 +253,7 @@ private void initialize(MinecraftConnection connection) throws Throwable {
     });
   }
 
+  @SuppressFBWarnings("NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE")
   private void connectToServer(Logger logger, ConnectedPlayer player, MinecraftConnection connection) {
     if (connection.getProtocolVersion().compareTo(ProtocolVersion.MINECRAFT_1_20_2) < 0) {
       try {
@@ -262,20 +262,14 @@ private void connectToServer(Logger logger, ConnectedPlayer player, MinecraftCon
       } 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));
+    } else if (connection.getState() == StateRegistry.PLAY) {
+      // Synchronize with the client to ensure that it will not corrupt CONFIG state with PLAY packets
+      ((LimboSessionHandlerImpl) connection.getActiveSessionHandler())
+          .disconnectToConfig(() -> this.connectToServer(logger, player, connection));
 
-        return;
-      }
-
-      connection.setActiveSessionHandler(StateRegistry.CONFIG,
+      return; // Re-running this method due to synchronization with the client
+    } else {
+      this.plugin.setActiveSessionHandler(connection, StateRegistry.CONFIG,
           new ClientConfigSessionHandler(this.server, this.player));
     }