Forge --> vanilla switch fixed

src/main/java/org/adde0109/ambassador/forge/ForgeFMLConnectionType.java

@@ -30,7 +30,7 @@ public class ForgeFMLConnectionType implements ConnectionType {
 
   @Override
   public GameProfile addGameProfileTokensIfRequired(GameProfile original, PlayerInfoForwarding forwardingType) {
-    //This is meant for Arc light to parse
+    //This is meant for Arclight to parse
     if (forwardingType == PlayerInfoForwarding.LEGACY) {
       return original.addProperties(Collections.singleton(new GameProfile.Property("extraData", "\1FML" + netVersion + "\1", "")));
     } else {

src/main/java/org/adde0109/ambassador/velocity/VelocityEventHandler.java

@@ -7,7 +7,6 @@ import com.velocitypowered.api.event.connection.LoginEvent;
 import com.velocitypowered.api.event.connection.PostLoginEvent;
 import com.velocitypowered.api.event.player.*;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
-import com.velocitypowered.proxy.connection.backend.VelocityServerConnection;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
 import com.velocitypowered.proxy.protocol.StateRegistry;
 import org.adde0109.ambassador.Ambassador;
@@ -39,11 +38,6 @@ public class VelocityEventHandler {
     continuation.resume();
   }
 
-  @Subscribe(order = PostOrder.LAST)
-  public void onServerPreConnectEvent(ServerPreConnectEvent event, Continuation continuation) {
-    continuation.resume();
-  }
-
   @Subscribe(order = PostOrder.LAST)
   public void onPlayerChooseInitialServerEvent(PlayerChooseInitialServerEvent event, Continuation continuation) {
     ConnectedPlayer player = (ConnectedPlayer) event.getPlayer();

src/main/java/org/adde0109/ambassador/velocity/VelocityForgeClientConnectionPhase.java

@@ -1,7 +1,5 @@
 package org.adde0109.ambassador.velocity;
 
-import com.velocitypowered.api.event.player.KickedFromServerEvent;
-import com.velocitypowered.api.proxy.ServerConnection;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
@@ -10,7 +8,6 @@ import com.velocitypowered.proxy.connection.client.ClientConnectionPhase;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
 import com.velocitypowered.proxy.network.Connections;
 import com.velocitypowered.proxy.protocol.StateRegistry;
-import com.velocitypowered.proxy.protocol.packet.LoginPluginMessage;
 import com.velocitypowered.proxy.protocol.packet.LoginPluginResponse;
 import com.velocitypowered.proxy.protocol.packet.PluginMessage;
 import io.netty.buffer.Unpooled;
@@ -18,8 +15,10 @@ import net.kyori.adventure.text.Component;
 import org.adde0109.ambassador.Ambassador;
 import org.adde0109.ambassador.forge.FML2CRPMResetCompleteDecoder;
 import org.adde0109.ambassador.forge.ForgeConstants;
+import org.adde0109.ambassador.forge.ForgeFMLConnectionType;
 import org.adde0109.ambassador.forge.ForgeHandshakeUtils;
 import org.adde0109.ambassador.velocity.client.OutboundForgeHandshakeHolder;
+import org.adde0109.ambassador.velocity.client.OutboundSuccessHolder;
 
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.ScheduledFuture;
@@ -45,10 +44,14 @@ public enum VelocityForgeClientConnectionPhase implements ClientConnectionPhase
 
       //We unregister so no plugin sees this client while the client is being reset.
       ((VelocityServer) Ambassador.getInstance().server).unregisterConnection(player);
-      player.getConnectedServer().getConnection().getChannel().config().setAutoRead(false);
+      if (player.getConnectedServer() != null) {
+        player.getConnectedServer().disconnect();
+        player.setConnectedServer(null);
+      }
+      player.getConnectionInFlight().getConnection().getChannel().config().setAutoRead(false);
 
       connection.getChannel().pipeline().addBefore(Connections.MINECRAFT_DECODER, ForgeConstants.RESET_LISTENER,new FML2CRPMResetCompleteDecoder());
-      connection.getChannel().pipeline().addLast(ForgeConstants.FORGE_HANDSHAKE_HOLDER,new OutboundForgeHandshakeHolder());
+      connection.getChannel().pipeline().addAfter(Connections.MINECRAFT_ENCODER, ForgeConstants.FORGE_HANDSHAKE_HOLDER,new OutboundForgeHandshakeHolder());
 
       player.getConnection().setSessionHandler(new VelocityForgeHandshakeSessionHandler(player.getConnection().getSessionHandler(),player));
 
@@ -85,6 +88,14 @@ public enum VelocityForgeClientConnectionPhase implements ClientConnectionPhase
           player.setPhase(NOT_STARTED);
           //Send all held messages
           player.getConnection().getChannel().pipeline().remove(ForgeConstants.FORGE_HANDSHAKE_HOLDER);
+          player.getConnectionInFlight().getConnection().getChannel().config().setAutoRead(true);
+
+          if (!(server.getConnection().getType() instanceof ForgeFMLConnectionType)) {
+            MinecraftConnection connection = player.getConnection();
+            ((OutboundSuccessHolder) connection.getChannel().pipeline().get(ForgeConstants.SERVER_SUCCESS_LISTENER))
+                    .sendPacket();
+            connection.setState(StateRegistry.PLAY);
+          }
         }
         return true;
       } else {
@@ -96,12 +107,13 @@ public enum VelocityForgeClientConnectionPhase implements ClientConnectionPhase
 
 
 
-  public ServerConnection internalServerConnection;
+  public boolean vanillaMode = true;
 
   public boolean handle(ConnectedPlayer player, LoginPluginResponse response, VelocityServerConnection server) {
     player.setPhase(nextPhase());
 
     player.getConnectionInFlight().getConnection().write(response.retain());
+    vanillaMode = false;
     return true;
   }
 

src/main/java/org/adde0109/ambassador/velocity/backend/ForgeLoginSessionHandler.java

@@ -3,16 +3,16 @@ package org.adde0109.ambassador.velocity.backend;
 import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
-import com.velocitypowered.proxy.connection.backend.BackendConnectionPhases;
+import com.velocitypowered.proxy.connection.backend.*;
-import com.velocitypowered.proxy.connection.backend.LoginSessionHandler;
+import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
-import com.velocitypowered.proxy.connection.backend.TransitionSessionHandler;
-import com.velocitypowered.proxy.connection.backend.VelocityServerConnection;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
+import com.velocitypowered.proxy.protocol.StateRegistry;
 import com.velocitypowered.proxy.protocol.packet.Disconnect;
 import com.velocitypowered.proxy.protocol.packet.LoginPluginMessage;
 import com.velocitypowered.proxy.protocol.packet.ServerLoginSuccess;
 import net.kyori.adventure.text.Component;
 import org.adde0109.ambassador.forge.ForgeConstants;
+import org.adde0109.ambassador.forge.ForgeFMLConnectionType;
 import org.adde0109.ambassador.velocity.client.OutboundSuccessHolder;
 
 public class ForgeLoginSessionHandler implements MinecraftSessionHandler {
@@ -45,7 +45,19 @@ public class ForgeLoginSessionHandler implements MinecraftSessionHandler {
     if ((serverConnection.getPhase() instanceof VelocityForgeBackendConnectionPhase phase)) {
       phase.onLoginSuccess(serverConnection,serverConnection.getPlayer());
     }
+
+    ConnectedPlayer player = serverConnection.getPlayer();
+    if (!(serverConnection.getConnection().getType() instanceof ForgeFMLConnectionType)) {
+      player.getPhase().resetConnectionPhase(player);
+    } else {
+      MinecraftConnection connection = player.getConnection();
+      ((OutboundSuccessHolder) connection.getChannel().pipeline().get(ForgeConstants.SERVER_SUCCESS_LISTENER))
+              .sendPacket();
+      connection.setState(StateRegistry.PLAY);
+    }
+
     original.handle(packet);
+
     serverConnection.getConnection().setSessionHandler(
             new ForgePlaySessionHandler((TransitionSessionHandler) serverConnection
                     .getConnection().getSessionHandler(),serverConnection));

src/main/java/org/adde0109/ambassador/velocity/backend/VelocityForgeBackendConnectionPhase.java

@@ -3,14 +3,9 @@ package org.adde0109.ambassador.velocity.backend;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.backend.BackendConnectionPhase;
 import com.velocitypowered.proxy.connection.backend.VelocityServerConnection;
-import com.velocitypowered.proxy.connection.client.ClientConnectionPhase;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
-import com.velocitypowered.proxy.protocol.StateRegistry;
 import com.velocitypowered.proxy.protocol.packet.LoginPluginMessage;
-import io.netty.channel.PendingWriteQueue;
-import org.adde0109.ambassador.forge.ForgeConstants;
 import org.adde0109.ambassador.velocity.VelocityForgeClientConnectionPhase;
-import org.adde0109.ambassador.velocity.client.OutboundSuccessHolder;
 
 public enum VelocityForgeBackendConnectionPhase implements BackendConnectionPhase {
   NOT_STARTED() {
@@ -23,11 +18,15 @@ public enum VelocityForgeBackendConnectionPhase implements BackendConnectionPhas
     @Override
     public void onLoginSuccess(VelocityServerConnection serverCon, ConnectedPlayer player) {
       serverCon.setConnectionPhase(VelocityForgeBackendConnectionPhase.COMPLETE);
+    }
 
-      MinecraftConnection connection = player.getConnection();
+    @Override
-      ((OutboundSuccessHolder) connection.getChannel().pipeline().get(ForgeConstants.SERVER_SUCCESS_LISTENER))
+    void onTransitionToNewPhase(VelocityServerConnection connection) {
-              .sendPacket();
+      MinecraftConnection mc = connection.getConnection();
-      connection.setState(StateRegistry.PLAY);
+      if (mc != null) {
+        //This looks ugly. But unless the player didn't have a FML marker, we're fine.
+        mc.setType(connection.getPlayer().getConnection().getType());
+      }
     }
   },
 
@@ -45,7 +44,7 @@ public enum VelocityForgeBackendConnectionPhase implements BackendConnectionPhas
 
   public void handle(VelocityServerConnection server, ConnectedPlayer player, LoginPluginMessage message) {
 
-    VelocityForgeBackendConnectionPhase newPhase = nextPhase();
+    VelocityForgeBackendConnectionPhase newPhase = getNewPhase(server,message);
 
     server.setConnectionPhase(newPhase);
 
@@ -56,11 +55,24 @@ public enum VelocityForgeBackendConnectionPhase implements BackendConnectionPhas
   }
 
   public void onLoginSuccess(VelocityServerConnection serverCon, ConnectedPlayer player) {
-
   }
+
+  void onTransitionToNewPhase(VelocityServerConnection connection) {
+  }
+
   VelocityForgeBackendConnectionPhase nextPhase() {
     return this;
   }
+
+  private VelocityForgeBackendConnectionPhase getNewPhase(VelocityServerConnection serverConnection,
+                                                       LoginPluginMessage packet) {
+    VelocityForgeBackendConnectionPhase phaseToTransitionTo = nextPhase();
+    if (phaseToTransitionTo != this) {
+      phaseToTransitionTo.onTransitionToNewPhase(serverConnection);
+    }
+    return phaseToTransitionTo;
+  }
+
   public boolean consideredComplete() {
     return false;
   }