Handle forge player when no forge servers are available

build.gradle

@@ -4,7 +4,7 @@ plugins {
 }
 
 group 'org.adde0109'
-version '0.2.1-SNAPSHOT'
+version '0.2.2'
 
 repositories {
     maven {

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

@@ -22,7 +22,7 @@ import org.slf4j.Logger;
 import java.nio.file.Path;
 import java.util.*;
 
-@Plugin(id = "ambassador", name = "Ambassador", version = "0.2.1-SNAPSHOT", authors = {"adde0109"})
+@Plugin(id = "ambassador", name = "Ambassador", version = "0.2.2", authors = {"adde0109"})
 public class Ambassador {
 
   private final ProxyServer server;
@@ -68,6 +68,10 @@ public class Ambassador {
       forgeServerConnection.getHandshake().whenComplete((msg, ex) -> {
         if (ex != null) {
           //The server was forge but aren't right now. Or it's just offline.
+          if (ex instanceof ForgeHandshakeUtils.HandshakeReceiver.HandshakeNotAvailableException) {
+            //It's not running ambassador so it should be unregistered.
+            forgeHandshakeHandler.unRegisterForgeServer(forgeServerConnection.getServer());
+          }
           continuation.resume();
         } else {
           Optional<ForgeConnection> forgeConnection = forgeHandshakeHandler.getForgeConnection(event.getPlayer());
@@ -77,7 +81,8 @@ public class Ambassador {
             event.getPlayer().sendMessage(Component.text("This server requires Forge!", NamedTextColor.RED));
             continuation.resume();
           } else if (forgeConnection.isPresent()) {
-            if (msg.equals(forgeConnection.get().getTransmittedHandshake())) {
+            if (forgeConnection.get().getTransmittedHandshake().isPresent()
+                && msg.equals(forgeConnection.get().getTransmittedHandshake().get())) {
               //The client's registry is the same as the server's
               continuation.resume();
             } else {

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

@@ -5,8 +5,10 @@ import com.velocitypowered.api.event.Continuation;
 import com.velocitypowered.api.event.player.ServerLoginPluginMessageEvent;
 import com.velocitypowered.api.proxy.LoginPhaseConnection;
 import com.velocitypowered.api.proxy.messages.MinecraftChannelIdentifier;
+import com.velocitypowered.api.proxy.server.RegisteredServer;
 import java.io.EOFException;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
 public class ForgeConnection {
@@ -16,65 +18,59 @@ public class ForgeConnection {
   private byte[] recivedClientModlist;
   private static byte[] recivedClientACK;
 
-  private ForgeHandshakeUtils.CachedServerHandshake transmittedHandshake;
+  private Optional<ForgeHandshakeUtils.CachedServerHandshake> transmittedHandshake = Optional.empty();
+  private Optional<RegisteredServer> syncedTo = Optional.empty();
 
 
-  private ForgeConnection(LoginPhaseConnection connection) {
+  public ForgeConnection(LoginPhaseConnection connection) {
     this.connection = connection;
   }
 
+  public static CompletableFuture<Boolean> testIfForge(LoginPhaseConnection connection) {
+    CompletableFuture<Boolean> future = new CompletableFuture<>();
 
-  public static CompletableFuture<ForgeConnection> sync(LoginPhaseConnection connection,
+    byte[] testPacket = ForgeHandshakeUtils.generateTestPacket();
-                                                        ForgeServerConnection forgeServerConnection,
+    connection.sendLoginPluginMessage(MinecraftChannelIdentifier.create("fml", "loginwrapper"), testPacket,
-                                                        Continuation continuation) {
+        responseBody -> {
-    CompletableFuture<ForgeConnection> future = new CompletableFuture<>();
+          future.complete(responseBody != null);
-    ForgeConnection forgeConnection = new ForgeConnection(connection);
+        });
-    forgeServerConnection.getHandshake().whenComplete((msg, ex) -> {
+    return future;
+  }
+
+
+
+  public CompletableFuture<Boolean> sync(ForgeServerConnection forgeServerConnection) {
+    CompletableFuture<Boolean> future = new CompletableFuture<>();
+    forgeServerConnection.getHandshake().whenComplete((msg,ex) -> {
       if (ex != null) {
-        future.completeExceptionally(ex);
+        future.complete(false);
-      } else {
-        forgeConnection.sendModlist(msg.modListPacket).thenAccept((response) -> {
-          if (response != null) {
-            future.complete(forgeConnection);
-          } else {
-            future.complete(null);
-          }
-        });
-        forgeConnection.sendOther(msg.otherPackets).thenAccept((response) -> {
-          if (response != null) {
-            future.complete(forgeConnection);
-          } else {
-            future.complete(null);
-          }
-        });
-        forgeConnection.transmittedHandshake = msg;
       }
-      //Write
+        sendModlist(msg.modListPacket).thenAccept((response) -> {
-      continuation.resume();
+          recivedClientModlist = response;
+        });
+        sendOther(msg.otherPackets).thenAccept((response) -> {
+          ForgeConnection.recivedClientACK = response;
+          transmittedHandshake = Optional.of(msg);
+          syncedTo = Optional.of(forgeServerConnection.getServer());
+        });
+      future.complete(true);
     });
     return future;
   }
 
-  public CompletableFuture<byte[]> sendModlist(byte[] modListPacket) {
+  private CompletableFuture<byte[]> sendModlist(byte[] modListPacket) {
     CompletableFuture<byte[]> future = new CompletableFuture<>();
     connection.sendLoginPluginMessage(MinecraftChannelIdentifier.create("fml", "loginwrapper"), modListPacket,
-        responseBody -> {
+        future::complete);
-          recivedClientModlist = responseBody;
-          future.complete(recivedClientModlist);
-        });
     return future;
   }
 
-  CompletableFuture<byte[]> sendOther(List<byte[]> otherPackets) {
+  private CompletableFuture<byte[]> sendOther(List<byte[]> otherPackets) {
     CompletableFuture<byte[]> future = new CompletableFuture<>();
     for (int i = 0; i < otherPackets.size(); i++) {
       connection.sendLoginPluginMessage(MinecraftChannelIdentifier.create("fml", "loginwrapper"), otherPackets.get(i),
           (i < (otherPackets.size() - 1)) ? responseBody -> {
-          } : responseBody -> {
+          } : future::complete);
-            if (responseBody != null)
-              recivedClientACK = responseBody;
-            future.complete(responseBody);
-          });
     }
     return future;
   }
@@ -100,7 +96,7 @@ public class ForgeConnection {
     return connection;
   }
 
-  public ForgeHandshakeUtils.CachedServerHandshake getTransmittedHandshake() {
+  public Optional<ForgeHandshakeUtils.CachedServerHandshake> getTransmittedHandshake() {
     return transmittedHandshake;
   }
 

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

@@ -50,26 +50,27 @@ public class ForgeHandshakeHandler {
     }
     RegisteredServer defaultServer = config.getServer(event.getConnection().getProtocolVersion().getProtocol());
 
-          if (defaultServer == null) {
+    ForgeConnection forgeConnection = new ForgeConnection((LoginPhaseConnection) event.getConnection());
-            continuation.resume();
+    ForgeConnection.testIfForge((LoginPhaseConnection) event.getConnection())
-            return;
+        .thenAccept((isForge) -> {
-          }
+          registerForgeConnection(forgeConnection);
+        });
 
-
+    if (defaultServer != null) {
-
+      //If a connection does not already exist, create one.
-          //If a connection does not already exist, create one.
+      if (!forgeServerConnectionMap.containsKey(defaultServer)) {
-          if (!forgeServerConnectionMap.containsKey(defaultServer)) {
+        forgeServerConnectionMap.put(defaultServer, new ForgeServerConnection(defaultServer));
-            forgeServerConnectionMap.put(defaultServer, new ForgeServerConnection(defaultServer));
+      }
-          }
+      //Forge Handshake
-
+      forgeConnection.sync(forgeServerConnectionMap.get(defaultServer)).thenAccept((done) -> {
-          ForgeServerConnection forgeServerConnection = forgeServerConnectionMap.get(defaultServer);
+        continuation.resume();
-
+      });
-          //Syncing - continuation is forwarded to this method
+    } else {
-          ForgeConnection.sync((LoginPhaseConnection) event.getConnection(),forgeServerConnection,continuation).thenAccept(
+      continuation.resume();
-              this::onSyncComplete);
+    }
   }
 
-  private void onSyncComplete(ForgeConnection forgeConnection) {
+  private void registerForgeConnection(ForgeConnection forgeConnection) {
     if (forgeConnection != null) {
       incomingForgeConnections.values().removeIf((c) -> !c.getConnection().isActive());
       incomingForgeConnections.put(forgeConnection.getConnection().getRemoteAddress(), forgeConnection);
@@ -92,6 +93,9 @@ public class ForgeHandshakeHandler {
   public void registerForgeServer(RegisteredServer server, ForgeServerConnection forgeServerConnection) {
     forgeServerConnectionMap.put(server,forgeServerConnection);
   }
+  public void unRegisterForgeServer(RegisteredServer server) {
+    forgeServerConnectionMap.remove(server);
+  }
 
   @Subscribe
   public void onServerLoginPluginMessageEvent(ServerLoginPluginMessageEvent event, Continuation continuation) {

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

@@ -1,6 +1,8 @@
 package org.adde0109.ambassador.forge;
 
 import com.google.common.io.ByteArrayDataInput;
+import com.google.common.io.ByteArrayDataOutput;
+import com.google.common.io.ByteStreams;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 import com.velocitypowered.api.proxy.server.ServerPing;
 import com.velocitypowered.api.util.ModInfo;
@@ -28,6 +30,36 @@ public class ForgeHandshakeUtils {
     return i;
   }
 
+  public static void writeVarInt(ByteArrayDataOutput stream,int i) {
+    while((i & -128) != 0) {
+      stream.writeByte(i & 127 | 128);
+      i >>>= 7;
+    }
+
+    stream.writeByte(i);
+  }
+
+  public static void writeUtf(ByteArrayDataOutput stream,String p_211400_1_) {
+    byte[] abyte = p_211400_1_.getBytes(StandardCharsets.UTF_8);
+      writeVarInt(stream,abyte.length);
+      stream.write(abyte);
+  }
+
+  public static byte[] generateTestPacket() {
+    ByteArrayDataOutput dataAndPacketIdStream = ByteStreams.newDataOutput();
+    writeVarInt(dataAndPacketIdStream,4);
+    writeUtf(dataAndPacketIdStream,"ambassadortestpacket");
+    writeVarInt(dataAndPacketIdStream,0);
+
+    ByteArrayDataOutput stream = ByteStreams.newDataOutput();
+    byte[] dataAndPacketId = dataAndPacketIdStream.toByteArray();
+    writeUtf(stream,"fml:handshake");
+    writeVarInt(stream,dataAndPacketId.length);
+    stream.write(dataAndPacketId);
+
+    return stream.toByteArray();
+  }
+
   public static class HandshakeReceiver {
 
     private int partLength;
@@ -45,7 +77,7 @@ public class ForgeHandshakeUtils {
 
     private HandshakeReceiver(ServerPing serverPing) throws Exception {
       if ((serverPing.getModinfo().isEmpty()) || (!Objects.equals(serverPing.getModinfo().get().getType(), "ambassador"))) {
-        throw new Exception("The specified Forge server is not running the Forge-side version of this plugin!");
+        throw new HandshakeNotAvailableException("The specified Forge server is not running the Forge-side version of this plugin!");
       }
 
       ModInfo.Mod pair = serverPing.getModinfo().orElseThrow(IllegalAccessError::new).getMods().get(0);
@@ -160,7 +192,11 @@ public class ForgeHandshakeUtils {
       return list;
     }
 
-
+    public static class HandshakeNotAvailableException extends Exception {
+      HandshakeNotAvailableException(String errorMessage) {
+        super(errorMessage);
+      }
+    }
 
   }
   public static class CachedServerHandshake {