Skip to content

Commit

Permalink
support proxy protocol in more cases
Browse files Browse the repository at this point in the history
  • Loading branch information
erebe committed Aug 10, 2024
1 parent 8c4d091 commit dff2433
Show file tree
Hide file tree
Showing 5 changed files with 28 additions and 22 deletions.
25 changes: 14 additions & 11 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -520,10 +520,11 @@ fn parse_tunnel_arg(arg: &str) -> Result<LocalToRemote, io::Error> {
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),
Expand Down Expand Up @@ -557,9 +558,11 @@ fn parse_tunnel_arg(arg: &str) -> Result<LocalToRemote, io::Error> {
})
}
"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),
})
Expand Down Expand Up @@ -604,15 +607,15 @@ fn parse_reverse_tunnel_arg(arg: &str) -> Result<LocalToRemote, io::Error> {
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 { .. }
| LocalProtocol::ReverseHttpProxy { .. }
| 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),
Expand Down Expand Up @@ -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
{
Expand Down Expand Up @@ -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 { .. }
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
4 changes: 2 additions & 2 deletions src/restrictions/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 { .. }
Expand All @@ -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 { .. }
Expand Down
5 changes: 4 additions & 1 deletion src/tunnel/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@ pub enum LocalProtocol {
Udp {
timeout: Option<Duration>,
},
Stdio,
Stdio {
proxy_protocol: bool,
},
Socks5 {
timeout: Option<Duration>,
credentials: Option<(String, String)>,
Expand Down Expand Up @@ -53,6 +55,7 @@ pub enum LocalProtocol {
},
Unix {
path: PathBuf,
proxy_protocol: bool,
},
}

Expand Down
2 changes: 1 addition & 1 deletion src/tunnel/server/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 { .. }
Expand Down
14 changes: 7 additions & 7 deletions src/tunnel/transport/jwt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit dff2433

Please sign in to comment.