Restruction. Tested

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

@@ -1,9 +1,11 @@
 package org.adde0109.ambassador;
 
+import com.google.common.io.ByteArrayDataInput;
 import com.google.inject.Inject;
 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.ServerLoginPluginMessageEvent;
 import com.velocitypowered.api.event.proxy.ProxyInitializeEvent;
 import com.velocitypowered.api.plugin.Plugin;
 import com.velocitypowered.api.plugin.annotation.DataDirectory;
@@ -11,8 +13,9 @@ import com.velocitypowered.api.proxy.LoginPhaseConnection;
 import com.velocitypowered.api.proxy.ProxyServer;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 
+import java.io.EOFException;
 import java.net.InetSocketAddress;
-import java.util.concurrent.atomic.AtomicBoolean;
+import org.adde0109.ambassador.event.PreSyncEvent;
 import org.slf4j.Logger;
 
 import java.nio.file.Path;
@@ -29,8 +32,8 @@ public class Ambassador {
 
   private static ForgeHandshakeDataHandler forgeHandshakeDataHandler;
 
-  public Map<RegisteredServer, ForgeServerConnection> forgeServerConnectionMap;
-  public Map<InetSocketAddress,ForgeConnection> incomingForgeConnections;
+  public Map<RegisteredServer, ForgeServerConnection> forgeServerConnectionMap = new HashMap<RegisteredServer,ForgeServerConnection>();
+  public Map<InetSocketAddress,ForgeConnection> incomingForgeConnections = new HashMap<InetSocketAddress,ForgeConnection>();;
 
   @Inject
   public Ambassador(ProxyServer server, Logger logger, @DataDirectory Path dataDirectory) {
@@ -67,7 +70,7 @@ public class Ambassador {
           }
           RegisteredServer newServer = e.getResult().getServer().get();
 
-          ForgeConnection forgeConnection =  new ForgeConnection((LoginPhaseConnection) event.getConnection());
+
 
           //If a connection does not already exist, create one.
           if (!forgeServerConnectionMap.containsKey(newServer)) {
@@ -76,38 +79,38 @@ public class Ambassador {
 
           ForgeServerConnection forgeServerConnection = forgeServerConnectionMap.get(newServer);
 
-          //Syncing
-          forgeServerConnection.getHandshake().whenComplete((msg,ex) -> {
-            if (ex != null) {
-              logger.warn("Could not sync player '" + event.getUsername() + "' to server '"
-                  + forgeServerConnection.getServerInfo().getName() +"' Cause: " + ex.getMessage());
-            } else {
-              forgeConnection.sendModlist(msg.modListPacket).thenAccept((response) -> {
-                if (response != null) {
-                  forgeServerConnection.setDefaultClientModlist(response);
-                }
-              });
-              forgeConnection.sendOther(msg.otherPackets).thenAccept((response) -> {
-                if (response != null) {
-                  forgeServerConnection.setDefaultClientACK(response);
-                }
-                onSyncComplete(forgeConnection);
-              });
-            }
-            //Writes the messages
-            continuation.resume();
-          });
+          //Syncing - continuation is forwarded to this method
+          ForgeConnection.sync((LoginPhaseConnection) event.getConnection(),forgeServerConnection,continuation).thenAccept(
+              this::onSyncComplete);
         });
   }
 
   public void onSyncComplete(ForgeConnection forgeConnection) {
-    if (forgeConnection.isModded()) {
+    if (forgeConnection != null) {
       incomingForgeConnections.values().removeIf((c) -> !c.getConnection().isActive());
       incomingForgeConnections.put(forgeConnection.getConnection().getRemoteAddress(), forgeConnection);
     }
   }
 
 
+  private ForgeConnection getForgeConnection(InetSocketAddress socketAddress) {
+    incomingForgeConnections.values().removeIf((c) -> !c.getConnection().isActive());
+    return incomingForgeConnections.get(socketAddress);
+  }
+
+
+
+  @Subscribe
+  public void onServerLoginPluginMessageEvent(ServerLoginPluginMessageEvent event, Continuation continuation) {
+    //Only respond the servers that we can respond to
+    if(!forgeServerConnectionMap.containsKey(event.getConnection().getServer())) {
+      continuation.resume();
+      return;
+    }
+    //Grab the connection responsible for this - no pun intended
+    ForgeServerConnection connection = forgeServerConnectionMap.get(event.getConnection().getServer());
+    connection.handle(event,continuation);
+  }
 
 
 

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

@@ -1,5 +1,6 @@
 package org.adde0109.ambassador;
 
+import com.velocitypowered.api.event.Continuation;
 import com.velocitypowered.api.proxy.InboundConnection;
 import com.velocitypowered.api.proxy.LoginPhaseConnection;
 import com.velocitypowered.api.proxy.messages.MinecraftChannelIdentifier;
@@ -13,19 +14,48 @@ public class ForgeConnection {
 
   private byte[] recivedClientModlist;
 
-  private boolean isModded = false;
+  private ForgeHandshakeDataHandler.CachedServerHandshake transmittedHandshake;
 
-  ForgeConnection(LoginPhaseConnection connection) {
+
+  private ForgeConnection(LoginPhaseConnection connection) {
     this.connection = connection;
   }
 
+
+  public static CompletableFuture<ForgeConnection> sync(LoginPhaseConnection connection, ForgeServerConnection forgeServerConnection, Continuation continuation) {
+    CompletableFuture<ForgeConnection> future = new CompletableFuture<ForgeConnection>();
+    ForgeConnection forgeConnection = new ForgeConnection(connection);
+    forgeServerConnection.getHandshake().whenComplete((msg,ex) -> {
+      if (ex != null) {
+        future.completeExceptionally(ex);
+      } else {
+        forgeConnection.sendModlist(msg.modListPacket).thenAccept((response) -> {
+          if (response != null) {
+            forgeServerConnection.setDefaultClientModlist(response);
+            future.complete(forgeConnection);
+          } else {
+            future.complete(null);
+          }
+        });
+        forgeConnection.sendOther(msg.otherPackets).thenAccept((response) -> {
+          if (response != null) {
+            forgeServerConnection.setDefaultClientACK(response);
+            future.complete(forgeConnection);
+          } else {
+            future.complete(null);
+          }
+        });
+      }
+      //Write
+      continuation.resume();
+    });
+    return future;
+  }
+
   public CompletableFuture<byte[]> sendModlist(byte[] modListPacket) {
     CompletableFuture<byte[]> future = new CompletableFuture<byte[]>();
     connection.sendLoginPluginMessage(MinecraftChannelIdentifier.create("fml","loginwrapper"), modListPacket,
         responseBody -> {
-      if (responseBody != null) {
-        isModded = true;
-      }
       recivedClientModlist = responseBody;
       future.complete(recivedClientModlist);
         });
@@ -41,9 +71,6 @@ public class ForgeConnection {
     return future;
   }
 
-  public boolean isModded() {
-    return isModded;
-  }
   public LoginPhaseConnection getConnection() {
     return connection;
   }

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

@@ -46,95 +46,6 @@ public class ForgeHandshakeDataHandler {
   }
 
 
-  @Subscribe
-  public void onServerPostConnectEvent(ServerPostConnectEvent event) {
-    RegisteredServer server = getSyncedServer(event.getPlayer()).orElse(null);
-    if(server != null) {
-      event.getPlayer().sendMessage(Component.text("Synced to Server: " + server.getServerInfo().getName()));
-    }
-  }
-
-
-  public Optional<RegisteredServer> getSyncedServer(Player player) {
-    return getSyncedServer(player.getRemoteAddress());
-  }
-
-  private Optional<RegisteredServer> getSyncedServer(InetSocketAddress socketAddress) {
-    syncedConnections.keySet().removeIf(c -> !c.isActive());
-    InboundConnection key = syncedConnections.keySet().stream()
-        .filter((c) -> c.getRemoteAddress()==socketAddress).findFirst().orElse(null);
-      return Optional.ofNullable(syncedConnections.get(key));
-  }
-
-  @Subscribe
-  public void onServerLoginPluginMessageEvent(ServerLoginPluginMessageEvent event, Continuation continuation) {
-    if((!recivedClientModlist.containsKey(event.getConnection().getServer())) || (recivedClientACK == null)) {
-      continuation.resume();
-      return;
-    }
-    ByteArrayDataInput data = event.contentsAsDataStream();
-    if(data.skipBytes(PACKET_LENGTH_INDEX) != PACKET_LENGTH_INDEX) {  //Channel Identifier
-      continuation.resumeWithException(new EOFException());
-      return;
-    }
-    readVarInt(data); //Length
-    int packetID = readVarInt(data);
-
-    if(packetID == 1) {
-      event.setResult(ServerLoginPluginMessageEvent.ResponseResult.reply(recivedClientModlist.get(event.getConnection().getServer())));
-    }
-    else {
-      event.setResult(ServerLoginPluginMessageEvent.ResponseResult.reply(recivedClientACK));
-    }
-    continuation.resume();
-
-  }
-
-
-  public CompletableFuture<CachedServerHandshake> getHandshake(RegisteredServer handshakeServer) {
-    CompletableFuture<CachedServerHandshake> future;
-    if((handshakeServer.getPlayersConnected().isEmpty()) || (!cachedServerHandshake.containsKey(handshakeServer))) {
-      handshakeReceiver receiver = new handshakeReceiver(handshakeServer, logger);
-      future = receiver.downloadHandshake();
-      future.thenAccept(p -> {
-        cachedServerHandshake.put(handshakeServer,p);
-      });
-      return future;
-    }
-    else {
-      future = new CompletableFuture<>();
-      future.complete(cachedServerHandshake.get(handshakeServer));
-      return future;
-    }
-  }
-
-  //Should return ForgePlayer instead of Boolean in the future...
-  CompletableFuture<Boolean> sendModlist(byte[] modListPacket, LoginPhaseConnection connection, RegisteredServer server) {
-    CompletableFuture<Boolean> future = new CompletableFuture<Boolean>();
-    connection.sendLoginPluginMessage(MinecraftChannelIdentifier.create("fml","loginwrapper"), modListPacket, responseBody ->  {
-      if (responseBody != null) {
-        recivedClientModlist.put(server,responseBody);
-        future.complete(true);
-      } else {
-        future.complete(false);
-      }
-    });
-    return future;
-  }
-
-  CompletableFuture<Void> sendOther(List<byte[]> otherPackets, LoginPhaseConnection connection, RegisteredServer server) {
-    CompletableFuture<Void> future = new CompletableFuture<Void>();
-    for (int i = 0;i<otherPackets.size();i++) {
-      connection.sendLoginPluginMessage(MinecraftChannelIdentifier.create("fml","loginwrapper"), otherPackets.get(i),
-              (i<(otherPackets.size()-1)) ? responseBody -> recivedClientACK = responseBody : responseBody -> {
-        syncedConnections.put(connection,server);
-        syncedConnections.keySet().removeIf(c -> !c.isActive());
-        future.complete(null);
-              });
-    }
-    return future;
-  }
-
   public static int readVarInt(ByteArrayDataInput stream) {
     int i = 0;
     int j = 0;

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

@@ -1,18 +1,21 @@
 package org.adde0109.ambassador;
 
+import com.google.common.io.ByteArrayDataInput;
+import com.velocitypowered.api.event.Continuation;
+import com.velocitypowered.api.event.player.ServerLoginPluginMessageEvent;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 import com.velocitypowered.api.proxy.server.ServerInfo;
+import java.io.EOFException;
 import java.util.concurrent.CompletableFuture;
 import org.slf4j.Logger;
 
 public class ForgeServerConnection {
 
+  private static final int PACKET_LENGTH_INDEX = 14;    //length of "fml:handshake"+1
   private final Ambassador ambassador;
   private final Logger logger;
   private final RegisteredServer handshakeServer;
-
   private ForgeHandshakeDataHandler.CachedServerHandshake handshake;
-
   private byte[] defaultClientModlist;
   private byte[] defaultClientACK;
 
@@ -43,6 +46,24 @@ public class ForgeServerConnection {
     }
   }
 
+  public void handle(ServerLoginPluginMessageEvent event, Continuation continuation) {
+    ByteArrayDataInput data = event.contentsAsDataStream();
+    if(data.skipBytes(PACKET_LENGTH_INDEX) != PACKET_LENGTH_INDEX) {  //Channel Identifier
+      continuation.resumeWithException(new EOFException());
+      return;
+    }
+    ForgeHandshakeDataHandler.readVarInt(data); //Length
+    int packetID = ForgeHandshakeDataHandler.readVarInt(data);
+
+    if(packetID == 1) {
+      event.setResult(ServerLoginPluginMessageEvent.ResponseResult.reply(defaultClientModlist));
+    }
+    else {
+      event.setResult(ServerLoginPluginMessageEvent.ResponseResult.reply(defaultClientACK));
+    }
+    continuation.resume();
+  }
+
   public void setDefaultClientModlist(byte[] modlist) {
     this.defaultClientModlist = modlist;
   }

src/main/java/org/adde0109/ambassador/event/PostSyncEvent.java

@@ -1,4 +1,4 @@
-package org.adde0109.ambassador;
+package org.adde0109.ambassador.event;
 
 public class PostSyncEvent {
 

src/main/java/org/adde0109/ambassador/event/PreSyncEvent.java

@@ -1,4 +1,4 @@
-package org.adde0109.ambassador;
+package org.adde0109.ambassador.event;
 
 import com.velocitypowered.api.event.ResultedEvent;
 import com.velocitypowered.api.event.player.ServerPreConnectEvent;