From a47f2f8d1a069a815ea750a49a593ea23138ce38 Mon Sep 17 00:00:00 2001
From: Adrian Bergqvist <adde0109@gmail.com>
Date: Tue, 28 Jun 2022 19:46:09 +0200
Subject: [PATCH] More Server Switching fixes

---
 .../org/adde0109/ambassador/Ambassador.java   | 23 +++++++++++++++++++
 .../adde0109/ambassador/AmbassadorConfig.java |  4 ++++
 .../Forge/ForgeHandshakeHandler.java          |  3 ++-
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/adde0109/ambassador/Ambassador.java b/src/main/java/org/adde0109/ambassador/Ambassador.java
index 20e9bd2..f5bafc4 100644
--- a/src/main/java/org/adde0109/ambassador/Ambassador.java
+++ b/src/main/java/org/adde0109/ambassador/Ambassador.java
@@ -2,7 +2,9 @@ package org.adde0109.ambassador;
 
 import com.google.inject.Inject;
 import com.velocitypowered.api.event.Continuation;
+import com.velocitypowered.api.event.PostOrder;
 import com.velocitypowered.api.event.Subscribe;
+import com.velocitypowered.api.event.player.PlayerChooseInitialServerEvent;
 import com.velocitypowered.api.event.player.ServerPreConnectEvent;
 import com.velocitypowered.api.event.proxy.ProxyInitializeEvent;
 import com.velocitypowered.api.plugin.Plugin;
@@ -10,6 +12,7 @@ import com.velocitypowered.api.plugin.annotation.DataDirectory;
 import com.velocitypowered.api.proxy.ProxyServer;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 
+import net.kyori.adventure.text.Component;
 import org.adde0109.ambassador.Forge.ForgeConnection;
 import org.adde0109.ambassador.Forge.ForgeHandshakeHandler;
 import org.adde0109.ambassador.Forge.ForgeServerConnection;
@@ -75,6 +78,7 @@ public class Ambassador {
             continuation.resume();
           }
         }
+        //Register newly discovered forge server
         if (forgeServerConnectionOptional.isEmpty()) {
           forgeServerConnection.setDefaultClientModlist(forgeConnection.get().getRecivedClientModlist());
           forgeServerConnection.setDefaultClientACK(ForgeConnection.getRecivedClientACK());
@@ -82,9 +86,28 @@ public class Ambassador {
         }
 
       });
+      //If vanilla tries to connect to forge
+    } else if (forgeServerConnectionOptional.isPresent() && (event.getPreviousServer() != null)){
+      event.setResult(ServerPreConnectEvent.ServerResult.denied());
+      event.getPlayer().sendMessage(Component.text("This server requires Forge!"));
+      continuation.resume();
     } else {
       continuation.resume();
     }
   }
 
+
+  @Subscribe
+  public void onPlayerChooseInitialServerEvent(PlayerChooseInitialServerEvent event, Continuation continuation) {
+    //Only handle Forge connections
+    if((event.getInitialServer().isPresent()) && (forgeHandshakeHandler.getForgeConnection(event.getPlayer()).isPresent())) {
+      //Forge client
+      ForgeConnection forgeConnection = forgeHandshakeHandler.getForgeConnection(event.getPlayer()).get();
+      if (config.getForced(forgeConnection.getConnection().getProtocolVersion().getProtocol())) {
+        event.setInitialServer(config.getServer(forgeConnection.getConnection().getProtocolVersion().getProtocol()));
+      }
+    }
+    continuation.resume();
+  }
+
 }
diff --git a/src/main/java/org/adde0109/ambassador/AmbassadorConfig.java b/src/main/java/org/adde0109/ambassador/AmbassadorConfig.java
index ae97fab..3bfb033 100644
--- a/src/main/java/org/adde0109/ambassador/AmbassadorConfig.java
+++ b/src/main/java/org/adde0109/ambassador/AmbassadorConfig.java
@@ -35,6 +35,10 @@ public class AmbassadorConfig {
     return settings.differentiators.get(protocolVersion).handshakeServer;
   }
 
+  public boolean getForced (int protocolVersion) {
+    return settings.differentiators.get(protocolVersion).forced;
+  }
+
   public boolean shouldHandle(int protocolVersion) {
     return settings.differentiators.containsKey(protocolVersion);
   }
diff --git a/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java b/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java
index cdb3695..7e8c416 100644
--- a/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java
+++ b/src/main/java/org/adde0109/ambassador/Forge/ForgeHandshakeHandler.java
@@ -86,7 +86,8 @@ public class ForgeHandshakeHandler {
   @Subscribe
   public void onServerLoginPluginMessageEvent(ServerLoginPluginMessageEvent event, Continuation continuation) {
     //Only respond the servers that we can respond to
-    if(!forgeServerConnectionMap.containsKey(event.getConnection().getServer())) {
+    if((!forgeServerConnectionMap.containsKey(event.getConnection().getServer())
+        || (getForgeConnection(event.getConnection().getPlayer()).isEmpty()))) {
       continuation.resume();
       return;
     }