Tweaked logic and added disconnect handling

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

@@ -45,15 +45,23 @@ 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);
 
       connection.getChannel().pipeline().addBefore(Connections.MINECRAFT_DECODER, ForgeConstants.RESET_LISTENER,new FML2CRPMResetCompleteDecoder());
       connection.getChannel().pipeline().addLast(ForgeConstants.FORGE_HANDSHAKE_HOLDER,new OutboundForgeHandshakeHolder());
 
+      player.getConnection().setSessionHandler(new VelocityForgeHandshakeSessionHandler(player.getConnection().getSessionHandler(),player));
+
       connection.write(new PluginMessage("fml:handshake", Unpooled.wrappedBuffer(ForgeHandshakeUtils.generatePluginResetPacket())));
 
       player.setPhase(WAITING_RESET);
       WAITING_RESET.onTransitionToNewPhase(player);
     }
+
+    @Override
+    public boolean consideredComplete() {
+      return true;
+    }
   },
   WAITING_RESET {
 
@@ -77,8 +85,6 @@ public enum VelocityForgeClientConnectionPhase implements ClientConnectionPhase
           player.setPhase(NOT_STARTED);
           //Send all held messages
           player.getConnection().getChannel().pipeline().remove(ForgeConstants.FORGE_HANDSHAKE_HOLDER);
-
-          player.getConnection().setSessionHandler(new VelocityForgeHandshakeSessionHandler(player.getConnection().getSessionHandler(),player));
         }
         return true;
       } else {
@@ -109,6 +115,13 @@ public enum VelocityForgeClientConnectionPhase implements ClientConnectionPhase
     return this;
   }
 
+  @Override
+  public boolean consideredComplete() {
+    return false;
+  }
+
+
+
 /*
   public void handleKick(KickedFromServerEvent event) {
     //If kicked before the client has entered PLAY and has been reset.

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

@@ -1,14 +1,19 @@
 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.LoginSessionHandler;
 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.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.velocity.client.OutboundSuccessHolder;
 
 public class ForgeLoginSessionHandler implements MinecraftSessionHandler {
 
@@ -47,9 +52,23 @@ public class ForgeLoginSessionHandler implements MinecraftSessionHandler {
     return true;
   }
 
+  @Override
+  public boolean handle(Disconnect packet) {
+    if (!serverConnection.getPlayer().getPhase().consideredComplete()) {
+      serverConnection.getPlayer().handleConnectionException(serverConnection.getServer(), packet,false);
+      return true;
+    }
+    return original.handle(packet);
+  }
 
   @Override
   public void disconnected() {
+    if (!serverConnection.getPlayer().getPhase().consideredComplete()) {
+      serverConnection.getPlayer().handleConnectionException(serverConnection.getServer(),
+              Disconnect.create(Component.text("Probably mismatched mods"),
+                      serverConnection.getPlayer().getProtocolVersion()),false);
+      return;
+    }
     original.disconnected();
   }
 

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

@@ -10,6 +10,7 @@ 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() {
@@ -24,7 +25,8 @@ public enum VelocityForgeBackendConnectionPhase implements BackendConnectionPhas
       serverCon.setConnectionPhase(VelocityForgeBackendConnectionPhase.COMPLETE);
 
       MinecraftConnection connection = player.getConnection();
-      connection.getChannel().pipeline().remove(ForgeConstants.SERVER_SUCCESS_LISTENER);
+      ((OutboundSuccessHolder) connection.getChannel().pipeline().get(ForgeConstants.SERVER_SUCCESS_LISTENER))
+              .sendPacket();
       connection.setState(StateRegistry.PLAY);
     }
   },

src/main/java/org/adde0109/ambassador/velocity/client/OutboundSuccessHolder.java

@@ -8,6 +8,12 @@ import io.netty.channel.ChannelPromise;
 public class OutboundSuccessHolder extends ChannelOutboundHandlerAdapter {
 
   private ServerLoginSuccess packet;
+  private ChannelHandlerContext ctx;
+
+  @Override
+  public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
+    this.ctx = ctx;
+  }
 
   @Override
   public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
@@ -18,11 +24,8 @@ public class OutboundSuccessHolder extends ChannelOutboundHandlerAdapter {
     }
   }
 
-  @Override
-  public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
-    if (ctx.channel().isActive()) {
-      ctx.write(packet, ctx.voidPromise());
-      ctx.flush();
-    }
+  public void sendPacket() {
+    ctx.write(packet, ctx.voidPromise());
+    ctx.flush();
   }
 }