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, pub nodelay: Option, + pub enable_proxy_protocol: Option, } 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 { + 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::( 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( bind_addr: String, + enable_proxy_protocol: bool, mut data_ch_rx: mpsc::Receiver, data_ch_req_tx: mpsc::UnboundedSender, shutdown_rx: broadcast::Receiver, @@ -637,6 +643,11 @@ async fn run_tcp_connection_pool( 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;