From 8a00f859936ceaa01c1f13a79494fa44c8aecfaf Mon Sep 17 00:00:00 2001
From: Adrian Bergqvist <adde0109@gmail.com>
Date: Mon, 27 Jun 2022 19:32:16 +0200
Subject: [PATCH] Now detects when server switch is possible

---
 .../org/adde0109/ambassador/Ambassador.java   | 29 ++++++++++++++++---
 .../ambassador/Forge/ForgeConnection.java     |  5 ++++
 .../Forge/ForgeHandshakeHandler.java          | 15 ++++++----
 3 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/src/main/java/org/adde0109/ambassador/Ambassador.java b/src/main/java/org/adde0109/ambassador/Ambassador.java
index 9dbad48..216fa48 100644
--- a/src/main/java/org/adde0109/ambassador/Ambassador.java
+++ b/src/main/java/org/adde0109/ambassador/Ambassador.java
@@ -12,6 +12,7 @@ import com.velocitypowered.api.proxy.server.RegisteredServer;
 
 import org.adde0109.ambassador.Forge.ForgeConnection;
 import org.adde0109.ambassador.Forge.ForgeHandshakeHandler;
+import org.adde0109.ambassador.Forge.ForgeServerConnection;
 import org.checkerframework.checker.index.qual.PolyUpperBound;
 import org.slf4j.Logger;
 
@@ -52,11 +53,31 @@ public class Ambassador {
 
   @Subscribe
   public void onServerPreConnectEvent(ServerPreConnectEvent event, Continuation continuation) {
-    ForgeConnection forgeConnection = forgeHandshakeHandler.getForgeConnection(event.getPlayer());
-    if (forgeConnection != null) {
-
+    Optional<ForgeConnection> forgeConnection = forgeHandshakeHandler.getForgeConnection(event.getPlayer());
+    Optional<ForgeServerConnection> forgeServerConnectionOptional = forgeHandshakeHandler.getForgeServerConnection(event.getOriginalServer());
+    if (forgeConnection.isPresent()) {
+      ForgeServerConnection forgeServerConnection;
+      if (forgeServerConnectionOptional.isEmpty()) {
+        forgeServerConnection = new ForgeServerConnection(event.getOriginalServer(), logger);
+      } else {
+        forgeServerConnection = forgeServerConnectionOptional.get();
+      }
+      forgeServerConnection.getHandshake().whenComplete((msg, ex) -> {
+        if (ex != null) {
+          continuation.resume();
+        } else {
+          if (Arrays.equals(msg.modListPacket,forgeConnection.get().getTransmittedHandshake().modListPacket)) {
+            continuation.resume();
+          } else {
+            event.setResult(ServerPreConnectEvent.ServerResult.denied());
+            logger.warn("Resync needed");
+            continuation.resume();
+          }
+        }
+      });
+    } else {
+      continuation.resume();
     }
-    continuation.resume();
   }
 
 }
diff --git a/src/main/java/org/adde0109/ambassador/Forge/ForgeConnection.java b/src/main/java/org/adde0109/ambassador/Forge/ForgeConnection.java
index 25dfbe3..3b6cd95 100644
--- a/src/main/java/org/adde0109/ambassador/Forge/ForgeConnection.java
+++ b/src/main/java/org/adde0109/ambassador/Forge/ForgeConnection.java
@@ -43,6 +43,7 @@ public class ForgeConnection {
             future.complete(null);
           }
         });
+        forgeConnection.transmittedHandshake = msg;
       }
       //Write
       continuation.resume();
@@ -73,4 +74,8 @@ public class ForgeConnection {
     return connection;
   }
 
+  public ForgeHandshakeUtils.CachedServerHandshake getTransmittedHandshake() {
+    return transmittedHandshake;
+  }
+
 }
diff --git a/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java b/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java
index 8f7fd24..38f8c45 100644
--- a/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java
+++ b/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java
@@ -11,6 +11,7 @@ import com.velocitypowered.api.proxy.server.RegisteredServer;
 import java.net.InetSocketAddress;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import org.adde0109.ambassador.AmbassadorConfig;
 import org.slf4j.Logger;
 
@@ -20,9 +21,9 @@ public class ForgeHandshakeHandler {
   private final ProxyServer server;
   private final Logger logger;
 
-  public Map<RegisteredServer, ForgeServerConnection>
+  private Map<RegisteredServer, ForgeServerConnection>
       forgeServerConnectionMap = new HashMap<RegisteredServer,ForgeServerConnection>();
-  public Map<InetSocketAddress,ForgeConnection> incomingForgeConnections = new HashMap<InetSocketAddress,ForgeConnection>();
+  private Map<InetSocketAddress,ForgeConnection> incomingForgeConnections = new HashMap<InetSocketAddress,ForgeConnection>();
 
 
   public ForgeHandshakeHandler(AmbassadorConfig config, ProxyServer server, Logger logger) {
@@ -65,13 +66,17 @@ public class ForgeHandshakeHandler {
     }
   }
 
-  public ForgeConnection getForgeConnection(Player player) {
+  public Optional<ForgeConnection> getForgeConnection(Player player) {
     return getForgeConnection(player.getRemoteAddress());
   }
 
-  private ForgeConnection getForgeConnection(InetSocketAddress socketAddress) {
+  private Optional<ForgeConnection> getForgeConnection(InetSocketAddress socketAddress) {
     incomingForgeConnections.values().removeIf((c) -> !c.getConnection().isActive());
-    return incomingForgeConnections.get(socketAddress);
+    return Optional.ofNullable(incomingForgeConnections.get(socketAddress));
+  }
+
+  public Optional<ForgeServerConnection> getForgeServerConnection(RegisteredServer registeredServer) {
+    return Optional.ofNullable(forgeServerConnectionMap.get(registeredServer));
   }