Better handling of vanilla -> forge connections

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

@@ -65,48 +65,43 @@ public class Ambassador {
 
   @Subscribe
   public void onServerPreConnectEvent(ServerPreConnectEvent event, Continuation continuation) {
-    Optional<ForgeConnection> forgeConnection = forgeHandshakeHandler.getForgeConnection(event.getPlayer());
     Optional<ForgeServerConnection> forgeServerConnectionOptional = forgeHandshakeHandler.getForgeServerConnection(event.getOriginalServer());
-    if (forgeConnection.isPresent()) {
+    if (forgeServerConnectionOptional.isPresent()) {
-      ForgeServerConnection forgeServerConnection;
+      //Check 1; Check if the server is already known to us. Check if the client is compatible.
-      if (forgeServerConnectionOptional.isEmpty()) {
+      ForgeServerConnection forgeServerConnection = forgeServerConnectionOptional.get();
-        forgeServerConnection = new ForgeServerConnection(event.getOriginalServer(), logger);
-      } else {
-        forgeServerConnection = forgeServerConnectionOptional.get();
-      }
       forgeServerConnection.getHandshake().whenComplete((msg, ex) -> {
         if (ex != null) {
+          //The server was forge but aren't right now. Or it's just offline.
           continuation.resume();
         } else {
-          if (msg.equals(forgeConnection.get().getTransmittedHandshake())) {
+          Optional<ForgeConnection> forgeConnection = forgeHandshakeHandler.getForgeConnection(event.getPlayer());
-            continuation.resume();
+          if (forgeConnection.isEmpty() && (event.getPlayer().getCurrentServer().isPresent())) {
-          } else {
+            //If vanilla tries to connect to a server we know is forge
             event.setResult(ServerPreConnectEvent.ServerResult.denied());
-            logger.warn("Resync needed");
+            event.getPlayer().sendMessage(Component.text("This server requires Forge!", NamedTextColor.RED));
+            continuation.resume();
+          } else if (forgeConnection.isPresent()) {
+            if (msg.equals(forgeConnection.get().getTransmittedHandshake())) {
+              //The client's registry is the same as the server's
+              continuation.resume();
+            } else {
+              event.setResult(ServerPreConnectEvent.ServerResult.denied());
+              logger.warn("Resync needed");
+              continuation.resume();
+            }
+          } else {
+            //If the initial server is forge while the client is vanilla.
+            //Can't handle, just let it pass.
             continuation.resume();
           }
         }
-        //Register newly discovered forge server
-        if (forgeServerConnectionOptional.isEmpty()) {
-          forgeHandshakeHandler.registerForgeServer(event.getOriginalServer(), forgeServerConnection);
-        }
-
       });
-      //If vanilla tries to connect to forge
-    } else if (forgeServerConnectionOptional.isPresent() && (event.getPlayer().getCurrentServer().isPresent())){
-      event.setResult(ServerPreConnectEvent.ServerResult.denied());
-      event.getPlayer().sendMessage(Component.text("This server requires Forge!", NamedTextColor.RED));
-      continuation.resume();
     } else {
+      //The server is not known to us.
       continuation.resume();
     }
   }
 
-  @Subscribe
-  public void onKickedFromServerEvent(KickedFromServerEvent event, Continuation continuation) {
-    continuation.resume();
-  }
-
   @Subscribe
   public void onPlayerChooseInitialServerEvent(PlayerChooseInitialServerEvent event, Continuation continuation) {
     //Only handle Forge connections

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

@@ -4,16 +4,23 @@ import com.google.common.io.ByteArrayDataInput;
 import com.velocitypowered.api.event.Continuation;
 import com.velocitypowered.api.event.Subscribe;
 import com.velocitypowered.api.event.connection.PreLoginEvent;
+import com.velocitypowered.api.event.player.KickedFromServerEvent;
 import com.velocitypowered.api.event.player.ServerLoginPluginMessageEvent;
 import com.velocitypowered.api.proxy.LoginPhaseConnection;
 import com.velocitypowered.api.proxy.Player;
 import com.velocitypowered.api.proxy.ProxyServer;
+import com.velocitypowered.api.proxy.messages.ChannelIdentifier;
+import com.velocitypowered.api.proxy.messages.MinecraftChannelIdentifier;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 import java.io.EOFException;
 import java.net.InetSocketAddress;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import net.kyori.adventure.text.Component;
+import net.kyori.adventure.text.format.NamedTextColor;
 import org.adde0109.ambassador.AmbassadorConfig;
 import org.slf4j.Logger;
 
@@ -27,6 +34,9 @@ public class ForgeHandshakeHandler {
       forgeServerConnectionMap = new HashMap<>();
   private final Map<InetSocketAddress,ForgeConnection> incomingForgeConnections = new HashMap<>();
 
+  private final List<Player> doomedConnections = new ArrayList<>();
+  private static final ChannelIdentifier LOGIN_WRAPPER_ID = MinecraftChannelIdentifier.create("fml","loginwrapper");
+
 
   public ForgeHandshakeHandler(AmbassadorConfig config, ProxyServer server, Logger logger) {
     this.config = config;
@@ -87,11 +97,41 @@ public class ForgeHandshakeHandler {
 
   @Subscribe
   public void onServerLoginPluginMessageEvent(ServerLoginPluginMessageEvent event, Continuation continuation) {
+    if (!event.getIdentifier().equals(LOGIN_WRAPPER_ID)) {
+      continuation.resume();
+      return;
+    }
+    //Check 2
+    if (getForgeServerConnection(event.getConnection().getServer()).isEmpty()) {
+      registerForgeServer(event.getConnection().getServer(),
+          new ForgeServerConnection(event.getConnection().getServer(), logger));
+    }
+
     if (incomingForgeConnections.containsKey(event.getConnection().getPlayer().getRemoteAddress())) {
       incomingForgeConnections.get(event.getConnection().getPlayer().getRemoteAddress())
           .handleServerHandshakePacket(event,continuation);
     } else {
-      continuation.resume();
+        continuation.resume();
     }
   }
+
+  @Subscribe
+  public void onKickedFromServerEvent(KickedFromServerEvent event, Continuation continuation) {
+    if (getForgeServerConnection(event.getServer()).isPresent() && getForgeConnection(event.getPlayer()).isEmpty()) {
+      //Turns out the server the vanilla client is connecting to is forge. Let's handle the connection error.
+      logger.info("Vanilla player {} tried to connect to forge server {}. The connection error can be ignored.",
+          event.getPlayer(),event.getServer().getServerInfo().getName());
+      KickedFromServerEvent.ServerKickResult result = event.getResult();
+      Component component = Component.text("The server you were trying to connect to requires Forge to be installed.", NamedTextColor.RED);
+      if (result instanceof KickedFromServerEvent.DisconnectPlayer) {
+        event.setResult(KickedFromServerEvent.DisconnectPlayer.create(component));
+      } else if (result instanceof KickedFromServerEvent.RedirectPlayer) {
+        RegisteredServer redirectServer = ((KickedFromServerEvent.RedirectPlayer)event.getResult()).getServer();
+        event.setResult(KickedFromServerEvent.RedirectPlayer.create(redirectServer,component));
+      } else if (result instanceof KickedFromServerEvent.Notify) {
+        event.setResult(KickedFromServerEvent.Notify.create(component));
+      }
+    }
+    continuation.resume();
+  }
 }