Add option to send HAProxy header with RCON connections

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
 }