Add send_proxy_v2 option to send HAProxy header to server

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);