Seperated catcher and reset listener

build.gradle

@@ -4,7 +4,7 @@ plugins {
 }
 
 group 'org.adde0109'
-version '1.0.9-alpha'
+version '1.0.11-alpha'
 
 repositories {
     maven {

src/main/java/org/adde0109/ambassador/Ambassador.java

@@ -22,7 +22,7 @@ import org.slf4j.Logger;
 
 import java.nio.file.Path;
 
-@Plugin(id = "ambassador", name = "Ambassador", version = "1.0.9-alpha", authors = {"adde0109"})
+@Plugin(id = "ambassador", name = "Ambassador", version = "1.0.11-alpha", authors = {"adde0109"})
 public class Ambassador {
 
   public ProxyServer server;

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

@@ -18,12 +18,14 @@ import org.adde0109.ambassador.velocity.VelocityForgeClientConnectionPhase;
 import org.adde0109.ambassador.velocity.VelocityForgeHandshakeSessionHandler;
 import org.adde0109.ambassador.velocity.VelocityLoginPayloadManager;
 
-import java.util.NoSuchElementException;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
 
 public class FML2CRPMClientConnectionPhase extends VelocityForgeClientConnectionPhase {
   private static String OUTBOUND_CATCHER_NAME = "ambassador-catcher";
+  private static String RESET_LISTENER = "ambassador-reset-listener";
 
   //TODO: Use modData inside ConnectedPlayer instead
   public byte[] modListData;
@@ -47,22 +49,21 @@ public class FML2CRPMClientConnectionPhase extends VelocityForgeClientConnection
     connection.setSessionHandler(new VelocityForgeHandshakeSessionHandler(connection.getSessionHandler(),player));
 
 
-    if (connection.getState() == StateRegistry.LOGIN) {
-      getPayloadManager().sendPayload("fml:loginwrapper", Unpooled.wrappedBuffer(ForgeHandshakeUtils.generateResetPacket()));
-    } else {
-      connection.write(new PluginMessage("fml:handshake",Unpooled.wrappedBuffer(ForgeHandshakeUtils.generatePluginResetPacket())));
-      connection.setState(StateRegistry.LOGIN);
-    }
-    getPayloadManager().listenFor(98).thenAccept((response) -> {
-      this.clientPhase = ClientPhase.HANDSHAKE;
-      future.complete(true);
-    });
 
-    this.clientPhase = null;
+    ScheduledFuture<?> scheduledFuture = connection.eventLoop().schedule(()-> {
-    connection.getChannel().pipeline().addBefore(Connections.HANDLER,OUTBOUND_CATCHER_NAME,new FML2CRPMConnectionHandler(() -> {
+      connection.getChannel().pipeline().remove(OUTBOUND_CATCHER_NAME);
-      player.getConnection().setState(StateRegistry.PLAY);
       future.complete(false);
+    },5, TimeUnit.SECONDS);
+    connection.getChannel().pipeline().addBefore(Connections.HANDLER,RESET_LISTENER,new FML2CRPMResetCompleteListener(() -> {
+      if (scheduledFuture.cancel(false)) {
+        connection.setState(StateRegistry.LOGIN);
+        this.clientPhase = ClientPhase.HANDSHAKE;
+        future.complete(true);
+      }
     }));
+    connection.write(new PluginMessage("fml:handshake",Unpooled.wrappedBuffer(ForgeHandshakeUtils.generatePluginResetPacket())));
+    this.clientPhase = null;
+    connection.getChannel().pipeline().addBefore(Connections.HANDLER,OUTBOUND_CATCHER_NAME,new FML2CRPMOutboundCatcher());
     return future;
   }
   public void complete(VelocityServer server, ConnectedPlayer player, MinecraftConnection connection) {
@@ -80,10 +81,8 @@ public class FML2CRPMClientConnectionPhase extends VelocityForgeClientConnection
 
     connection.setState(StateRegistry.PLAY);
     connection.setSessionHandler(((VelocityForgeHandshakeSessionHandler) connection.getSessionHandler()).getOriginal());
-    try {
+
-      connection.getChannel().pipeline().remove(OUTBOUND_CATCHER_NAME);
+    connection.getChannel().pipeline().remove(OUTBOUND_CATCHER_NAME);
-    } catch (NoSuchElementException ignored) {
-    }
   }
 
   public void handleKick(KickedFromServerEvent event) {

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

@@ -1,24 +1,17 @@
 package org.adde0109.ambassador.forge;
 
 import com.velocitypowered.proxy.protocol.packet.LoginPluginMessage;
-import com.velocitypowered.proxy.protocol.packet.LoginPluginResponse;
 import com.velocitypowered.proxy.protocol.packet.ServerLoginSuccess;
-import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPromise;
 import io.netty.util.ReferenceCountUtil;
-import org.jetbrains.annotations.NotNull;
 
 import java.util.*;
 
-public class FML2CRPMConnectionHandler extends ChannelDuplexHandler {
+public class FML2CRPMOutboundCatcher extends ChannelOutboundHandlerAdapter {
 
   private final Map<ChannelPromise, Object> catchedPackets = Collections.synchronizedMap(new LinkedHashMap<>());
-  private final Runnable abort;
-
-  public FML2CRPMConnectionHandler(Runnable abort) {
-    this.abort = abort;
-  }
 
   @Override
   public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
@@ -48,15 +41,4 @@ public class FML2CRPMConnectionHandler extends ChannelDuplexHandler {
       catchedPackets.put(promise,msg);
     }
   }
-
-  @Override
-  public void channelRead(@NotNull ChannelHandlerContext ctx, @NotNull Object msg) throws Exception {
-    if (!(msg instanceof LoginPluginResponse)) {
-      abort.run();
-      ctx.pipeline().remove(this);
-      ctx.pipeline().fireChannelRead(msg);
-    } else {
-      ctx.fireChannelRead(msg);
-    }
-  }
 }

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

@@ -0,0 +1,29 @@
+package org.adde0109.ambassador.forge;
+
+import com.velocitypowered.proxy.protocol.ProtocolUtils;
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+
+public class FML2CRPMResetCompleteListener extends ChannelInboundHandlerAdapter {
+
+  final Runnable whenComplete;
+
+  public FML2CRPMResetCompleteListener(Runnable whenComplete) {
+    this.whenComplete = whenComplete;
+  }
+
+  @Override
+  public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+    if (msg instanceof ByteBuf buf) {
+      int originalReaderIndex = buf.readerIndex();
+      int packetId = ProtocolUtils.readVarInt(buf);
+      if (packetId == 0x02 && buf.readableBytes() > 1) {
+        ctx.pipeline().remove(this);
+        whenComplete.run();
+        return;
+      }
+    }
+    ctx.fireChannelRead(msg);
+  }
+}