From dff243369ceb4cbd281e0a4d7d9ab72f16f17ecd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=A3rebe=20-=20Romain=20GERARD?= Date: Sat, 10 Aug 2024 10:58:05 +0200 Subject: [PATCH] support proxy protocol in more cases --- src/main.rs | 25 ++++++++++++++----------- src/restrictions/types.rs | 4 ++-- src/tunnel/mod.rs | 5 ++++- src/tunnel/server/server.rs | 2 +- src/tunnel/transport/jwt.rs | 14 +++++++------- 5 files changed, 28 insertions(+), 22 deletions(-) diff --git a/src/main.rs b/src/main.rs index 5d3a86a..bcec0a9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -520,10 +520,11 @@ fn parse_tunnel_arg(arg: &str) -> Result { format!("cannot parse unix socket path from {}", arg), )); }; - let (dest_host, dest_port, _options) = parse_tunnel_dest(remote)?; + let (dest_host, dest_port, options) = parse_tunnel_dest(remote)?; Ok(LocalToRemote { local_protocol: LocalProtocol::Unix { path: PathBuf::from(path), + proxy_protocol: get_proxy_protocol(&options), }, local: SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 0, 0, 0)), remote: (dest_host, dest_port), @@ -557,9 +558,11 @@ fn parse_tunnel_arg(arg: &str) -> Result { }) } "stdio" => { - let (dest_host, dest_port, _options) = parse_tunnel_dest(tunnel_info)?; + let (dest_host, dest_port, options) = parse_tunnel_dest(tunnel_info)?; Ok(LocalToRemote { - local_protocol: LocalProtocol::Stdio, + local_protocol: LocalProtocol::Stdio { + proxy_protocol: get_proxy_protocol(&options), + }, local: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from(0), 0)), remote: (dest_host, dest_port), }) @@ -604,7 +607,7 @@ fn parse_reverse_tunnel_arg(arg: &str) -> Result { credentials, proxy_protocol: _proxy_protocol, } => LocalProtocol::ReverseHttpProxy { timeout, credentials }, - LocalProtocol::Unix { path } => LocalProtocol::ReverseUnix { path }, + LocalProtocol::Unix { path, .. } => LocalProtocol::ReverseUnix { path }, LocalProtocol::ReverseTcp { .. } | LocalProtocol::ReverseUdp { .. } | LocalProtocol::ReverseSocks5 { .. } @@ -612,7 +615,7 @@ fn parse_reverse_tunnel_arg(arg: &str) -> Result { | LocalProtocol::ReverseUnix { .. } | LocalProtocol::TProxyTcp | LocalProtocol::TProxyUdp { .. } - | LocalProtocol::Stdio => { + | LocalProtocol::Stdio { .. } => { return Err(io::Error::new( ErrorKind::InvalidInput, format!("Cannot use {:?} as reverse tunnels {}", proto.local_protocol, arg), @@ -710,7 +713,7 @@ async fn main() -> anyhow::Result<()> { if args .local_to_remote .iter() - .filter(|x| x.local_protocol == LocalProtocol::Stdio) + .filter(|x| matches!(x.local_protocol, LocalProtocol::Stdio { .. })) .count() > 0 { @@ -937,7 +940,7 @@ async fn main() -> anyhow::Result<()> { LocalProtocol::ReverseUnix { .. } => { panic!("Unix socket is not available for non Unix platform") } - LocalProtocol::Stdio + LocalProtocol::Stdio { .. } | LocalProtocol::TProxyTcp | LocalProtocol::TProxyUdp { .. } | LocalProtocol::Tcp { .. } @@ -975,9 +978,9 @@ async fn main() -> anyhow::Result<()> { }); } #[cfg(unix)] - LocalProtocol::Unix { path } => { + LocalProtocol::Unix { path, proxy_protocol } => { use crate::tunnel::listeners::UnixTunnelListener; - let server = UnixTunnelListener::new(path, tunnel.remote.clone(), false).await?; // TODO: support proxy protocol + let server = UnixTunnelListener::new(path, tunnel.remote.clone(), *proxy_protocol).await?; tokio::spawn(async move { if let Err(err) = client.run_tunnel(server).await { error!("{:?}", err); @@ -1035,8 +1038,8 @@ async fn main() -> anyhow::Result<()> { }); } - LocalProtocol::Stdio => { - let (server, mut handle) = new_stdio_listener(tunnel.remote.clone(), false).await?; // TODO: support proxy protocol + LocalProtocol::Stdio { proxy_protocol } => { + let (server, mut handle) = new_stdio_listener(tunnel.remote.clone(), *proxy_protocol).await?; tokio::spawn(async move { if let Err(err) = client.run_tunnel(server).await { error!("{:?}", err); diff --git a/src/restrictions/types.rs b/src/restrictions/types.rs index f58897e..d000a46 100644 --- a/src/restrictions/types.rs +++ b/src/restrictions/types.rs @@ -157,7 +157,7 @@ impl From<&LocalProtocol> for ReverseTunnelConfigProtocol { match value { LocalProtocol::Tcp { .. } | LocalProtocol::Udp { .. } - | LocalProtocol::Stdio + | LocalProtocol::Stdio { .. } | LocalProtocol::Socks5 { .. } | LocalProtocol::TProxyTcp { .. } | LocalProtocol::TProxyUdp { .. } @@ -178,7 +178,7 @@ impl From<&LocalProtocol> for TunnelConfigProtocol { | LocalProtocol::ReverseUdp { .. } | LocalProtocol::ReverseSocks5 { .. } | LocalProtocol::ReverseUnix { .. } - | LocalProtocol::Stdio + | LocalProtocol::Stdio { .. } | LocalProtocol::Socks5 { .. } | LocalProtocol::TProxyTcp { .. } | LocalProtocol::TProxyUdp { .. } diff --git a/src/tunnel/mod.rs b/src/tunnel/mod.rs index 6987fe4..1aec9ec 100644 --- a/src/tunnel/mod.rs +++ b/src/tunnel/mod.rs @@ -22,7 +22,9 @@ pub enum LocalProtocol { Udp { timeout: Option, }, - Stdio, + Stdio { + proxy_protocol: bool, + }, Socks5 { timeout: Option, credentials: Option<(String, String)>, @@ -53,6 +55,7 @@ pub enum LocalProtocol { }, Unix { path: PathBuf, + proxy_protocol: bool, }, } diff --git a/src/tunnel/server/server.rs b/src/tunnel/server/server.rs index 538be2c..ee6f9a7 100644 --- a/src/tunnel/server/server.rs +++ b/src/tunnel/server/server.rs @@ -266,7 +266,7 @@ impl WsServer { error!("Received an unsupported target protocol {:?}", remote); Err(anyhow::anyhow!("Invalid upgrade request")) } - LocalProtocol::Stdio + LocalProtocol::Stdio { .. } | LocalProtocol::Socks5 { .. } | LocalProtocol::TProxyTcp | LocalProtocol::TProxyUdp { .. } diff --git a/src/tunnel/transport/jwt.rs b/src/tunnel/transport/jwt.rs index 74bc416..a09f216 100644 --- a/src/tunnel/transport/jwt.rs +++ b/src/tunnel/transport/jwt.rs @@ -33,17 +33,17 @@ impl JwtTunnelConfig { p: match dest.protocol { LocalProtocol::Tcp { .. } => dest.protocol.clone(), LocalProtocol::Udp { .. } => dest.protocol.clone(), - LocalProtocol::Stdio => LocalProtocol::Tcp { proxy_protocol: false }, - LocalProtocol::Socks5 { .. } => LocalProtocol::Tcp { proxy_protocol: false }, - LocalProtocol::HttpProxy { .. } => dest.protocol.clone(), - LocalProtocol::ReverseTcp => LocalProtocol::ReverseTcp, + LocalProtocol::ReverseTcp => dest.protocol.clone(), LocalProtocol::ReverseUdp { .. } => dest.protocol.clone(), LocalProtocol::ReverseSocks5 { .. } => dest.protocol.clone(), - LocalProtocol::TProxyTcp => LocalProtocol::Tcp { proxy_protocol: false }, - LocalProtocol::TProxyUdp { timeout } => LocalProtocol::Udp { timeout }, - LocalProtocol::Unix { .. } => LocalProtocol::Tcp { proxy_protocol: false }, LocalProtocol::ReverseUnix { .. } => dest.protocol.clone(), LocalProtocol::ReverseHttpProxy { .. } => dest.protocol.clone(), + LocalProtocol::TProxyTcp => unreachable!("cannot use tproxy tcp as destination protocol"), + LocalProtocol::TProxyUdp { .. } => unreachable!("cannot use tproxy udp as destination protocol"), + LocalProtocol::Stdio { .. } => unreachable!("cannot use stdio as destination protocol"), + LocalProtocol::Unix { .. } => unreachable!("canont use unix as destination protocol"), + LocalProtocol::Socks5 { .. } => unreachable!("cannot use socks5 as destination protocol"), + LocalProtocol::HttpProxy { .. } => unreachable!("cannot use http proxy as destination protocol"), }, r: dest.host.to_string(), rp: dest.port,