Add send_proxy_v2 option to send HAProxy header to server

Cargo.lock

@@ -391,6 +391,12 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "doc-comment"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10"
+
 [[package]]
 name = "dotenv"
 version = "0.15.0"
@@ -771,6 +777,7 @@ dependencies = [
  "named-binary-tag",
  "notify",
  "pretty_env_logger",
+ "proxy-protocol",
  "quartz_nbt",
  "rand 0.8.4",
  "rcon",
@@ -1090,6 +1097,16 @@ dependencies = [
  "unicode-xid",
 ]
 
+[[package]]
+name = "proxy-protocol"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0e50c72c21c738f5c5f350cc33640aee30bf7cd20f9d9da20ed41bce2671d532"
+dependencies = [
+ "bytes",
+ "snafu",
+]
+
 [[package]]
 name = "quartz_nbt"
 version = "0.2.4"
@@ -1394,6 +1411,27 @@ version = "0.4.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9def91fd1e018fe007022791f865d0ccc9b3a0d5001e01aabb8b40e46000afb5"
 
+[[package]]
+name = "snafu"
+version = "0.6.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eab12d3c261b2308b0d80c26fffb58d17eba81a4be97890101f416b478c79ca7"
+dependencies = [
+ "doc-comment",
+ "snafu-derive",
+]
+
+[[package]]
+name = "snafu-derive"
+version = "0.6.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1508efa03c362e23817f96cde18abed596a25219a8b2c66e8db33c03543d315b"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "socket2"
 version = "0.4.2"

Cargo.toml

@@ -45,6 +45,7 @@ log = "0.4"
 minecraft-protocol = { git = "https://github.com/timvisee/rust-minecraft-protocol", rev = "356ea54" }
 notify = "4.0"
 pretty_env_logger = "0.4"
+proxy-protocol = "0.5"
 rand = "0.8"
 serde = "1.0"
 serde_json = "1.0"

TODO.md

@@ -1,5 +1,6 @@
 # TODO
 
+- Describe `send_proxy_v2` usage on docs page
 - Better organize code
 - Resolve TODOs in code
 - Don't drop errors, handle everywhere where needed (some were dropped while

res/lazymc.toml

@@ -51,6 +51,9 @@ command = "java -Xmx1G -Xms1G -jar server.jar --nogui"
 # On connect, clients show a 'Disconnected' message rather than the ban reason.
 #drop_banned_ips = false
 
+# Add HAProxy v2 header to proxied connections.
+#send_proxy_v2 = false
+
 [time]
 # Sleep after number of seconds.
 #sleep_after = 60
@@ -105,6 +108,9 @@ command = "java -Xmx1G -Xms1G -jar server.jar --nogui"
 # The target server will receive original client handshake and login request as received by lazymc.
 #address = "127.0.0.1:25565"
 
+# Add HAProxy v2 header to proxied connections.
+#send_proxy_v2 = false
+
 [join.lobby]
 # Lobby occupation method.
 # The client joins a fake lobby server with an empty world, floating in space.

src/config.rs

@@ -188,6 +188,10 @@ pub struct Server {
     /// Drop connections from banned IPs.
     #[serde(default)]
     pub drop_banned_ips: bool,
+
+    /// Add HAProxy v2 header to proxied connections.
+    #[serde(default)]
+    pub send_proxy_v2: bool,
 }
 
 /// Time configuration.
@@ -333,12 +337,17 @@ pub struct JoinForward {
     /// IP and port to forward to.
     #[serde(deserialize_with = "to_socket_addrs")]
     pub address: SocketAddr,
+
+    /// Add HAProxy v2 header to proxied connections.
+    #[serde(default)]
+    pub send_proxy_v2: bool,
 }
 
 impl Default for JoinForward {
     fn default() -> Self {
         Self {
             address: "127.0.0.1:25565".parse().unwrap(),
+            send_proxy_v2: false,
         }
     }
 }

src/join/forward.rs

@@ -4,6 +4,7 @@ use bytes::BytesMut;
 use tokio::net::TcpStream;
 
 use crate::config::*;
+use crate::proxy::ProxyHeader;
 use crate::service;
 
 use super::MethodResult;
@@ -20,6 +21,7 @@ pub async fn occupy(
 
     service::server::route_proxy_address_queue(
         inbound,
+        ProxyHeader::Proxy.not_none(config.join.forward.send_proxy_v2),
         config.join.forward.address,
         inbound_history.clone(),
     );

src/lobby.rs

@@ -15,6 +15,7 @@ use minecraft_protocol::version::v1_17_1::game::{
     PlayerPositionAndLook, Respawn, SetTitleSubtitle, SetTitleText, SetTitleTimes, TimeUpdate,
 };
 use nbt::CompoundTag;
+use tokio::io::AsyncWriteExt;
 use tokio::net::tcp::{ReadHalf, WriteHalf};
 use tokio::net::TcpStream;
 use tokio::select;
@@ -128,7 +129,10 @@ pub async fn serve(
             // Wait for server to come online, then set up new connection to it
             stage_wait(client, &server, &config, &mut writer).await?;
             let (server_client, mut outbound, mut server_buf) =
-                connect_to_server(client_info, &config).await?;
+                connect_to_server(client_info, &inbound, &config).await?;
+            let (returned_reader, returned_writer) = inbound.split();
+            reader = returned_reader;
+            writer = returned_writer;
 
             // Grab join game packet from server
             let join_game =
@@ -562,11 +566,12 @@ async fn wait_for_server(server: &Server, config: &Config) -> Result<(), ()> {
 /// This will initialize the connection to the play state. Client details are used.
 async fn connect_to_server(
     client_info: ClientInfo,
+    inbound: &TcpStream,
     config: &Config,
 ) -> Result<(Client, TcpStream, BytesMut), ()> {
     time::timeout(
         SERVER_CONNECT_TIMEOUT,
-        connect_to_server_no_timeout(client_info, config),
+        connect_to_server_no_timeout(client_info, inbound, config),
     )
     .await
     .map_err(|_| {
@@ -580,6 +585,7 @@ async fn connect_to_server(
 // TODO: clean this up
 async fn connect_to_server_no_timeout(
     client_info: ClientInfo,
+    inbound: &TcpStream,
     config: &Config,
 ) -> Result<(Client, TcpStream, BytesMut), ()> {
     // Open connection
@@ -588,6 +594,15 @@ async fn connect_to_server_no_timeout(
         .await
         .map_err(|_| ())?;
 
+    // Add proxy header
+    if config.server.send_proxy_v2 {
+        trace!(target: "lazymc::lobby", "Sending local proxy header for server connection");
+        outbound
+            .write_all(&proxy::stream_proxy_header(inbound).map_err(|_| ())?)
+            .await
+            .map_err(|_| ())?;
+    }
+
     // Construct temporary server client
     let tmp_client = match outbound.local_addr() {
         Ok(addr) => Client::new(addr),

src/proxy.rs

@@ -1,6 +1,9 @@
 use std::error::Error;
 use std::net::SocketAddr;
 
+use bytes::BytesMut;
+use proxy_protocol::version2::{ProxyAddresses, ProxyCommand, ProxyTransportProtocol};
+use proxy_protocol::EncodeError;
 use tokio::io;
 use tokio::io::AsyncWriteExt;
 use tokio::net::TcpStream;
@@ -8,8 +11,12 @@ use tokio::net::TcpStream;
 use crate::net;
 
 /// Proxy the inbound stream to a target address.
-pub async fn proxy(inbound: TcpStream, addr_target: SocketAddr) -> Result<(), Box<dyn Error>> {
-    proxy_with_queue(inbound, addr_target, &[]).await
+pub async fn proxy(
+    inbound: TcpStream,
+    proxy_header: ProxyHeader,
+    addr_target: SocketAddr,
+) -> Result<(), Box<dyn Error>> {
+    proxy_with_queue(inbound, proxy_header, addr_target, &[]).await
 }
 
 /// Proxy the inbound stream to a target address.
@@ -17,12 +24,26 @@ pub async fn proxy(inbound: TcpStream, addr_target: SocketAddr) -> Result<(), Bo
 /// Send the queue to the target server before proxying.
 pub async fn proxy_with_queue(
     inbound: TcpStream,
+    proxy_header: ProxyHeader,
     addr_target: SocketAddr,
     queue: &[u8],
 ) -> Result<(), Box<dyn Error>> {
     // Set up connection to server
     // TODO: on connect fail, ping server and redirect to serve_status if offline
-    let outbound = TcpStream::connect(addr_target).await?;
+    let mut outbound = TcpStream::connect(addr_target).await?;
+
+    // Add proxy header
+    match proxy_header {
+        ProxyHeader::None => {}
+        ProxyHeader::Local => {
+            let header = local_proxy_header()?;
+            outbound.write_all(&header).await?;
+        }
+        ProxyHeader::Proxy => {
+            let header = stream_proxy_header(&inbound)?;
+            outbound.write_all(&header).await?;
+        }
+    }
 
     // Start proxy on both streams
     proxy_inbound_outbound_with_queue(inbound, outbound, &[], queue).await
@@ -71,3 +92,75 @@ pub async fn proxy_inbound_outbound_with_queue(
 
     Ok(())
 }
+
+/// Proxy header.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum ProxyHeader {
+    /// Do not add proxy header.
+    None,
+
+    /// Header for locally initiated connection.
+    Local,
+
+    /// Header for proxied connection.
+    Proxy,
+}
+
+impl ProxyHeader {
+    /// Changes to `None` if `false` if given.
+    ///
+    /// `None` stays `None`.
+    pub fn not_none(self, not_none: bool) -> Self {
+        if not_none {
+            self
+        } else {
+            Self::None
+        }
+    }
+}
+
+/// Get the proxy header for a locally initiated connection.
+///
+/// This header may be sent over the outbound stream to signal client information.
+pub fn local_proxy_header() -> Result<BytesMut, EncodeError> {
+    // Build proxy header
+    let header = proxy_protocol::ProxyHeader::Version2 {
+        command: ProxyCommand::Local,
+        transport_protocol: ProxyTransportProtocol::Stream,
+        addresses: ProxyAddresses::Unspec,
+    };
+
+    proxy_protocol::encode(header)
+}
+
+/// Get the proxy header for the given inbound stream.
+///
+/// This header may be sent over the outbound stream to signal client information.
+pub fn stream_proxy_header(inbound: &TcpStream) -> Result<BytesMut, EncodeError> {
+    // Get peer and local address
+    let peer = inbound
+        .peer_addr()
+        .expect("Peer address not known for TCP stream");
+    let local = inbound
+        .local_addr()
+        .expect("Local address not known for TCP stream");
+
+    // Build proxy header
+    let header = proxy_protocol::ProxyHeader::Version2 {
+        command: ProxyCommand::Proxy,
+        transport_protocol: ProxyTransportProtocol::Stream,
+        addresses: match (peer, local) {
+            (SocketAddr::V4(source), SocketAddr::V4(destination)) => ProxyAddresses::Ipv4 {
+                source,
+                destination,
+            },
+            (SocketAddr::V6(source), SocketAddr::V6(destination)) => ProxyAddresses::Ipv6 {
+                source,
+                destination,
+            },
+            (_, _) => unreachable!(),
+        },
+    };
+
+    proxy_protocol::encode(header)
+}

src/service/server.rs

@@ -7,7 +7,7 @@ use tokio::net::{TcpListener, TcpStream};
 
 use crate::config::Config;
 use crate::proto::client::Client;
-use crate::proxy;
+use crate::proxy::{self, ProxyHeader};
 use crate::server::{self, Server};
 use crate::service;
 use crate::status;
@@ -114,7 +114,12 @@ fn route_status(inbound: TcpStream, config: Arc<Config>, server: Arc<Server>, pe
 #[inline]
 fn route_proxy(inbound: TcpStream, config: Arc<Config>) {
     // When server is online, proxy all
-    let service = proxy::proxy(inbound, config.server.address).map(|r| {
+    let service = proxy::proxy(
+        inbound,
+        ProxyHeader::Proxy.not_none(config.server.send_proxy_v2),
+        config.server.address,
+    )
+    .map(|r| {
         if let Err(err) = r {
             warn!(target: "lazymc", "Failed to proxy: {}", err);
         }
@@ -126,15 +131,25 @@ fn route_proxy(inbound: TcpStream, config: Arc<Config>) {
 /// Route inbound TCP stream to proxy with queued data, spawning a new task.
 #[inline]
 pub fn route_proxy_queue(inbound: TcpStream, config: Arc<Config>, queue: BytesMut) {
-    route_proxy_address_queue(inbound, config.server.address, queue);
+    route_proxy_address_queue(
+        inbound,
+        ProxyHeader::Proxy.not_none(config.server.send_proxy_v2),
+        config.server.address,
+        queue,
+    );
 }
 
 /// Route inbound TCP stream to proxy with given address and queued data, spawning a new task.
 #[inline]
-pub fn route_proxy_address_queue(inbound: TcpStream, addr: SocketAddr, queue: BytesMut) {
+pub fn route_proxy_address_queue(
+    inbound: TcpStream,
+    proxy_header: ProxyHeader,
+    addr: SocketAddr,
+    queue: BytesMut,
+) {
     // When server is online, proxy all
     let service = async move {
-        proxy::proxy_with_queue(inbound, addr, &queue)
+        proxy::proxy_with_queue(inbound, proxy_header, addr, &queue)
             .map(|r| {
                 if let Err(err) = r {
                     warn!(target: "lazymc", "Failed to proxy: {}", err);