From fe69711818c2d803c3bc231a7098510e1bb85eb7 Mon Sep 17 00:00:00 2001
From: Adrian Bergqvist <adde0109@gmail.com>
Date: Thu, 15 Sep 2022 08:23:24 +0200
Subject: [PATCH] WIP: listen to response

---
 .../forge/ForgeFML2ClientConnectionPhase.java    | 16 ++++++++++++----
 .../ambassador/forge/ForgeHandshakeHandler.java  |  2 +-
 .../VelocityForgeClientConnectionPhase.java      |  5 +++--
 .../VelocityForgeHandshakeSessionHandler.java    | 13 +++++++++----
 4 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/main/java/org/adde0109/ambassador/forge/ForgeFML2ClientConnectionPhase.java b/src/main/java/org/adde0109/ambassador/forge/ForgeFML2ClientConnectionPhase.java
index 538e272..5769e43 100644
--- a/src/main/java/org/adde0109/ambassador/forge/ForgeFML2ClientConnectionPhase.java
+++ b/src/main/java/org/adde0109/ambassador/forge/ForgeFML2ClientConnectionPhase.java
@@ -1,5 +1,6 @@
 package org.adde0109.ambassador.forge;
 
+import com.velocitypowered.api.event.Continuation;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.client.ClientConnectionPhase;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
@@ -17,30 +18,37 @@ public class ForgeFML2ClientConnectionPhase implements VelocityForgeClientConnec
   private final MinecraftConnection connection;
   private boolean isResettable;
   private ByteBuf modListData;
+  private Continuation whenComplete;
   ForgeFML2ClientConnectionPhase(MinecraftConnection connection) {
     this.connection = connection;
   }
   @Override
-  public void handleLogin(@Nullable ForgeHandshakeUtils.CachedServerHandshake handshake) {
+  public void handleLogin(ForgeHandshakeUtils.CachedServerHandshake handshake, Continuation continuation) {
+    VelocityForgeHandshakeSessionHandler sessionHandler = new VelocityForgeHandshakeSessionHandler(this);
     if(handshake == null) {
       this.connection.write(new LoginPluginMessage(98,"fml:loginwrapper", Unpooled.wrappedBuffer(ForgeHandshakeUtils.generateResetPacket())));
+      sessionHandler.listen(98);
     } else {
       this.connection.delayedWrite(new LoginPluginMessage(1,"fml:loginwrapper", Unpooled.wrappedBuffer(handshake.modListPacket)));
+      sessionHandler.listen(1);
       for (int i = 0;i<handshake.otherPackets.size();i++) {
         this.connection.delayedWrite(new LoginPluginMessage(i+2,"fml:loginwrapper", Unpooled.wrappedBuffer(handshake.otherPackets.get(i))));
+        sessionHandler.listen(i+2);
       }
-      //TODO:Register sent packets
-      this.connection.setSessionHandler(new VelocityForgeHandshakeSessionHandler(this));
+      this.connection.setSessionHandler(sessionHandler);
       this.connection.flush();
     }
   }
 
   @Override
-  public void handle(LoginPluginResponse packet) {
+  public void handle(LoginPluginResponse packet, boolean lastMessage) {
     if (packet.getId() == 98) {
       isResettable = packet.isSuccess();
     } else {
       modListData = packet.content().retain();
     }
+    if (lastMessage) {
+      whenComplete.resume();
+    }
   }
 }
diff --git a/src/main/java/org/adde0109/ambassador/forge/ForgeHandshakeHandler.java b/src/main/java/org/adde0109/ambassador/forge/ForgeHandshakeHandler.java
index 4a46d21..8da119e 100644
--- a/src/main/java/org/adde0109/ambassador/forge/ForgeHandshakeHandler.java
+++ b/src/main/java/org/adde0109/ambassador/forge/ForgeHandshakeHandler.java
@@ -42,7 +42,7 @@ public class ForgeHandshakeHandler {
       if (ex != null) {
         ambassador.logger.warn("Forge player, " + player.getUsername() + ", is entering vanilla-mode because of: " + ex.getMessage());
       }
-      phase.handleLogin(msg);
+      phase.handleLogin(msg,continuation);
     }, player.getConnection().eventLoop());
   }
 
diff --git a/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeClientConnectionPhase.java b/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeClientConnectionPhase.java
index 9d592ab..9918e0a 100644
--- a/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeClientConnectionPhase.java
+++ b/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeClientConnectionPhase.java
@@ -1,5 +1,6 @@
 package org.adde0109.ambassador.velocity;
 
+import com.velocitypowered.api.event.Continuation;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.client.ClientConnectionPhase;
 import com.velocitypowered.proxy.protocol.packet.LoginPluginResponse;
@@ -9,8 +10,8 @@ import javax.annotation.Nullable;
 
 public interface VelocityForgeClientConnectionPhase extends ClientConnectionPhase {
 
-  default void handleLogin(@Nullable ForgeHandshakeUtils.CachedServerHandshake initialHandshake) {
+  default void handleLogin(@Nullable ForgeHandshakeUtils.CachedServerHandshake initialHandshake, Continuation continuation) {
   }
-  default void handle(LoginPluginResponse packet) {
+  default void handle(LoginPluginResponse packet, boolean lastMessage) {
   }
 }
diff --git a/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeHandshakeSessionHandler.java b/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeHandshakeSessionHandler.java
index 48653cf..6b9051d 100644
--- a/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeHandshakeSessionHandler.java
+++ b/src/main/java/org/adde0109/ambassador/velocity/VelocityForgeHandshakeSessionHandler.java
@@ -1,11 +1,14 @@
 package org.adde0109.ambassador.velocity;
 
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
-import com.velocitypowered.proxy.connection.client.ClientConnectionPhase;
 import com.velocitypowered.proxy.protocol.packet.LoginPluginResponse;
 
+import java.util.ArrayList;
+import java.util.List;
+
 public class VelocityForgeHandshakeSessionHandler implements MinecraftSessionHandler {
 
+  private final ArrayList<Integer> listenerList = new ArrayList();
   private final VelocityForgeClientConnectionPhase phase;
   public VelocityForgeHandshakeSessionHandler(VelocityForgeClientConnectionPhase phase) {
     this.phase = phase;
@@ -13,10 +16,12 @@ public class VelocityForgeHandshakeSessionHandler implements MinecraftSessionHan
 
   @Override
   public boolean handle(LoginPluginResponse packet) {
-    //TODO: Check if we sent it
-    if (packet.getId() == 1 && packet.getId() == 98) {
-      phase.handle(packet);
+    if (listenerList.removeIf(id -> id.equals(packet.getId()))) {
+      phase.handle(packet, listenerList.isEmpty());
     }
     return true;
   }
+  public void listen(int id) {
+    listenerList.add(id);
+  }
 }