Merge pull request #123 from mikroskeem/feature/native-forced-hosts

[WIP] Forced hosts configuration support

api/src/main/java/com/velocitypowered/api/proxy/config/ProxyConfig.java

@@ -66,6 +66,12 @@ public interface ProxyConfig {
      */
     List<String> getAttemptConnectionOrder();
 
+    /**
+     * Get forced servers mapped to given virtual host
+     * @return list of server names
+     */
+    Map<String, List<String>> getForcedHosts();
+
     /**
      * Get the minimum compression threshold for packets
      * @return the compression threshold

proxy/src/main/java/com/velocitypowered/proxy/config/AnnotatedConfig.java

@@ -120,7 +120,7 @@ public class AnnotatedConfig {
                     if (field.getAnnotation(IsMap.class) != null) { // check if field is map
                         Map<String, ?> map = (Map<String, ?>) field.get(toSave);
                         for (Entry<String, ?> entry : map.entrySet()) {
-                            lines.add(entry.getKey() + " = " + toString(entry.getValue())); // Save a map data
+                            lines.add(safeKey(entry.getKey()) + " = " + toString(entry.getValue())); // Save a map data
                         }
                         lines.add(""); //Add empty line
                         continue;
@@ -165,6 +165,13 @@ public class AnnotatedConfig {
         return value != null ? value.toString() : "null";
     }
 
+    private String safeKey(String key) {
+        if(key.contains(".") && !(key.indexOf('"') == 0 && key.lastIndexOf('"') == (key.length() - 1))) {
+            return '"' + key + '"';
+        }
+        return key;
+    }
+
     /**
      * Saves lines to file
      *

proxy/src/main/java/com/velocitypowered/proxy/config/VelocityConfiguration.java

@@ -2,6 +2,7 @@ package com.velocitypowered.proxy.config;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableList;
 import com.moandjiezana.toml.Toml;
 import com.velocitypowered.api.proxy.config.ProxyConfig;
 import com.velocitypowered.api.util.Favicon;
@@ -65,6 +66,9 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
     @Table("[servers]")
     private final Servers servers;
 
+    @Table("[forced-hosts]")
+    private final ForcedHosts forcedHosts;
+
     @Table("[advanced]")
     private final Advanced advanced;
 
@@ -76,15 +80,16 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
     @Ignore
     private Favicon favicon;
 
-    public VelocityConfiguration(Servers servers, Advanced advanced, Query query) {
+    public VelocityConfiguration(Servers servers, ForcedHosts forcedHosts, Advanced advanced, Query query) {
         this.servers = servers;
+        this.forcedHosts = forcedHosts;
         this.advanced = advanced;
         this.query = query;
     }
 
     private VelocityConfiguration(String bind, String motd, int showMaxPlayers, boolean onlineMode,
             boolean announceForge, PlayerInfoForwarding playerInfoForwardingMode, byte[] forwardingSecret,
-            Servers servers, Advanced advanced, Query query) {
+            Servers servers, ForcedHosts forcedHosts, Advanced advanced, Query query) {
         this.bind = bind;
         this.motd = motd;
         this.showMaxPlayers = showMaxPlayers;
@@ -93,6 +98,7 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
         this.playerInfoForwardingMode = playerInfoForwardingMode;
         this.forwardingSecret = forwardingSecret;
         this.servers = servers;
+        this.forcedHosts = forcedHosts;
         this.advanced = advanced;
         this.query = query;
     }
@@ -153,6 +159,21 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
                     valid = false;
                 }
             }
+
+            for (Map.Entry<String, List<String>> entry : forcedHosts.getForcedHosts().entrySet()) {
+                if (entry.getValue().isEmpty()) {
+                    logger.error("Forced host '{}' does not contain any servers", entry.getKey());
+                    valid = false;
+                    continue;
+                }
+
+                for (String server : entry.getValue()) {
+                    if (!servers.getServers().containsKey(server)) {
+                        logger.error("Server '{}' for forced host '{}' does not exist", server, entry.getKey());
+                        valid = false;
+                    }
+                }
+            }
         }
 
         try {
@@ -257,6 +278,10 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
         return servers.getAttemptConnectionOrder();
     }
 
+    public Map<String, List<String>> getForcedHosts() {
+        return forcedHosts.getForcedHosts();
+    }
+
     public int getCompressionThreshold() {
         return advanced.getCompressionThreshold();
     }
@@ -301,6 +326,7 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
                 .add("forwardingSecret", forwardingSecret)
                 .add("announceForge", announceForge)
                 .add("servers", servers)
+                .add("forcedHosts", forcedHosts)
                 .add("advanced", advanced)
                 .add("query", query)
                 .add("favicon", favicon)
@@ -311,7 +337,7 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
         Toml toml;
         if (!path.toFile().exists()) {
             getLogger().info("No velocity.toml found, creating one for you...");
-            return new VelocityConfiguration(new Servers(), new Advanced(), new Query());
+            return new VelocityConfiguration(new Servers(), new ForcedHosts(), new Advanced(), new Query());
         } else {
             try (Reader reader = Files.newBufferedReader(path, StandardCharsets.UTF_8)) {
                 toml = new Toml().read(reader);
@@ -319,6 +345,7 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
         }
 
         Servers servers = new Servers(toml.getTable("servers"));
+        ForcedHosts forcedHosts = new ForcedHosts(toml.getTable("forced-hosts"));
         Advanced advanced = new Advanced(toml.getTable("advanced"));
         Query query = new Query(toml.getTable("query"));
         byte[] forwardingSecret = toml.getString("forwarding-secret", "5up3r53cr3t")
@@ -333,6 +360,7 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
                 PlayerInfoForwarding.valueOf(toml.getString("player-info-forwarding-mode", "MODERN").toUpperCase()),
                 forwardingSecret,
                 servers,
+                forcedHosts,
                 advanced,
                 query
         );
@@ -421,6 +449,61 @@ public class VelocityConfiguration extends AnnotatedConfig implements ProxyConfi
 
     }
 
+    private static class ForcedHosts {
+        @IsMap
+        @Comment("Configure your forced hosts here.")
+        private Map<String, List<String>> forcedHosts = ImmutableMap.of(
+                "lobby.example.com", ImmutableList.of("lobby"),
+                "factions.example.com", ImmutableList.of("factions"),
+                "minigames.example.com", ImmutableList.of("minigames")
+        );
+
+        private ForcedHosts() {}
+
+        private ForcedHosts(Toml toml) {
+            if (toml != null) {
+                Map<String, List<String>> forcedHosts = new HashMap<>();
+                for (Map.Entry<String, Object> entry : toml.entrySet()) {
+                    if (entry.getValue() instanceof String) {
+                        forcedHosts.put(stripQuotes(entry.getKey()), ImmutableList.of((String) entry.getValue()));
+                    } else if (entry.getValue() instanceof List) {
+                        forcedHosts.put(stripQuotes(entry.getKey()), ImmutableList.copyOf((List<String>) entry.getValue()));
+                    } else {
+                        throw new IllegalStateException("Invalid value of type " + entry.getValue().getClass() + " in forced hosts!");
+                    }
+                }
+                this.forcedHosts = ImmutableMap.copyOf(forcedHosts);
+            }
+        }
+
+        private ForcedHosts(Map<String, List<String>> forcedHosts) {
+            this.forcedHosts = forcedHosts;
+        }
+
+        private Map<String, List<String>> getForcedHosts() {
+            return forcedHosts;
+        }
+
+        private void setForcedHosts(Map<String, List<String>> forcedHosts) {
+            this.forcedHosts = forcedHosts;
+        }
+
+        private static String stripQuotes(String key) {
+            int lastIndex;
+            if (key.indexOf('"') == 0 && (lastIndex = key.lastIndexOf('"')) == (key.length() - 1)) {
+                return key.substring(1, lastIndex);
+            }
+            return key;
+        }
+
+        @Override
+        public String toString() {
+            return "ForcedHosts{" +
+                    "forcedHosts=" + forcedHosts +
+                    '}';
+        }
+    }
+
     private static class Advanced {
 
         @Comment({

proxy/src/main/java/com/velocitypowered/proxy/connection/client/ConnectedPlayer.java

@@ -42,12 +42,16 @@ import net.kyori.text.serializer.ComponentSerializers;
 import net.kyori.text.serializer.PlainComponentSerializer;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 import java.net.InetSocketAddress;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -71,6 +75,9 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
     private final VelocityTabList tabList;
     private final VelocityServer server;
 
+    @MonotonicNonNull
+    private List<String> serversToTry = null;
+    
     ConnectedPlayer(VelocityServer server, GameProfile profile, MinecraftConnection connection, InetSocketAddress virtualHost) {
         this.server = server;
         this.tabList = new VelocityTabList(connection);
@@ -344,7 +351,15 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
     }
 
     Optional<RegisteredServer> getNextServerToTry() {
-        List<String> serversToTry = server.getConfiguration().getAttemptConnectionOrder();
+        if (serversToTry == null) {
+            String virtualHost = getVirtualHost().map(InetSocketAddress::getHostString).orElse("");
+            serversToTry = server.getConfiguration().getForcedHosts().getOrDefault(virtualHost, Collections.emptyList());
+        }
+
+        if (serversToTry.isEmpty()) {
+            serversToTry = server.getConfiguration().getAttemptConnectionOrder();
+        }
+
         if (tryIndex >= serversToTry.size()) {
             return Optional.empty();
         }