Resync tracker WIP

build.gradle

@@ -4,7 +4,7 @@ plugins {
 }
 
 group 'org.adde0109'
-version '0.3.2'
+version '0.3.3'
 
 repositories {
     maven {

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

@@ -18,6 +18,7 @@ import org.adde0109.ambassador.forge.ForgeConnection;
 import org.adde0109.ambassador.forge.ForgeHandshakeHandler;
 import org.adde0109.ambassador.forge.ForgeHandshakeUtils;
 import org.adde0109.ambassador.forge.ForgeServerConnection;
+import org.adde0109.ambassador.forge.ForgeServerSwitchHandler;
 import org.bstats.velocity.Metrics;
 import org.slf4j.Logger;
 
@@ -27,13 +28,14 @@ import java.util.*;
 @Plugin(id = "ambassador", name = "Ambassador", version = "0.3.2", authors = {"adde0109"})
 public class Ambassador {
 
-  private final ProxyServer server;
+  public ProxyServer server;
-  private final Logger logger;
+  public final Logger logger;
+  public AmbassadorConfig config;
   private final Metrics.Factory metricsFactory;
   private final Path dataDirectory;
-  private AmbassadorConfig config;
 
-  private ForgeHandshakeHandler forgeHandshakeHandler;
+  public ForgeHandshakeHandler forgeHandshakeHandler;
+  public ForgeServerSwitchHandler forgeServerSwitchHandler;
 
 
 
@@ -51,8 +53,10 @@ public class Ambassador {
 
     config = AmbassadorConfig.readOrCreateConfig(dataDirectory,server,logger);
     if(config != null) {
-      forgeHandshakeHandler = new ForgeHandshakeHandler(config, server, logger);
+      forgeHandshakeHandler = new ForgeHandshakeHandler(this);
+      forgeServerSwitchHandler = new ForgeServerSwitchHandler(this);
       server.getEventManager().register(this, forgeHandshakeHandler);
+      server.getEventManager().register(this,forgeServerSwitchHandler);
     }
     else {
       logger.warn("Ambassador will be disabled because of errors");
@@ -66,64 +70,12 @@ public class Ambassador {
     AmbassadorConfig c = AmbassadorConfig.readOrCreateConfig(dataDirectory,server,logger);
     if (config != null) {
       config = c;
-      forgeHandshakeHandler.setConfig(config);
       logger.info("Successfully reloaded the config");
     } else {
       logger.warn("Using the old config");
     }
   }
 
-  @Subscribe
-  public void onServerPreConnectEvent(ServerPreConnectEvent event, Continuation continuation) {
-    Optional<ForgeServerConnection> forgeServerConnectionOptional = forgeHandshakeHandler.getForgeServerConnection(event.getOriginalServer());
-    if (forgeServerConnectionOptional.isPresent()) {
-      //Check 1; Check if the server is already known to us. Check if the client is compatible.
-      ForgeServerConnection 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.
-          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());
-          if (forgeConnection.isEmpty() && (event.getPlayer().getCurrentServer().isPresent())) {
-            //If vanilla tries to connect to a server we know is forge
-            event.setResult(ServerPreConnectEvent.ServerResult.denied());
-            event.getPlayer().sendMessage(Component.text("This server requires Forge!", NamedTextColor.RED));
-            continuation.resume();
-          } else if (forgeConnection.isPresent()) {
-
-            //To make legacy forwarding work
-            List<GameProfile.Property> properties = new ArrayList<>(event.getPlayer().getGameProfileProperties());
-            properties.add(new GameProfile.Property("extraData", "\1FML2\1",""));
-            event.getPlayer().setGameProfileProperties(properties);
-
-            if (forgeConnection.get().getTransmittedHandshake().isPresent()
-                    && forgeConnection.get().getRecivedClientModlist().isPresent()
-                && msg.equals(forgeConnection.get().getTransmittedHandshake().get())) {
-              //The client's registry is the same as the server's
-              continuation.resume();
-            } else {
-              event.setResult(ServerPreConnectEvent.ServerResult.denied());
-              logger.info("Kicking {} because of re-sync needed", event.getPlayer());
-              event.getPlayer().disconnect(Component.text("Please reconnect"));
-              continuation.resume();
-            }
-          } else {
-            //If the initial server is forge while the client is vanilla.
-            //Can't handle, just let it pass.
-            continuation.resume();
-          }
-        }
-      });
-    } else {
-      //The server is not known to us.
-      continuation.resume();
-    }
-  }
 
   @Subscribe
   public void onPlayerChooseInitialServerEvent(PlayerChooseInitialServerEvent event, Continuation continuation) {

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

@@ -1,6 +1,7 @@
 package org.adde0109.ambassador.forge;
 
 import com.velocitypowered.api.event.Continuation;
+import com.velocitypowered.api.event.PostOrder;
 import com.velocitypowered.api.event.Subscribe;
 import com.velocitypowered.api.event.connection.PreLoginEvent;
 import com.velocitypowered.api.event.player.KickedFromServerEvent;
@@ -20,14 +21,13 @@ import java.util.Optional;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.TranslatableComponent;
 import net.kyori.adventure.text.format.NamedTextColor;
+import org.adde0109.ambassador.Ambassador;
 import org.adde0109.ambassador.AmbassadorConfig;
 import org.slf4j.Logger;
 
 public class ForgeHandshakeHandler {
 
-  private AmbassadorConfig config;
+  private final Ambassador ambassador;
-  private final ProxyServer server;
-  private final Logger logger;
 
   private final Map<RegisteredServer, ForgeServerConnection>
       forgeServerConnectionMap = new HashMap<>();
@@ -36,21 +36,19 @@ public class ForgeHandshakeHandler {
   private static final ChannelIdentifier LOGIN_WRAPPER_ID = MinecraftChannelIdentifier.create("fml","loginwrapper");
 
 
-  public ForgeHandshakeHandler(AmbassadorConfig config, ProxyServer server, Logger logger) {
+  public ForgeHandshakeHandler(Ambassador ambassador) {
-    this.config = config;
+    this.ambassador = ambassador;
-    this.server = server;
-    this.logger = logger;
   }
 
-  @Subscribe
+  @Subscribe(order = PostOrder.LAST)
   public void onPreLoginEvent(PreLoginEvent event, Continuation continuation) {
-    if (!config.shouldHandle(event.getConnection().getProtocolVersion().getProtocol())) {
+    if (!ambassador.config.shouldHandle(event.getConnection().getProtocolVersion().getProtocol()) || !event.getResult().isAllowed()) {
       continuation.resume();
       return;
     }
-    RegisteredServer defaultServer = config.getServer(event.getConnection().getProtocolVersion().getProtocol());
+    RegisteredServer defaultServer = ambassador.config.getServer(event.getConnection().getProtocolVersion().getProtocol());
 
-    ForgeConnection forgeConnection = new ForgeConnection((LoginPhaseConnection) event.getConnection(), logger);
+    ForgeConnection forgeConnection = new ForgeConnection((LoginPhaseConnection) event.getConnection(), ambassador.logger);
     forgeConnection.testIfForge((LoginPhaseConnection) event.getConnection())
         .thenAccept((isForge) -> {
           if (isForge)
@@ -98,10 +96,6 @@ public class ForgeHandshakeHandler {
     forgeServerConnectionMap.remove(server);
   }
 
-  public void setConfig(AmbassadorConfig config) {
-    this.config = config;
-  }
-
   @Subscribe
   public void onServerLoginPluginMessageEvent(ServerLoginPluginMessageEvent event, Continuation continuation) {
     if (!event.getIdentifier().equals(LOGIN_WRAPPER_ID)) {
@@ -132,7 +126,7 @@ public class ForgeHandshakeHandler {
         if (((TranslatableComponent) reason).key().equals("multiplayer.disconnect.unexpected_query_response")) {
           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.",
+            ambassador.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);

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

@@ -0,0 +1,120 @@
+package org.adde0109.ambassador.forge;
+
+import com.velocitypowered.api.event.Continuation;
+import com.velocitypowered.api.event.PostOrder;
+import com.velocitypowered.api.event.Subscribe;
+import com.velocitypowered.api.event.player.ServerPreConnectEvent;
+import com.velocitypowered.api.proxy.server.RegisteredServer;
+import com.velocitypowered.api.scheduler.ScheduledTask;
+import com.velocitypowered.api.scheduler.TaskStatus;
+import com.velocitypowered.api.util.GameProfile;
+
+import java.net.InetSocketAddress;
+import java.util.*;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import net.kyori.adventure.text.Component;
+import net.kyori.adventure.text.format.NamedTextColor;
+import org.adde0109.ambassador.Ambassador;
+
+public class ForgeServerSwitchHandler {
+
+  private final Ambassador ambassador;
+  public final ReSyncTracker reSyncTracker;
+
+  public ForgeServerSwitchHandler(Ambassador ambassador) {
+    this.ambassador = ambassador;
+    this.reSyncTracker = new ReSyncTracker();
+  }
+
+
+  @Subscribe(order = PostOrder.LAST)
+  public void onServerPreConnectEvent(ServerPreConnectEvent event, Continuation continuation) {
+    if (!event.getResult().isAllowed()) {
+      continuation.resume();
+      return;
+    }
+    Optional<ForgeServerConnection> forgeServerConnectionOptional = ambassador.forgeHandshakeHandler.getForgeServerConnection(event.getOriginalServer());
+    if (forgeServerConnectionOptional.isPresent()) {
+      //Check 1; Check if the server is already known to us. Check if the client is compatible.
+      ForgeServerConnection 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.
+          if (ex instanceof ForgeHandshakeUtils.HandshakeReceiver.HandshakeNotAvailableException) {
+            //It's not running ambassador, so it should be unregistered.
+            ambassador.forgeHandshakeHandler.unRegisterForgeServer(forgeServerConnection.getServer());
+          }
+          continuation.resume();
+        } else {
+          Optional<ForgeConnection> forgeConnection = ambassador.forgeHandshakeHandler.getForgeConnection(event.getPlayer());
+          if (forgeConnection.isEmpty() && (event.getPlayer().getCurrentServer().isPresent())) {
+            //If vanilla tries to connect to a server we know is forge
+            event.setResult(ServerPreConnectEvent.ServerResult.denied());
+            event.getPlayer().sendMessage(Component.text("This server requires Forge!", NamedTextColor.RED));
+            continuation.resume();
+          } else if (forgeConnection.isPresent()) {
+
+            //To make legacy forwarding work
+            List<GameProfile.Property> properties = new ArrayList<>(event.getPlayer().getGameProfileProperties());
+            properties.add(new GameProfile.Property("extraData", "\1FML2\1",""));
+            event.getPlayer().setGameProfileProperties(properties);
+
+            if (forgeConnection.get().getTransmittedHandshake().isPresent()
+                && forgeConnection.get().getRecivedClientModlist().isPresent()
+                && msg.equals(forgeConnection.get().getTransmittedHandshake().get())) {
+              //The client's registry is the same as the server's
+              continuation.resume();
+            } else {
+              event.setResult(ServerPreConnectEvent.ServerResult.denied());
+              ambassador.logger.info("Kicking {} because of re-sync needed", event.getPlayer());
+              event.getPlayer().disconnect(Component.text("Please reconnect"));
+              reSyncTracker.put(event.getPlayer().getRemoteAddress(),event.getOriginalServer());
+              continuation.resume();
+            }
+          } else {
+            //If the initial server is forge while the client is vanilla.
+            //Can't handle, just let it pass.
+            continuation.resume();
+          }
+        }
+      });
+    } else {
+      //The server is not known to us.
+      continuation.resume();
+    }
+  }
+  class ReSyncTracker {
+    private static final int TIMEOUT = 2;
+    private static final int TICK_TIME = 15;
+    private Optional<ScheduledTask> scheduledTask = Optional.empty();
+    private final AtomicInteger counter = new AtomicInteger();
+    private Map<InetSocketAddress, Integer> reSyncTimeoutMap = new HashMap<>();
+    private Map<InetSocketAddress,RegisteredServer> reSyncTargetMap = new HashMap<>();
+
+    public void tick() {
+      int c = counter.incrementAndGet();
+      if (reSyncTimeoutMap.values().removeIf((v) -> c>=v))
+        reSyncTargetMap.keySet().removeIf((k) -> !reSyncTargetMap.containsKey(k));
+      //Remove if the reSyncTargetMap is empty
+      if (reSyncTargetMap.isEmpty() && scheduledTask.isPresent())
+        if (scheduledTask.get().status() == TaskStatus.SCHEDULED)
+          scheduledTask = Optional.empty();
+    }
+
+    public void put(InetSocketAddress inetSocketAddress, RegisteredServer registeredServer) {
+      reSyncTimeoutMap.put(inetSocketAddress, counter.get()+TIMEOUT);
+      reSyncTargetMap.put(inetSocketAddress, registeredServer);
+      //Start if not already started
+      if (scheduledTask.isPresent())
+        if (scheduledTask.get().status() == TaskStatus.CANCELLED || scheduledTask.get().status() == TaskStatus.FINISHED)
+          scheduledTask = Optional.of(ambassador.server.getScheduler().buildTask((ambassador), this::tick)
+                  .repeat(TICK_TIME, TimeUnit.SECONDS).schedule());
+    }
+    public RegisteredServer remove(InetSocketAddress inetSocketAddress) {
+      reSyncTimeoutMap.remove(inetSocketAddress);
+      return reSyncTargetMap.remove(inetSocketAddress);
+    }
+  }
+}