From d4e89dbdda9cde355c3e365c67d5428c541afb2f Mon Sep 17 00:00:00 2001
From: Jackson <theminecoder.dev@gmail.com>
Date: Mon, 21 Oct 2024 11:09:55 +1100
Subject: [PATCH] Ability to specify servers from the command line (#1445)

---
 .../velocitypowered/proxy/ProxyOptions.java   | 54 +++++++++++++++++++
 .../velocitypowered/proxy/VelocityServer.java | 10 +++-
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java b/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java
index 64a49917d..c660bf50e 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java
@@ -17,11 +17,18 @@
 
 package com.velocitypowered.proxy;
 
+import com.velocitypowered.api.proxy.server.ServerInfo;
 import java.io.IOException;
+import java.net.InetSocketAddress;
 import java.util.Arrays;
+import java.util.List;
+
+import com.velocitypowered.proxy.util.AddressUtil;
 import joptsimple.OptionParser;
 import joptsimple.OptionSet;
 import joptsimple.OptionSpec;
+import joptsimple.ValueConversionException;
+import joptsimple.ValueConverter;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -35,6 +42,8 @@ public final class ProxyOptions {
   private final boolean help;
   private final @Nullable Integer port;
   private final @Nullable Boolean haproxy;
+  private final boolean ignoreConfigServers;
+  private final List<ServerInfo> servers;
 
   ProxyOptions(final String[] args) {
     final OptionParser parser = new OptionParser();
@@ -49,11 +58,20 @@ public final class ProxyOptions {
             "Choose whether to enable haproxy protocol. "
                     + "The configuration haproxy protocol will be ignored.")
         .withRequiredArg().ofType(Boolean.class);
+    final OptionSpec<ServerInfo> servers = parser.accepts("add-server",
+            "Define a server mapping. "
+                    + "You must ensure that server name is not also registered in the config or use --ignore-config-servers.")
+        .withRequiredArg().withValuesConvertedBy(new ServerInfoConverter());
+    final OptionSpec<Void> ignoreConfigServers = parser.accepts("ignore-config-servers",
+            "Skip registering servers from the config file. "
+                    + "Useful in dynamic setups or with the --add-server flag.");
     final OptionSet set = parser.parse(args);
 
     this.help = set.has(help);
     this.port = port.value(set);
     this.haproxy = haproxy.value(set);
+    this.servers = servers.values(set);
+    this.ignoreConfigServers = set.has(ignoreConfigServers);
 
     if (this.help) {
       try {
@@ -75,4 +93,40 @@ public final class ProxyOptions {
   public @Nullable Boolean isHaproxy() {
     return this.haproxy;
   }
+
+  public boolean isIgnoreConfigServers() {
+    return this.ignoreConfigServers;
+  }
+
+  public List<ServerInfo> getServers() {
+    return this.servers;
+  }
+
+  private static class ServerInfoConverter implements ValueConverter<ServerInfo> {
+
+    @Override
+    public ServerInfo convert(String s) {
+      String[] split = s.split(":", 2);
+      if (split.length < 2) {
+        throw new ValueConversionException("Invalid server format. Use <name>:<address>");
+      }
+      InetSocketAddress address;
+      try {
+        address = AddressUtil.parseAddress(split[1]);
+      } catch (IllegalStateException e) {
+        throw new ValueConversionException("Invalid hostname for server flag with name: " + split[0]);
+      }
+      return new ServerInfo(split[0], address);
+    }
+
+    @Override
+    public Class<? extends ServerInfo> valueType() {
+      return ServerInfo.class;
+    }
+
+    @Override
+    public String valuePattern() {
+      return "name>:<address";
+    }
+  }
 }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
index 2563f984f..cecc440c9 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
@@ -275,8 +275,14 @@ public class VelocityServer implements ProxyServer, ForwardingAudience {
 
     this.doStartupConfigLoad();
 
-    for (Map.Entry<String, String> entry : configuration.getServers().entrySet()) {
-      servers.register(new ServerInfo(entry.getKey(), AddressUtil.parseAddress(entry.getValue())));
+    for (ServerInfo cliServer : options.getServers()) {
+      servers.register(cliServer);
+    }
+
+    if (!options.isIgnoreConfigServers()) {
+      for (Map.Entry<String, String> entry : configuration.getServers().entrySet()) {
+        servers.register(new ServerInfo(entry.getKey(), AddressUtil.parseAddress(entry.getValue())));
+      }
     }
 
     ipAttemptLimiter = Ratelimiters.createWithMilliseconds(configuration.getLoginRatelimit());