From c013a20e52f092fd87c47f54015fa282567ef6ad Mon Sep 17 00:00:00 2001
From: Shashi Shekhar <shshekhar93@gmail.com>
Date: Wed, 10 Apr 2024 23:27:20 +0530
Subject: [PATCH] Add proxy protocol support

---
 examples/proxy_protocol/client.toml | 11 +++++++++++
 examples/proxy_protocol/server.toml | 12 ++++++++++++
 src/config.rs                       |  1 +
 src/helper.rs                       | 15 +++++++++++++++
 src/server.rs                       | 13 ++++++++++++-
 5 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 examples/proxy_protocol/client.toml
 create mode 100644 examples/proxy_protocol/server.toml

diff --git a/examples/proxy_protocol/client.toml b/examples/proxy_protocol/client.toml
new file mode 100644
index 00000000..43dc90f0
--- /dev/null
+++ b/examples/proxy_protocol/client.toml
@@ -0,0 +1,11 @@
+# rathole configuration for proxy protocol enabled client
+# 
+# The client configuration is essentially unaffected, since the proxy
+# protocol header would be transarently passed to the downstream server.
+
+[client]
+remote_addr = "localhost:2333"
+default_token = "123"
+
+[client.services.foo1]
+local_addr = "127.0.0.1:80"
diff --git a/examples/proxy_protocol/server.toml b/examples/proxy_protocol/server.toml
new file mode 100644
index 00000000..ad801a01
--- /dev/null
+++ b/examples/proxy_protocol/server.toml
@@ -0,0 +1,12 @@
+# rathole configuration for proxy protocol enabled client
+# 
+# The service configuration has an additional `enable_proxy_protocol` boolean field.
+# Not setting this field defaults its value to `false` at runtime.
+
+[server]
+bind_addr = "0.0.0.0:2333"
+default_token = "123"
+
+[server.services.foo1]
+bind_addr = "0.0.0.0:5202"
+enable_proxy_protocol = true
diff --git a/src/config.rs b/src/config.rs
index ca85fc20..9ea550b7 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -102,6 +102,7 @@ pub struct ServerServiceConfig {
     pub bind_addr: String,
     pub token: Option<MaskedString>,
     pub nodelay: Option<bool>,
+    pub enable_proxy_protocol: Option<bool>,
 }
 
 impl ServerServiceConfig {
diff --git a/src/helper.rs b/src/helper.rs
index a292969f..56fe4881 100644
--- a/src/helper.rs
+++ b/src/helper.rs
@@ -164,3 +164,18 @@ where
     conn.flush().await.with_context(|| "Failed to flush data")?;
     Ok(())
 }
+
+pub fn generate_proxy_protocol_v1_header(s: &TcpStream) -> Result<String> {
+    let local_addr = s.local_addr()?;
+    let remote_addr = s.peer_addr()?;
+    let proto = if local_addr.is_ipv4() { "TCP4" } else { "TCP6" };
+    let header = format!(
+        "PROXY {} {} {} {} {}\r\n", 
+        proto, 
+        remote_addr.ip(), 
+        local_addr.ip(), 
+        remote_addr.port(), 
+        local_addr.port()
+    );
+    Ok(header)
+}
diff --git a/src/server.rs b/src/server.rs
index a4c49482..572f7eda 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,7 +1,7 @@
 use crate::config::{Config, ServerConfig, ServerServiceConfig, ServiceType, TransportType};
 use crate::config_watcher::{ConfigChange, ServerServiceChange};
 use crate::constants::{listen_backoff, UDP_BUFFER_SIZE};
-use crate::helper::{retry_notify_with_deadline, write_and_flush};
+use crate::helper::{generate_proxy_protocol_v1_header, retry_notify_with_deadline, write_and_flush};
 use crate::multi_map::MultiMap;
 use crate::protocol::Hello::{ControlChannelHello, DataChannelHello};
 use crate::protocol::{
@@ -427,11 +427,16 @@ where
 
         let shutdown_rx_clone = shutdown_tx.subscribe();
         let bind_addr = service.bind_addr.clone();
+        let enable_proxy_protocol = service.enable_proxy_protocol.unwrap_or_default();
+        if enable_proxy_protocol {
+            debug!("Proxy protocol is enabled");
+        }
         match service.service_type {
             ServiceType::Tcp => tokio::spawn(
                 async move {
                     if let Err(e) = run_tcp_connection_pool::<T>(
                         bind_addr,
+                        enable_proxy_protocol,
                         data_ch_rx,
                         data_ch_req_tx,
                         shutdown_rx_clone,
@@ -625,6 +630,7 @@ fn tcp_listen_and_send(
 #[instrument(skip_all)]
 async fn run_tcp_connection_pool<T: Transport>(
     bind_addr: String,
+    enable_proxy_protocol: bool,
     mut data_ch_rx: mpsc::Receiver<T::Stream>,
     data_ch_req_tx: mpsc::UnboundedSender<bool>,
     shutdown_rx: broadcast::Receiver<bool>,
@@ -637,6 +643,11 @@ async fn run_tcp_connection_pool<T: Transport>(
             if let Some(mut ch) = data_ch_rx.recv().await {
                 if write_and_flush(&mut ch, &cmd).await.is_ok() {
                     tokio::spawn(async move {
+                        if enable_proxy_protocol {
+                            let proxy_proto_header = generate_proxy_protocol_v1_header(&visitor).unwrap();
+                            let _ = ch.write_all(&proxy_proto_header.into_bytes()).await;
+                            let _ = ch.flush().await;
+                        }
                         let _ = copy_bidirectional(&mut ch, &mut visitor).await;
                     });
                     break;