Add option to send HAProxy header with RCON connections

Cargo.lock

@@ -762,6 +762,7 @@ name = "lazymc"
 version = "0.2.2"
 dependencies = [
  "anyhow",
+ "async-std",
  "base64",
  "bytes",
  "chrono",
@@ -1295,8 +1296,7 @@ dependencies = [
 [[package]]
 name = "rcon"
 version = "0.5.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "465a6f903164a399084787547a026b83e7937bc576d8acdbd9e41ebf5de90a85"
+source = "git+https://github.com/timvisee/rust-rcon?rev=6fc6726#6fc67266e2e026aa52a2ce14bfb30ed69ec5b241"
 dependencies = [
  "async-std",
  "bytes",

Cargo.toml

@@ -24,7 +24,7 @@ default = ["rcon", "lobby"]
 # RCON support
 # Allow use of RCON to manage (stop) server.
 # Required on Windows.
-rcon = ["rust_rcon"]
+rcon = ["rust_rcon", "async-std"]
 
 # Lobby support
 # Add lobby join method, keeps client in fake lobby world until server is ready.
@@ -56,7 +56,8 @@ toml = "0.5"
 version-compare = "0.1"
 
 # Feature: rcon
-rust_rcon = { package = "rcon", version = "0.5", optional = true }
+rust_rcon = { package = "rcon", git = "https://github.com/timvisee/rust-rcon", rev = "6fc6726", optional = true }
+async-std = { version = "1.9.0", deafult-features = false, optional = true }
 
 # Feature: lobby
 named-binary-tag = { version = "0.6", optional = true }

res/lazymc.toml

@@ -157,6 +157,9 @@ command = "java -Xmx1G -Xms1G -jar server.jar --nogui"
 #password = ""
 #randomize_password = true
 
+# Add HAProxy v2 header to RCON connections.
+#send_proxy_v2 = false
+
 [advanced]
 # Automatically update values in Minecraft server.properties file as required.
 #rewrite_server_properties = true

src/config.rs

@@ -410,6 +410,9 @@ pub struct Rcon {
 
     /// Randomize server RCON password on each start.
     pub randomize_password: bool,
+
+    /// Add HAProxy v2 header to RCON connections.
+    pub send_proxy_v2: bool,
 }
 
 impl Default for Rcon {
@@ -419,6 +422,7 @@ impl Default for Rcon {
             port: 25575,
             password: "".into(),
             randomize_password: true,
+            send_proxy_v2: false,
         }
     }
 }

src/mc/rcon.rs

@@ -1,8 +1,13 @@
 use std::time::Duration;
 
+use async_std::net::TcpStream;
+use async_std::prelude::*;
 use rust_rcon::{Connection, Error as RconError};
 use tokio::time;
 
+use crate::config::Config;
+use crate::proxy;
+
 /// Minecraft RCON quirk.
 ///
 /// Wait this time between RCON operations.
@@ -17,16 +22,39 @@ pub struct Rcon {
 
 impl Rcon {
     /// Connect to a host.
-    pub async fn connect(addr: &str, pass: &str) -> Result<Self, Box<dyn std::error::Error>> {
+    pub async fn connect(
+        config: &Config,
+        addr: &str,
+        pass: &str,
+    ) -> Result<Self, Box<dyn std::error::Error>> {
+        // Connect to our TCP stream
+        let mut stream = TcpStream::connect(addr).await?;
+
+        // Add proxy header
+        if config.rcon.send_proxy_v2 {
+            trace!(target: "lazymc::rcon", "Sending local proxy header for RCON connection");
+            stream.write_all(&proxy::local_proxy_header()?).await?;
+        }
+
         // Start connection
         let con = Connection::builder()
             .enable_minecraft_quirks(true)
-            .connect(addr, pass)
+            .connect_stream(stream, pass)
             .await?;
 
         Ok(Self { con })
     }
 
+    /// Connect to a host from the given configuration.
+    pub async fn connect_config(config: &Config) -> Result<Self, Box<dyn std::error::Error>> {
+        // RCON address
+        let mut addr = config.server.address;
+        addr.set_port(config.rcon.port);
+        let addr = addr.to_string();
+
+        Self::connect(config, &addr, &config.rcon.password).await
+    }
+
     /// Send command over RCON.
     pub async fn cmd(&mut self, cmd: &str) -> Result<String, RconError> {
         // Minecraft quirk

src/server.rs

@@ -498,13 +498,8 @@ async fn stop_server_rcon(config: &Config, server: &Server) -> bool {
         return false;
     }
 
-    // RCON address
-    let mut addr = config.server.address;
-    addr.set_port(config.rcon.port);
-    let addr = addr.to_string();
-
     // Create RCON client
-    let mut rcon = match Rcon::connect(&addr, &config.rcon.password).await {
+    let mut rcon = match Rcon::connect_config(&config).await {
         Ok(rcon) => rcon,
         Err(err) => {
             error!(target: "lazymc", "Failed to RCON server to sleep: {}", err);
@@ -522,11 +517,11 @@ async fn stop_server_rcon(config: &Config, server: &Server) -> bool {
     server.rcon_last_stop.lock().await.replace(Instant::now());
     server.update_state(State::Stopping, config).await;
 
-    drop(rcon_lock);
-
     // Gracefully close connection
     rcon.close().await;
 
+    drop(rcon_lock);
+
     true
 }