feat(transport): add websocket transport

Cargo.toml

@@ -11,7 +11,7 @@ build = "build.rs"
 include = ["src/**/*", "LICENSE", "README.md", "build.rs"]
 
 [features]
-default = ["server", "client", "tls", "noise", "hot-reload"]
+default = ["server", "client", "tls", "noise", "websocket", "hot-reload"]
 
 # Run as a server
 server = []
@@ -21,6 +21,8 @@ client = []
 tls = ["tokio-native-tls"]
 # Noise support
 noise = ["snowstorm", "base64"]
+# Websocket support
+websocket = ["tokio-tungstenite", "tokio-util", "futures-core", "futures-sink", "tls"]
 # Configuration hot-reload support
 hot-reload = ["notify"]
 
@@ -74,6 +76,10 @@ atty = "0.2"
 async-http-proxy = { version = "1.2", features = ["runtime-tokio", "basic-auth"] }
 async-socks5 = "0.5"
 url = { version = "2.2", features = ["serde"] }
+tokio-tungstenite = { version="0.20.1", optional = true}
+tokio-util = { version="0.7.9", optional = true, features = ["io"] }
+futures-core = { version="0.3.28", optional = true  }
+futures-sink = { version="0.3.28", optional = true }
 
 [build-dependencies]
 vergen = { version = "7.4.2", default-features = false, features = ["build", "git", "cargo"] }

Dockerfile

@@ -2,7 +2,7 @@ FROM rust:alpine as builder
 RUN apk add --no-cache musl-dev openssl openssl-dev pkgconfig
 WORKDIR /home/rust/src
 COPY . .
-RUN cargo build --locked --release --features client,server,noise,hot-reload
+RUN cargo build --locked --release --features client,server,noise,websocket,hot-reload
 RUN mkdir -p build-out/
 RUN cp target/release/rathole build-out/
 

README-zh.md

@@ -17,20 +17,20 @@ rathole，类似于 [frp](https://github.com/fatedier/frp) 和 [ngrok](https://g
 <!-- TOC -->
 
 - [rathole](#rathole)
-    - [Features](#features)
-    - [Quickstart](#quickstart)
-    - [Configuration](#configuration)
-        - [Logging](#logging)
-        - [Tuning](#tuning)
-    - [Benchmark](#benchmark)
-    - [Development Status](#development-status)
+  - [Features](#features)
+  - [Quickstart](#quickstart)
+  - [Configuration](#configuration)
+    - [Logging](#logging)
+    - [Tuning](#tuning)
+  - [Benchmark](#benchmark)
+  - [Development Status](#development-status)
 
 <!-- /TOC -->
 
 ## Features
 
-- **高性能** 具有更高的吞吐量，高并发下更稳定。见[Benchmark](#Benchmark)
-- **低资源消耗** 内存占用远低于同类工具。见[Benchmark](#Benchmark)。[二进制文件最小](docs/build-guide.md)可以到 **~500KiB**，可以部署在嵌入式设备如路由器上。
+- **高性能** 具有更高的吞吐量，高并发下更稳定。见[Benchmark](#benchmark)
+- **低资源消耗** 内存占用远低于同类工具。见[Benchmark](#benchmark)。[二进制文件最小](docs/build-guide.md)可以到 **~500KiB**，可以部署在嵌入式设备如路由器上。
 - **安全性** 每个服务单独强制鉴权。Server 和 Client 负责各自的配置。使用 Noise Protocol 可以简单地配置传输加密，而不需要自签证书。同时也支持 TLS。
 - **热重载** 支持配置文件热重载，动态修改端口转发服务。HTTP API 正在开发中。
 
@@ -91,7 +91,7 @@ local_addr = "127.0.0.1:22" # 需要被转发的服务的地址
 
 ## Configuration
 
-如果只有一个 `[server]` 和 `[client]` 块存在的话，`rathole` 可以根据配置文件的内容自动决定在服务器模式或客户端模式下运行，就像 [Quickstart](#Quickstart) 中的例子。
+如果只有一个 `[server]` 和 `[client]` 块存在的话，`rathole` 可以根据配置文件的内容自动决定在服务器模式或客户端模式下运行，就像 [Quickstart](#quickstart) 中的例子。
 
 但 `[client]` 和 `[server]` 块也可以放在一个文件中。然后在服务器端，运行 `rathole --server config.toml`。在客户端，运行 `rathole --client config.toml` 来明确告诉 `rathole` 运行模式。
 
@@ -126,6 +126,9 @@ pattern = "Noise_NK_25519_ChaChaPoly_BLAKE2s" # Optional. Default value as shown
 local_private_key = "key_encoded_in_base64" # Optional
 remote_public_key = "key_encoded_in_base64" # Optional
 
+[client.transport.websocket] # Necessary if `type` is "websocket"
+tls = true # If `true` then it will use settings in `client.transport.tls`
+
 [client.services.service1] # A service that needs forwarding. The name `service1` can change arbitrarily, as long as identical to the name in the server's configuration
 type = "tcp" # Optional. The protocol that needs forwarding. Possible values: ["tcp", "udp"]. Default: "tcp"
 token = "whatever" # Necessary if `client.default_token` not set
@@ -158,6 +161,9 @@ pattern = "Noise_NK_25519_ChaChaPoly_BLAKE2s"
 local_private_key = "key_encoded_in_base64"
 remote_public_key = "key_encoded_in_base64"
 
+[server.transport.websocket] # Necessary if `type` is "websocket"
+tls = true # If `true` then it will use settings in `server.transport.tls`
+
 [server.services.service1] # The service name must be identical to the client side
 type = "tcp" # Optional. Same as the client `[client.services.X.type]
 token = "whatever" # Necessary if `server.default_token` not set

README.md

@@ -18,20 +18,20 @@ rathole, like [frp](https://github.com/fatedier/frp) and [ngrok](https://github.
 <!-- TOC -->
 
 - [rathole](#rathole)
-    - [Features](#features)
-    - [Quickstart](#quickstart)
-    - [Configuration](#configuration)
-        - [Logging](#logging)
-        - [Tuning](#tuning)
-    - [Benchmark](#benchmark)
+  - [Features](#features)
+  - [Quickstart](#quickstart)
+  - [Configuration](#configuration)
+    - [Logging](#logging)
+    - [Tuning](#tuning)
+  - [Benchmark](#benchmark)
   - [Planning](#planning)
 
 <!-- /TOC -->
 
 ## Features
 
-- **High Performance** Much higher throughput can be achieved than frp, and more stable when handling a large volume of connections. See [Benchmark](#Benchmark)
-- **Low Resource Consumption** Consumes much fewer memory than similar tools. See [Benchmark](#Benchmark). [The binary can be](docs/build-guide.md) **as small as ~500KiB** to fit the constraints of devices, like embedded devices as routers.
+- **High Performance** Much higher throughput can be achieved than frp, and more stable when handling a large volume of connections. See [Benchmark](#benchmark)
+- **Low Resource Consumption** Consumes much fewer memory than similar tools. See [Benchmark](#benchmark). [The binary can be](docs/build-guide.md) **as small as ~500KiB** to fit the constraints of devices, like embedded devices as routers.
 - **Security** Tokens of services are mandatory and service-wise. The server and clients are responsible for their own configs. With the optional Noise Protocol, encryption can be configured at ease. No need to create a self-signed certificate! TLS is also supported.
 - **Hot Reload** Services can be added or removed dynamically by hot-reloading the configuration file. HTTP API is WIP.
 
@@ -93,7 +93,7 @@ To run `rathole` run as a background service on Linux, checkout the [systemd exa
 
 ## Configuration
 
-`rathole` can automatically determine to run in the server mode or the client mode, according to the content of the configuration file, if only one of `[server]` and `[client]` block is present, like the example in [Quickstart](#Quickstart).
+`rathole` can automatically determine to run in the server mode or the client mode, according to the content of the configuration file, if only one of `[server]` and `[client]` block is present, like the example in [Quickstart](#quickstart).
 
 But the `[client]` and `[server]` block can also be put in one file. Then on the server side, run `rathole --server config.toml` and on the client side, run `rathole --client config.toml` to explicitly tell `rathole` the running mode.
 
@@ -128,6 +128,9 @@ pattern = "Noise_NK_25519_ChaChaPoly_BLAKE2s" # Optional. Default value as shown
 local_private_key = "key_encoded_in_base64" # Optional
 remote_public_key = "key_encoded_in_base64" # Optional
 
+[client.transport.websocket] # Necessary if `type` is "websocket"
+tls = true # If `true` then it will use settings in `client.transport.tls`
+
 [client.services.service1] # A service that needs forwarding. The name `service1` can change arbitrarily, as long as identical to the name in the server's configuration
 type = "tcp" # Optional. The protocol that needs forwarding. Possible values: ["tcp", "udp"]. Default: "tcp"
 token = "whatever" # Necessary if `client.default_token` not set
@@ -160,6 +163,9 @@ pattern = "Noise_NK_25519_ChaChaPoly_BLAKE2s"
 local_private_key = "key_encoded_in_base64"
 remote_public_key = "key_encoded_in_base64"
 
+[server.transport.websocket] # Necessary if `type` is "websocket"
+tls = true # If `true` then it will use settings in `server.transport.tls`
+
 [server.services.service1] # The service name must be identical to the client side
 type = "tcp" # Optional. Same as the client `[client.services.X.type]
 token = "whatever" # Necessary if `server.default_token` not set

src/client.rs

@@ -24,6 +24,8 @@ use tracing::{debug, error, info, instrument, trace, warn, Instrument, Span};
 use crate::transport::NoiseTransport;
 #[cfg(feature = "tls")]
 use crate::transport::TlsTransport;
+#[cfg(feature = "websocket")]
+use crate::transport::WebsocketTransport;
 
 use crate::constants::{run_control_chan_backoff, UDP_BUFFER_SIZE, UDP_SENDQ_SIZE, UDP_TIMEOUT};
 
@@ -62,6 +64,15 @@ pub async fn run_client(
             #[cfg(not(feature = "noise"))]
             crate::helper::feature_not_compile("noise")
         }
+        TransportType::Websocket => {
+            #[cfg(feature = "websocket")]
+            {
+                let mut client = Client::<WebsocketTransport>::from(config).await?;
+                client.run(shutdown_rx, update_rx).await
+            }
+            #[cfg(not(feature = "websocket"))]
+            crate::helper::feature_not_compile("websocket")
+        }
     }
 }
 

src/config.rs

@@ -49,6 +49,8 @@ pub enum TransportType {
     Tls,
     #[serde(rename = "noise")]
     Noise,
+    #[serde(rename = "websocket")]
+    Websocket,
 }
 
 /// Per service config
@@ -75,8 +77,7 @@ impl ClientServiceConfig {
     }
 }
 
-#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
-#[derive(Default)]
+#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Default)]
 pub enum ServiceType {
     #[serde(rename = "tcp")]
     #[default]
@@ -85,8 +86,6 @@ pub enum ServiceType {
     Udp,
 }
 
-
-
 fn default_service_type() -> ServiceType {
     Default::default()
 }
@@ -136,6 +135,12 @@ pub struct NoiseConfig {
     // TODO: Maybe psk can be added
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(deny_unknown_fields)]
+pub struct WebsocketConfig {
+    pub tls: bool,
+}
+
 fn default_nodelay() -> bool {
     DEFAULT_NODELAY
 }
@@ -180,6 +185,7 @@ pub struct TransportConfig {
     pub tcp: TcpConfig,
     pub tls: Option<TlsConfig>,
     pub noise: Option<NoiseConfig>,
+    pub websocket: Option<WebsocketConfig>,
 }
 
 fn default_heartbeat_timeout() -> u64 {
@@ -313,6 +319,7 @@ impl Config {
                 // The check is done in transport
                 Ok(())
             }
+            TransportType::Websocket => Ok(()),
         }
     }
 

src/server.rs

@@ -27,6 +27,8 @@ use tracing::{debug, error, info, info_span, instrument, warn, Instrument, Span}
 use crate::transport::NoiseTransport;
 #[cfg(feature = "tls")]
 use crate::transport::TlsTransport;
+#[cfg(feature = "websocket")]
+use crate::transport::WebsocketTransport;
 
 type ServiceDigest = protocol::Digest; // SHA256 of a service name
 type Nonce = protocol::Digest; // Also called `session_key`
@@ -72,6 +74,15 @@ pub async fn run_server(
             #[cfg(not(feature = "noise"))]
             crate::helper::feature_not_compile("noise")
         }
+        TransportType::Websocket => {
+            #[cfg(feature = "websocket")]
+            {
+                let mut server = Server::<WebsocketTransport>::from(config).await?;
+                server.run(shutdown_rx, update_rx).await?;
+            }
+            #[cfg(not(feature = "websocket"))]
+            crate::helper::feature_not_compile("websocket")
+        }
     }
 
     Ok(())

src/transport/mod.rs

@@ -79,6 +79,11 @@ mod noise;
 #[cfg(feature = "noise")]
 pub use noise::NoiseTransport;
 
+#[cfg(feature = "websocket")]
+mod websocket;
+#[cfg(feature = "websocket")]
+pub use websocket::WebsocketTransport;
+
 #[derive(Debug, Clone, Copy)]
 struct Keepalive {
     // tcp_keepalive_time if the underlying protocol is TCP

src/transport/tls.rs

@@ -46,7 +46,7 @@ impl Transport for TlsTransport {
                 // if no trusted_root is specified, allow TlsConnector to use system default
                 let connector = native_tls::TlsConnector::builder().build()?;
                 Some(TlsConnector::from(connector))
-            },
+            }
         };
 
         let tls_acceptor = match config.pkcs12.as_ref() {