split

src/main.rs

@@ -12,7 +12,8 @@ use crate::tunnel::listeners::{
     new_stdio_listener, HttpProxyTunnelListener, Socks5TunnelListener, TcpTunnelListener, UdpTunnelListener,
 };
 use crate::tunnel::server::{TlsServerConfig, WsServer, WsServerConfig};
-use crate::tunnel::{to_host_port, LocalProtocol, RemoteAddr, TransportAddr, TransportScheme};
+use crate::tunnel::transport::{TransportAddr, TransportScheme};
+use crate::tunnel::{to_host_port, LocalProtocol, RemoteAddr};
 use anyhow::{anyhow, Context};
 use base64::Engine;
 use clap::Parser;

src/protocols/tls/server.rs

@@ -10,7 +10,7 @@ use tokio_rustls::client::TlsStream;
 
 use crate::tunnel::client::WsClientConfig;
 use crate::tunnel::server::TlsServerConfig;
-use crate::tunnel::TransportAddr;
+use crate::tunnel::transport::TransportAddr;
 use tokio_rustls::rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
 use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime};
 use tokio_rustls::rustls::server::WebPkiClientVerifier;

src/tunnel/client/client.rs

@@ -5,8 +5,8 @@ use crate::tunnel::connectors::TunnelConnector;
 use crate::tunnel::listeners::TunnelListener;
 use crate::tunnel::tls_reloader::TlsReloader;
 use crate::tunnel::transport::io::{TunnelReader, TunnelWriter};
-use crate::tunnel::transport::jwt_token_to_tunnel;
-use crate::tunnel::{RemoteAddr, TransportScheme};
+use crate::tunnel::transport::{jwt_token_to_tunnel, TransportScheme};
+use crate::tunnel::RemoteAddr;
 use anyhow::Context;
 use futures_util::pin_mut;
 use hyper::header::COOKIE;

src/tunnel/client/config.rs

@@ -1,5 +1,5 @@
 use crate::protocols::dns::DnsResolver;
-use crate::tunnel::TransportAddr;
+use crate::tunnel::transport::TransportAddr;
 use hyper::header::{HeaderName, HeaderValue};
 use once_cell::sync::Lazy;
 use parking_lot::RwLock;
@@ -29,17 +29,6 @@ pub struct WsClientConfig {
 }
 
 impl WsClientConfig {
-    pub const fn websocket_scheme(&self) -> &'static str {
-        match self.remote_addr.tls().is_some() {
-            false => "ws",
-            true => "wss",
-        }
-    }
-
-    pub fn websocket_host_url(&self) -> String {
-        format!("{}:{}", self.remote_addr.host(), self.remote_addr.port())
-    }
-
     pub fn tls_server_name(&self) -> ServerName<'static> {
         static INVALID_DNS_NAME: Lazy<DnsName> = Lazy::new(|| DnsName::try_from("dns-name-invalid.com").unwrap());
 

src/tunnel/mod.rs

@@ -3,14 +3,12 @@ pub mod connectors;
 pub mod listeners;
 pub mod server;
 mod tls_reloader;
-mod transport;
+pub mod transport;
 
-use crate::TlsClientConfig;
 use serde::{Deserialize, Serialize};
-use std::fmt::{Debug, Display, Formatter};
+use std::fmt::Debug;
 use std::net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6};
 use std::path::PathBuf;
-use std::str::FromStr;
 use std::time::Duration;
 use url::Host;
 
@@ -72,7 +70,7 @@ impl LocalProtocol {
     }
 
     pub const fn is_dynamic_reverse_tunnel(&self) -> bool {
-        matches!(self, |Self::ReverseSocks5 { .. }| Self::ReverseHttpProxy { .. })
+        matches!(self, Self::ReverseSocks5 { .. } | Self::ReverseHttpProxy { .. })
     }
 }
 
@@ -83,160 +81,6 @@ pub struct RemoteAddr {
     pub port: u16,
 }
 
-#[derive(Copy, Clone, Debug)]
-pub enum TransportScheme {
-    Ws,
-    Wss,
-    Http,
-    Https,
-}
-
-impl TransportScheme {
-    pub const fn values() -> &'static [Self] {
-        &[Self::Ws, Self::Wss, Self::Http, Self::Https]
-    }
-    pub const fn to_str(self) -> &'static str {
-        match self {
-            Self::Ws => "ws",
-            Self::Wss => "wss",
-            Self::Http => "http",
-            Self::Https => "https",
-        }
-    }
-
-    pub fn alpn_protocols(&self) -> Vec<Vec<u8>> {
-        match self {
-            Self::Ws => vec![],
-            Self::Wss => vec![b"http/1.1".to_vec()],
-            Self::Http => vec![],
-            Self::Https => vec![b"h2".to_vec()],
-        }
-    }
-}
-impl FromStr for TransportScheme {
-    type Err = ();
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "https" => Ok(Self::Https),
-            "http" => Ok(Self::Http),
-            "wss" => Ok(Self::Wss),
-            "ws" => Ok(Self::Ws),
-            _ => Err(()),
-        }
-    }
-}
-
-impl Display for TransportScheme {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        f.write_str(self.to_str())
-    }
-}
-
-#[derive(Clone)]
-pub enum TransportAddr {
-    Wss {
-        tls: TlsClientConfig,
-        scheme: TransportScheme,
-        host: Host,
-        port: u16,
-    },
-    Ws {
-        scheme: TransportScheme,
-        host: Host,
-        port: u16,
-    },
-    Https {
-        scheme: TransportScheme,
-        tls: TlsClientConfig,
-        host: Host,
-        port: u16,
-    },
-    Http {
-        scheme: TransportScheme,
-        host: Host,
-        port: u16,
-    },
-}
-
-impl Debug for TransportAddr {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        f.write_fmt(format_args!("{}://{}:{}", self.scheme(), self.host(), self.port()))
-    }
-}
-
-impl TransportAddr {
-    pub fn new(scheme: TransportScheme, host: Host, port: u16, tls: Option<TlsClientConfig>) -> Option<Self> {
-        match scheme {
-            TransportScheme::Https => Some(Self::Https {
-                scheme: TransportScheme::Https,
-                tls: tls?,
-                host,
-                port,
-            }),
-            TransportScheme::Http => Some(Self::Http {
-                scheme: TransportScheme::Http,
-                host,
-                port,
-            }),
-            TransportScheme::Wss => Some(Self::Wss {
-                scheme: TransportScheme::Wss,
-                tls: tls?,
-                host,
-                port,
-            }),
-            TransportScheme::Ws => Some(Self::Ws {
-                scheme: TransportScheme::Ws,
-                host,
-                port,
-            }),
-        }
-    }
-    pub const fn is_websocket(&self) -> bool {
-        matches!(self, Self::Ws { .. } | Self::Wss { .. })
-    }
-
-    pub const fn is_http2(&self) -> bool {
-        matches!(self, Self::Http { .. } | Self::Https { .. })
-    }
-
-    pub const fn tls(&self) -> Option<&TlsClientConfig> {
-        match self {
-            Self::Wss { tls, .. } => Some(tls),
-            Self::Https { tls, .. } => Some(tls),
-            Self::Ws { .. } => None,
-            Self::Http { .. } => None,
-        }
-    }
-
-    pub const fn host(&self) -> &Host {
-        match self {
-            Self::Wss { host, .. } => host,
-            Self::Ws { host, .. } => host,
-            Self::Https { host, .. } => host,
-            Self::Http { host, .. } => host,
-        }
-    }
-
-    pub const fn port(&self) -> u16 {
-        match self {
-            Self::Wss { port, .. } => *port,
-            Self::Ws { port, .. } => *port,
-            Self::Https { port, .. } => *port,
-            Self::Http { port, .. } => *port,
-        }
-    }
-
-    pub const fn scheme(&self) -> &TransportScheme {
-        match self {
-            Self::Wss { scheme, .. } => scheme,
-            Self::Ws { scheme, .. } => scheme,
-            Self::Https { scheme, .. } => scheme,
-            Self::Http { scheme, .. } => scheme,
-        }
-    }
-}
-
 pub fn to_host_port(addr: SocketAddr) -> (Host, u16) {
     match addr.ip() {
         IpAddr::V4(ip) => (Host::Ipv4(ip), addr.port()),

src/tunnel/server/reverse_tunnel.rs

@@ -1,6 +1,6 @@
 use crate::tunnel::listeners::TunnelListener;
 use crate::tunnel::RemoteAddr;
-use ahash::{HashMap, HashMapExt};
+use ahash::{AHashMap};
 use anyhow::anyhow;
 use futures_util::{pin_mut, StreamExt};
 use log::warn;
@@ -29,13 +29,13 @@ impl<T: TunnelListener> Clone for ReverseTunnelItem<T> {
 }
 
 pub struct ReverseTunnelServer<T: TunnelListener> {
-    servers: Arc<Mutex<HashMap<SocketAddr, ReverseTunnelItem<T>>>>,
+    servers: Arc<Mutex<AHashMap<SocketAddr, ReverseTunnelItem<T>>>>,
 }
 
 impl<T: TunnelListener> ReverseTunnelServer<T> {
     pub fn new() -> Self {
         Self {
-            servers: Arc::new(Mutex::new(HashMap::with_capacity(1))),
+            servers: Arc::new(Mutex::new(AHashMap::with_capacity(1))),
         }
     }
 

src/tunnel/transport/http2.rs

@@ -1,8 +1,8 @@
 use super::io::{TunnelRead, TunnelWrite, MAX_PACKET_LENGTH};
 use crate::tunnel::client::WsClient;
-use crate::tunnel::transport::headers_from_file;
 use crate::tunnel::transport::jwt::tunnel_to_jwt_token;
-use crate::tunnel::{RemoteAddr, TransportScheme};
+use crate::tunnel::transport::{headers_from_file, TransportScheme};
+use crate::tunnel::RemoteAddr;
 use anyhow::{anyhow, Context};
 use bytes::{Bytes, BytesMut};
 use http_body_util::{BodyExt, BodyStream, StreamBody};

src/tunnel/transport/mod.rs

@@ -9,11 +9,15 @@ use tracing::error;
 pub mod http2;
 pub mod io;
 mod jwt;
+mod types;
 pub mod websocket;
+
 pub use jwt::jwt_token_to_tunnel;
 pub use jwt::tunnel_to_jwt_token;
 pub use jwt::JwtTunnelConfig;
 pub use jwt::JWT_HEADER_PREFIX;
+pub use types::TransportAddr;
+pub use types::TransportScheme;
 
 #[allow(clippy::type_complexity)]
 #[inline]

src/tunnel/transport/types.rs

@@ -0,0 +1,151 @@
+use crate::tunnel::client::TlsClientConfig;
+use std::fmt::{Debug, Display, Formatter};
+use std::str::FromStr;
+use url::Host;
+
+#[derive(Copy, Clone, Debug)]
+pub enum TransportScheme {
+    Ws,
+    Wss,
+    Http,
+    Https,
+}
+
+impl TransportScheme {
+    pub const fn values() -> &'static [Self] {
+        &[Self::Ws, Self::Wss, Self::Http, Self::Https]
+    }
+    pub const fn to_str(self) -> &'static str {
+        match self {
+            Self::Ws => "ws",
+            Self::Wss => "wss",
+            Self::Http => "http",
+            Self::Https => "https",
+        }
+    }
+
+    pub fn alpn_protocols(&self) -> Vec<Vec<u8>> {
+        match self {
+            Self::Ws => vec![],
+            Self::Wss => vec![b"http/1.1".to_vec()],
+            Self::Http => vec![],
+            Self::Https => vec![b"h2".to_vec()],
+        }
+    }
+}
+impl FromStr for TransportScheme {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "https" => Ok(Self::Https),
+            "http" => Ok(Self::Http),
+            "wss" => Ok(Self::Wss),
+            "ws" => Ok(Self::Ws),
+            _ => Err(()),
+        }
+    }
+}
+
+impl Display for TransportScheme {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.to_str())
+    }
+}
+
+#[derive(Clone)]
+pub enum TransportAddr {
+    Wss {
+        tls: TlsClientConfig,
+        scheme: TransportScheme,
+        host: Host,
+        port: u16,
+    },
+    Ws {
+        scheme: TransportScheme,
+        host: Host,
+        port: u16,
+    },
+    Https {
+        scheme: TransportScheme,
+        tls: TlsClientConfig,
+        host: Host,
+        port: u16,
+    },
+    Http {
+        scheme: TransportScheme,
+        host: Host,
+        port: u16,
+    },
+}
+
+impl Debug for TransportAddr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.write_fmt(format_args!("{}://{}:{}", self.scheme(), self.host(), self.port()))
+    }
+}
+
+impl TransportAddr {
+    pub fn new(scheme: TransportScheme, host: Host, port: u16, tls: Option<TlsClientConfig>) -> Option<Self> {
+        match scheme {
+            TransportScheme::Https => Some(Self::Https {
+                scheme: TransportScheme::Https,
+                tls: tls?,
+                host,
+                port,
+            }),
+            TransportScheme::Http => Some(Self::Http {
+                scheme: TransportScheme::Http,
+                host,
+                port,
+            }),
+            TransportScheme::Wss => Some(Self::Wss {
+                scheme: TransportScheme::Wss,
+                tls: tls?,
+                host,
+                port,
+            }),
+            TransportScheme::Ws => Some(Self::Ws {
+                scheme: TransportScheme::Ws,
+                host,
+                port,
+            }),
+        }
+    }
+
+    pub const fn tls(&self) -> Option<&TlsClientConfig> {
+        match self {
+            Self::Wss { tls, .. } => Some(tls),
+            Self::Https { tls, .. } => Some(tls),
+            Self::Ws { .. } => None,
+            Self::Http { .. } => None,
+        }
+    }
+
+    pub const fn host(&self) -> &Host {
+        match self {
+            Self::Wss { host, .. } => host,
+            Self::Ws { host, .. } => host,
+            Self::Https { host, .. } => host,
+            Self::Http { host, .. } => host,
+        }
+    }
+
+    pub const fn port(&self) -> u16 {
+        match self {
+            Self::Wss { port, .. } => *port,
+            Self::Ws { port, .. } => *port,
+            Self::Https { port, .. } => *port,
+            Self::Http { port, .. } => *port,
+        }
+    }
+
+    pub const fn scheme(&self) -> &TransportScheme {
+        match self {
+            Self::Wss { scheme, .. } => scheme,
+            Self::Ws { scheme, .. } => scheme,
+            Self::Https { scheme, .. } => scheme,
+            Self::Http { scheme, .. } => scheme,
+        }
+    }
+}