From 363e3a7f94b1a8693893f4fe0349fb181b127bbb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=86=B0=E7=A0=9A=E7=82=BD?= <Bing_Yanchi@outlook.com>
Date: Sun, 21 Jan 2024 23:47:40 +0800
Subject: [PATCH] Add -haproxy-protocol startup parameter (#1136)

---
 .../java/com/velocitypowered/proxy/ProxyOptions.java  | 11 +++++++++++
 .../com/velocitypowered/proxy/VelocityServer.java     |  6 ++++++
 .../proxy/config/VelocityConfiguration.java           |  8 ++++++++
 3 files changed, 25 insertions(+)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java b/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java
index 30c29b6b9..3ab419ff4 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/ProxyOptions.java
@@ -34,6 +34,7 @@ public final class ProxyOptions {
   private static final Logger logger = LogManager.getLogger(ProxyOptions.class);
   private final boolean help;
   private final @Nullable Integer port;
+  private final @Nullable Boolean haproxy;
 
   ProxyOptions(final String[] args) {
     final OptionParser parser = new OptionParser();
@@ -43,10 +44,16 @@ public final class ProxyOptions {
     final OptionSpec<Integer> port = parser.acceptsAll(Arrays.asList("p", "port"),
             "Specify the bind port to be used. The configuration bind port will be ignored.")
         .withRequiredArg().ofType(Integer.class);
+    final OptionSpec<Boolean> haproxy = parser.acceptsAll(
+            Arrays.asList("haproxy", "haproxy-protocol"),
+            "Choose whether to enable haproxy protocol. "
+                    + "The configuration haproxy protocol will be ignored.")
+        .withRequiredArg().ofType(Boolean.class).defaultsTo(null);
     final OptionSet set = parser.parse(args);
 
     this.help = set.has(help);
     this.port = port.value(set);
+    this.haproxy = haproxy.value(set);
 
     if (this.help) {
       try {
@@ -64,4 +71,8 @@ public final class ProxyOptions {
   public @Nullable Integer getPort() {
     return this.port;
   }
+
+  public @Nullable Boolean isHaproxy() {
+    return this.haproxy;
+  }
 }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
index 8ccdbc508..c614d415e 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
@@ -257,6 +257,12 @@ public class VelocityServer implements ProxyServer, ForwardingAudience {
       this.cm.bind(configuration.getBind());
     }
 
+    final Boolean haproxy = this.options.isHaproxy();
+    if (haproxy != null) {
+      logger.debug("Overriding HAProxy protocol to {} from command line option", haproxy);
+      configuration.setProxyProtocol(haproxy);
+    }
+
     if (configuration.isQueryEnabled()) {
       this.cm.queryBind(configuration.getBind().getHostString(), configuration.getQueryPort());
     }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/config/VelocityConfiguration.java b/proxy/src/main/java/com/velocitypowered/proxy/config/VelocityConfiguration.java
index e55c0e9ae..f4a76919e 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/config/VelocityConfiguration.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/config/VelocityConfiguration.java
@@ -354,6 +354,10 @@ public class VelocityConfiguration implements ProxyConfig {
     return advanced.isProxyProtocol();
   }
 
+  public void setProxyProtocol(boolean proxyProtocol) {
+    advanced.setProxyProtocol(proxyProtocol);
+  }
+
   public boolean useTcpFastOpen() {
     return advanced.isTcpFastOpen();
   }
@@ -755,6 +759,10 @@ public class VelocityConfiguration implements ProxyConfig {
       return proxyProtocol;
     }
 
+    public void setProxyProtocol(boolean proxyProtocol) {
+      this.proxyProtocol = proxyProtocol;
+    }
+
     public boolean isTcpFastOpen() {
       return tcpFastOpen;
     }