cleanup

src/main.rs

@@ -12,14 +12,14 @@ use crate::tunnel::listeners::{
     new_stdio_listener, new_udp_listener, HttpProxyTunnelListener, Socks5TunnelListener, TcpTunnelListener,
 };
 use crate::tunnel::server::{TlsServerConfig, WsServer, WsServerConfig};
-use crate::tunnel::{to_host_port, RemoteAddr, TransportAddr, TransportScheme};
+use crate::tunnel::{to_host_port, LocalProtocol, RemoteAddr, TransportAddr, TransportScheme};
+use anyhow::{anyhow, Context};
 use base64::Engine;
 use clap::Parser;
 use hyper::header::HOST;
 use hyper::http::{HeaderName, HeaderValue};
 use log::debug;
 use parking_lot::{Mutex, RwLock};
-use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::fmt::Debug;
 use std::io;
@@ -376,66 +376,11 @@ struct Server {
     http_proxy_password: Option<String>,
 }
 
-#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
-enum LocalProtocol {
-    Tcp {
-        proxy_protocol: bool,
-    },
-    Udp {
-        timeout: Option<Duration>,
-    },
-    Stdio,
-    Socks5 {
-        timeout: Option<Duration>,
-        credentials: Option<(String, String)>,
-    },
-    TProxyTcp,
-    TProxyUdp {
-        timeout: Option<Duration>,
-    },
-    HttpProxy {
-        timeout: Option<Duration>,
-        credentials: Option<(String, String)>,
-        proxy_protocol: bool,
-    },
-    ReverseTcp,
-    ReverseUdp {
-        timeout: Option<Duration>,
-    },
-    ReverseSocks5 {
-        timeout: Option<Duration>,
-        credentials: Option<(String, String)>,
-    },
-    ReverseHttpProxy {
-        timeout: Option<Duration>,
-        credentials: Option<(String, String)>,
-    },
-    ReverseUnix {
-        path: PathBuf,
-    },
-    Unix {
-        path: PathBuf,
-    },
-}
-
-impl LocalProtocol {
-    pub const fn is_reverse_tunnel(&self) -> bool {
-        matches!(
-            self,
-            Self::ReverseTcp
-                | Self::ReverseUdp { .. }
-                | Self::ReverseSocks5 { .. }
-                | Self::ReverseUnix { .. }
-                | Self::ReverseHttpProxy { .. }
-        )
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct LocalToRemote {
     local_protocol: LocalProtocol,
     local: SocketAddr,
-    remote: (Host<String>, u16),
+    remote: (Host, u16),
 }
 
 fn parse_duration_sec(arg: &str) -> Result<Duration, io::Error> {
@@ -773,24 +718,7 @@ async fn main() -> anyhow::Result<()> {
                 TransportScheme::from_str(args.remote_addr.scheme()).expect("invalid scheme in server url");
             let tls = match transport_scheme {
                 TransportScheme::Ws | TransportScheme::Http => None,
-                TransportScheme::Wss => Some(TlsClientConfig {
-                    tls_connector: Arc::new(RwLock::new(
-                        tls::tls_connector(
-                            args.tls_verify_certificate,
-                            transport_scheme.alpn_protocols(),
-                            !args.tls_sni_disable,
-                            tls_certificate,
-                            tls_key,
-                        )
-                        .expect("Cannot create tls connector"),
-                    )),
-                    tls_sni_override: args.tls_sni_override,
-                    tls_verify_certificate: args.tls_verify_certificate,
-                    tls_sni_disabled: args.tls_sni_disable,
-                    tls_certificate_path: args.tls_certificate.clone(),
-                    tls_key_path: args.tls_private_key.clone(),
-                }),
-                TransportScheme::Https => Some(TlsClientConfig {
+                TransportScheme::Wss | TransportScheme::Https => Some(TlsClientConfig {
                     tls_connector: Arc::new(RwLock::new(
                         tls::tls_connector(
                             args.tls_verify_certificate,
@@ -824,25 +752,8 @@ async fn main() -> anyhow::Result<()> {
                     panic!("http headers file does not exists: {}", path.display());
                 }
             }
-            let http_proxy = if let Some(proxy) = args.http_proxy {
-                let mut proxy = if proxy.starts_with("http://") {
-                    Url::parse(&proxy).expect("Invalid http proxy url")
-                } else {
-                    Url::parse(&format!("http://{}", proxy)).expect("Invalid http proxy url")
-                };
 
-                if let Some(login) = args.http_proxy_login {
-                    proxy.set_username(login.as_str()).expect("Cannot set http proxy login");
-                }
-                if let Some(password) = args.http_proxy_password {
-                    proxy
-                        .set_password(Some(password.as_str()))
-                        .expect("Cannot set http proxy password");
-                }
-                Some(proxy)
-            } else {
-                None
-            };
+            let http_proxy = mk_http_proxy(args.http_proxy, args.http_proxy_login, args.http_proxy_password)?;
             let client_config = WsClientConfig {
                 remote_addr: TransportAddr::new(
                     TransportScheme::from_str(args.remote_addr.scheme()).unwrap(),
@@ -1176,26 +1087,7 @@ async fn main() -> anyhow::Result<()> {
                 restriction_cfg
             };
 
-            let http_proxy = if let Some(proxy) = args.http_proxy {
-                let mut proxy = if proxy.starts_with("http://") {
-                    Url::parse(&proxy).expect("Invalid http proxy url")
-                } else {
-                    Url::parse(&format!("http://{}", proxy)).expect("Invalid http proxy url")
-                };
-
-                if let Some(login) = args.http_proxy_login {
-                    proxy.set_username(login.as_str()).expect("Cannot set http proxy login");
-                }
-                if let Some(password) = args.http_proxy_password {
-                    proxy
-                        .set_password(Some(password.as_str()))
-                        .expect("Cannot set http proxy password");
-                }
-                Some(proxy)
-            } else {
-                None
-            };
-
+            let http_proxy = mk_http_proxy(args.http_proxy, args.http_proxy_login, args.http_proxy_password)?;
             let server_config = WsServerConfig {
                 socket_so_mark: args.socket_so_mark,
                 bind: args.remote_addr.socket_addrs(|| Some(8080)).unwrap()[0],
@@ -1230,3 +1122,33 @@ async fn main() -> anyhow::Result<()> {
     tokio::signal::ctrl_c().await.unwrap();
     Ok(())
 }
+
+fn mk_http_proxy(
+    http_proxy: Option<String>,
+    proxy_login: Option<String>,
+    proxy_password: Option<String>,
+) -> anyhow::Result<Option<Url>> {
+    let Some(proxy) = http_proxy else {
+        return Ok(None);
+    };
+
+    let mut proxy = if proxy.starts_with("http://") {
+        Url::parse(&proxy).with_context(|| "Invalid http proxy url")?
+    } else {
+        Url::parse(&format!("http://{}", proxy)).with_context(|| "Invalid http proxy url")?
+    };
+
+    if let Some(login) = proxy_login {
+        proxy
+            .set_username(login.as_str())
+            .map_err(|_| anyhow!("Cannot set http proxy login"))?;
+    }
+
+    if let Some(password) = proxy_password {
+        proxy
+            .set_password(Some(password.as_str()))
+            .map_err(|_| anyhow!("Cannot set http proxy password"))?;
+    }
+
+    Ok(Some(proxy))
+}

src/protocols/socks5/tcp_server.rs

@@ -1,5 +1,5 @@
 use super::udp_server::{Socks5UdpStream, Socks5UdpStreamWriter};
-use crate::LocalProtocol;
+use crate::tunnel::LocalProtocol;
 use anyhow::Context;
 use fast_socks5::server::{Config, DenyAuthentication, SimpleUserPassword, Socks5Server};
 use fast_socks5::util::target_addr::TargetAddr;

src/restrictions/types.rs

@@ -1,4 +1,4 @@
-use crate::LocalProtocol;
+use crate::tunnel::LocalProtocol;
 use ipnet::{IpNet, Ipv4Net, Ipv6Net};
 use regex::Regex;
 use serde::{Deserialize, Deserializer};

src/tunnel/client/cnx_pool.rs

@@ -1,7 +1,7 @@
 use crate::protocols;
 use crate::protocols::tls;
+use crate::tunnel::client::l4_transport_stream::TransportStream;
 use crate::tunnel::client::WsClientConfig;
-use crate::tunnel::TransportStream;
 use async_trait::async_trait;
 use bb8::ManageConnection;
 use std::ops::Deref;

src/tunnel/client/l4_transport_stream.rs

@@ -0,0 +1,61 @@
+use std::io::{Error, IoSlice};
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio::net::TcpStream;
+use tokio_rustls::client::TlsStream;
+
+pub enum TransportStream {
+    Plain(TcpStream),
+    Tls(TlsStream<TcpStream>),
+}
+
+impl AsyncRead for TransportStream {
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
+        match self.get_mut() {
+            Self::Plain(cnx) => Pin::new(cnx).poll_read(cx, buf),
+            Self::Tls(cnx) => Pin::new(cnx).poll_read(cx, buf),
+        }
+    }
+}
+
+impl AsyncWrite for TransportStream {
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Error>> {
+        match self.get_mut() {
+            Self::Plain(cnx) => Pin::new(cnx).poll_write(cx, buf),
+            Self::Tls(cnx) => Pin::new(cnx).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        match self.get_mut() {
+            Self::Plain(cnx) => Pin::new(cnx).poll_flush(cx),
+            Self::Tls(cnx) => Pin::new(cnx).poll_flush(cx),
+        }
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        match self.get_mut() {
+            Self::Plain(cnx) => Pin::new(cnx).poll_shutdown(cx),
+            Self::Tls(cnx) => Pin::new(cnx).poll_shutdown(cx),
+        }
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>],
+    ) -> Poll<Result<usize, Error>> {
+        match self.get_mut() {
+            Self::Plain(cnx) => Pin::new(cnx).poll_write_vectored(cx, bufs),
+            Self::Tls(cnx) => Pin::new(cnx).poll_write_vectored(cx, bufs),
+        }
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        match &self {
+            Self::Plain(cnx) => cnx.is_write_vectored(),
+            Self::Tls(cnx) => cnx.is_write_vectored(),
+        }
+    }
+}

src/tunnel/client/mod.rs

@@ -2,6 +2,7 @@
 mod client;
 mod cnx_pool;
 mod config;
+pub mod l4_transport_stream;
 
 pub use client::WsClient;
 pub use config::TlsClientConfig;

src/tunnel/connectors/sock5.rs

@@ -8,12 +8,12 @@ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
 use url::Url;
 
+use crate::protocols;
 use crate::protocols::dns::DnsResolver;
 use crate::protocols::udp;
 use crate::protocols::udp::WsUdpSocket;
 use crate::tunnel::connectors::TunnelConnector;
-use crate::tunnel::RemoteAddr;
-use crate::{protocols, LocalProtocol};
+use crate::tunnel::{LocalProtocol, RemoteAddr};
 
 pub struct Socks5TunnelConnector<'a> {
     so_mark: Option<u32>,

src/tunnel/listeners/http_proxy.rs

@@ -1,7 +1,6 @@
 use crate::protocols::http_proxy;
 use crate::protocols::http_proxy::HttpProxyListener;
-use crate::tunnel::RemoteAddr;
-use crate::LocalProtocol;
+use crate::tunnel::{LocalProtocol, RemoteAddr};
 use anyhow::{anyhow, Context};
 use std::net::SocketAddr;
 use std::pin::Pin;

src/tunnel/listeners/stdio.rs

@@ -1,6 +1,5 @@
 use crate::protocols::stdio;
-use crate::tunnel::RemoteAddr;
-use crate::LocalProtocol;
+use crate::tunnel::{LocalProtocol, RemoteAddr};
 use anyhow::{anyhow, Context};
 use std::pin::Pin;
 use std::task::Poll;

src/tunnel/listeners/tcp.rs

@@ -1,5 +1,5 @@
-use crate::tunnel::RemoteAddr;
-use crate::{protocols, LocalProtocol};
+use crate::protocols;
+use crate::tunnel::{LocalProtocol, RemoteAddr};
 use anyhow::{anyhow, Context};
 use std::net::SocketAddr;
 use std::pin::Pin;

src/tunnel/listeners/tproxy.rs

@@ -1,7 +1,7 @@
+use crate::protocols;
 use crate::protocols::udp;
 use crate::protocols::udp::{UdpStream, UdpStreamWriter};
-use crate::tunnel::{to_host_port, RemoteAddr};
-use crate::{protocols, LocalProtocol};
+use crate::tunnel::{to_host_port, LocalProtocol, RemoteAddr};
 use anyhow::{anyhow, Context};
 use std::io;
 use std::net::SocketAddr;

src/tunnel/listeners/udp.rs

@@ -1,7 +1,6 @@
 use crate::protocols::udp;
 use crate::protocols::udp::{UdpStream, UdpStreamWriter};
-use crate::tunnel::RemoteAddr;
-use crate::LocalProtocol;
+use crate::tunnel::{LocalProtocol, RemoteAddr};
 use anyhow::{anyhow, Context};
 use std::io;
 use std::net::SocketAddr;

src/tunnel/listeners/unix_sock.rs

@@ -1,7 +1,6 @@
 use crate::protocols::unix_sock;
 use crate::protocols::unix_sock::UnixListenerStream;
-use crate::tunnel::RemoteAddr;
-use crate::LocalProtocol;
+use crate::tunnel::{LocalProtocol, RemoteAddr};
 use anyhow::{anyhow, Context};
 use std::path::Path;
 use std::pin::Pin;

src/tunnel/mod.rs

@@ -5,21 +5,17 @@ pub mod server;
 mod tls_reloader;
 mod transport;
 
-use crate::{LocalProtocol, TlsClientConfig};
+use crate::TlsClientConfig;
 use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation};
 use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 use std::collections::HashSet;
 use std::fmt::{Debug, Display, Formatter};
-use std::io::{Error, IoSlice};
 use std::net::{IpAddr, SocketAddr};
 use std::ops::Deref;
-use std::pin::Pin;
+use std::path::PathBuf;
 use std::str::FromStr;
-use std::task::{Context, Poll};
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
-use tokio::net::TcpStream;
-use tokio_rustls::client::TlsStream;
+use std::time::Duration;
 use url::Host;
 use uuid::Uuid;
 
@@ -73,6 +69,61 @@ static JWT_DECODE: Lazy<(Validation, DecodingKey)> = Lazy::new(|| {
     (validation, DecodingKey::from_secret(JWT_SECRET))
 });
 
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub enum LocalProtocol {
+    Tcp {
+        proxy_protocol: bool,
+    },
+    Udp {
+        timeout: Option<Duration>,
+    },
+    Stdio,
+    Socks5 {
+        timeout: Option<Duration>,
+        credentials: Option<(String, String)>,
+    },
+    TProxyTcp,
+    TProxyUdp {
+        timeout: Option<Duration>,
+    },
+    HttpProxy {
+        timeout: Option<Duration>,
+        credentials: Option<(String, String)>,
+        proxy_protocol: bool,
+    },
+    ReverseTcp,
+    ReverseUdp {
+        timeout: Option<Duration>,
+    },
+    ReverseSocks5 {
+        timeout: Option<Duration>,
+        credentials: Option<(String, String)>,
+    },
+    ReverseHttpProxy {
+        timeout: Option<Duration>,
+        credentials: Option<(String, String)>,
+    },
+    ReverseUnix {
+        path: PathBuf,
+    },
+    Unix {
+        path: PathBuf,
+    },
+}
+
+impl LocalProtocol {
+    pub const fn is_reverse_tunnel(&self) -> bool {
+        matches!(
+            self,
+            Self::ReverseTcp
+                | Self::ReverseUdp { .. }
+                | Self::ReverseSocks5 { .. }
+                | Self::ReverseUnix { .. }
+                | Self::ReverseHttpProxy { .. }
+        )
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct RemoteAddr {
     pub protocol: LocalProtocol,
@@ -245,61 +296,6 @@ impl TryFrom<JwtTunnelConfig> for RemoteAddr {
     }
 }
 
-pub enum TransportStream {
-    Plain(TcpStream),
-    Tls(TlsStream<TcpStream>),
-}
-
-impl AsyncRead for TransportStream {
-    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
-        match self.get_mut() {
-            Self::Plain(cnx) => Pin::new(cnx).poll_read(cx, buf),
-            Self::Tls(cnx) => Pin::new(cnx).poll_read(cx, buf),
-        }
-    }
-}
-
-impl AsyncWrite for TransportStream {
-    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Error>> {
-        match self.get_mut() {
-            Self::Plain(cnx) => Pin::new(cnx).poll_write(cx, buf),
-            Self::Tls(cnx) => Pin::new(cnx).poll_write(cx, buf),
-        }
-    }
-
-    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
-        match self.get_mut() {
-            Self::Plain(cnx) => Pin::new(cnx).poll_flush(cx),
-            Self::Tls(cnx) => Pin::new(cnx).poll_flush(cx),
-        }
-    }
-
-    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
-        match self.get_mut() {
-            Self::Plain(cnx) => Pin::new(cnx).poll_shutdown(cx),
-            Self::Tls(cnx) => Pin::new(cnx).poll_shutdown(cx),
-        }
-    }
-
-    fn poll_write_vectored(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        bufs: &[IoSlice<'_>],
-    ) -> Poll<Result<usize, Error>> {
-        match self.get_mut() {
-            Self::Plain(cnx) => Pin::new(cnx).poll_write_vectored(cx, bufs),
-            Self::Tls(cnx) => Pin::new(cnx).poll_write_vectored(cx, bufs),
-        }
-    }
-
-    fn is_write_vectored(&self) -> bool {
-        match &self {
-            Self::Plain(cnx) => cnx.is_write_vectored(),
-            Self::Tls(cnx) => cnx.is_write_vectored(),
-        }
-    }
-}
-
 pub fn to_host_port(addr: SocketAddr) -> (Host, u16) {
     match addr.ip() {
         IpAddr::V4(ip) => (Host::Ipv4(ip), addr.port()),

src/tunnel/server/server.rs

@@ -15,8 +15,8 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::time::Duration;
 
-use crate::tunnel::RemoteAddr;
-use crate::{protocols, LocalProtocol};
+use crate::protocols;
+use crate::tunnel::{LocalProtocol, RemoteAddr};
 use hyper::body::Incoming;
 use hyper::server::conn::{http1, http2};
 use hyper::service::service_fn;