From 402398010edef47df6377b971e175c9a1ddde3b9 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Wed, 21 Nov 2018 15:57:21 -0500
Subject: [PATCH] Handle changing server addresses in config (not tested).

---
 .../velocitypowered/proxy/VelocityServer.java | 41 ++++++++++++++++++-
 .../connection/client/ConnectedPlayer.java    |  2 +-
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
index 78b327cef..49c69c190 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
@@ -3,6 +3,7 @@ package com.velocitypowered.proxy;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.velocitypowered.api.event.EventManager;
@@ -47,13 +48,16 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.security.KeyPair;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import net.kyori.text.Component;
 import net.kyori.text.TextComponent;
 import net.kyori.text.serializer.GsonComponentSerializer;
@@ -242,7 +246,8 @@ public class VelocityServer implements ProxyServer {
       return false;
     }
 
-    // Re-register servers
+    // Re-register servers. If we are replacing a server, we must evacuate players.
+    Collection<ConnectedPlayer> evacuate = new ArrayList<>();
     for (Map.Entry<String, String> entry : newConfiguration.getServers().entrySet()) {
       ServerInfo newInfo =
           new ServerInfo(entry.getKey(), AddressUtil.parseAddress(entry.getValue()));
@@ -250,10 +255,42 @@ public class VelocityServer implements ProxyServer {
       if (!rs.isPresent()) {
         servers.register(newInfo);
       } else if (!rs.get().getServerInfo().equals(newInfo)) {
-        throw new IllegalStateException("Unable to replace servers in flight!");
+        for (Player player : rs.get().getPlayersConnected()) {
+          if (!(player instanceof ConnectedPlayer)) {
+            throw new IllegalStateException("ConnectedPlayer not found for player " + player
+                + " in server " + rs.get().getServerInfo().getName());
+          }
+          evacuate.add((ConnectedPlayer) player);
+        }
+        servers.unregister(rs.get().getServerInfo());
+        servers.register(newInfo);
       }
     }
 
+    CountDownLatch latch = new CountDownLatch(evacuate.size());
+    for (ConnectedPlayer player : evacuate) {
+      Optional<RegisteredServer> next = player.getNextServerToTry();
+      if (next.isPresent()) {
+        player.createConnectionRequest(next.get()).connectWithIndication()
+            .whenComplete((success, ex) -> {
+              if (ex != null || success == null || !success) {
+                player.disconnect(TextComponent.of("Your server has been changed, but we could "
+                    + "not move you to any fallback servers."));
+              }
+              latch.countDown();
+            });
+      } else {
+        player.disconnect(TextComponent.of("Your server has been changed, but we could "
+            + "not move you to any fallback servers."));
+      }
+    }
+    try {
+      latch.await();
+    } catch (InterruptedException e) {
+      logger.error("Interrupted whilst moving players", e);
+      Thread.currentThread().interrupt();
+    }
+
     // If we have a new bind address, bind to it
     if (!configuration.getBind().equals(newConfiguration.getBind())) {
       this.cm.bind(newConfiguration.getBind());
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
index e06cdacfa..2c568da64 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java
@@ -375,7 +375,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
     }
   }
 
-  Optional<RegisteredServer> getNextServerToTry() {
+  public Optional<RegisteredServer> getNextServerToTry() {
     if (serversToTry == null) {
       String virtualHostStr = getVirtualHost().map(InetSocketAddress::getHostString).orElse("");
       serversToTry = server.getConfiguration().getForcedHosts()