feat: ws add padding

Author: cyejing

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "shuttle"
-version = "0.6.9"
+version = "0.6.10"
 edition = "2021"
 publish = false
 

Cargo.toml

@@ -21,4 +21,5 @@ check:
 
 publish:
 	cargo publish --registry crates-io --manifest-path shuttle-station/Cargo.toml
+	cargo doc
 

Makefile

@@ -1,6 +1,6 @@
 [package]
 name = "shuttle-station"
-version = "0.6.9"
+version = "0.6.10"
 edition = "2021"
 authors = ["Born <[email protected]>"]
 description = "shuttle-station"

shuttle-station/Cargo.toml

@@ -3,6 +3,7 @@ use std::fmt;
 use anyhow::Context;
 use async_trait::async_trait;
 use futures::SinkExt;
+use futures::StreamExt;
 use socks5_proto::Address;
 use tokio::net::TcpStream;
 use tokio_rustls::client::TlsStream;
@@ -35,6 +36,7 @@ pub struct TrojanDial {
 pub struct WebSocketDial {
     remote_addr: String,
     hash: String,
+    padding: bool,
 }
 
 impl TrojanDial {
@@ -49,8 +51,12 @@ impl TrojanDial {
 }
 
 impl WebSocketDial {
-    pub fn new(remote_addr: String, hash: String) -> Self {
-        Self { remote_addr, hash }
+    pub fn new(remote_addr: String, hash: String, padding: bool) -> Self {
+        Self {
+            remote_addr,
+            hash,
+            padding,
+        }
     }
 }
 
@@ -131,15 +137,24 @@ impl Dial<WebSocketCopyStream<MaybeTlsStream<TcpStream>>> for WebSocketDial {
             .await
             .context(format!("WebSocket can't connect remote {}", remote_addr))?;
 
-        let mut buf: Vec<u8> = vec![];
-        let req = trojan::Request::new(self.hash.clone(), Command::Connect, addr);
-        req.write_to_buf(&mut buf);
-
-        ws.send(Message::Binary(buf))
-            .await
-            .context("WebSocket can't send")?;
-        ws.flush().await?;
-
+        if self.padding {
+            let mut buf: Vec<u8> = vec![];
+            let req = trojan::Request::new(self.hash.clone(), Command::Padding, addr);
+            req.write_to_buf(&mut buf);
+            ws.send(Message::Binary(buf))
+                .await
+                .context("WebSocket can't send")?;
+            ws.flush().await?;
+            let _ = ws.next().await;
+        } else {
+            let mut buf: Vec<u8> = vec![];
+            let req = trojan::Request::new(self.hash.clone(), Command::Connect, addr);
+            req.write_to_buf(&mut buf);
+            ws.send(Message::Binary(buf))
+                .await
+                .context("WebSocket can't send")?;
+            ws.flush().await?;
+        }
         Ok(WebSocketCopyStream::new(ws))
     }
 }

shuttle-station/src/dial.rs

@@ -9,6 +9,12 @@ pub struct Padding {
 }
 
 impl Padding {
+    pub fn new(start: u16, end: u16) -> Self {
+        let mut rng = rand::thread_rng();
+        let len = rng.gen_range(start..end);
+        Self { len }
+    }
+
     pub async fn read_from<R>(stream: &mut R) -> anyhow::Result<Self>
     where
         R: AsyncRead + Unpin,
@@ -52,8 +58,6 @@ impl Padding {
 
 impl Default for Padding {
     fn default() -> Self {
-        let mut rng = rand::thread_rng();
-        let len = rng.gen_range(100..3000);
-        Self { len }
+        Self::new(2500, 4500)
     }
 }

shuttle-station/src/proto/padding.rs

@@ -89,7 +89,11 @@ async fn proxy_handle(cc: Arc<ClientConfig>, ts: TcpStream) {
         (ProxyMode::Websocket, _) => {
             ProxyConnection::new(
                 ts,
-                Box::new(WebSocketDial::new(cc.remote_addr.clone(), cc.hash.clone())),
+                Box::new(WebSocketDial::new(
+                    cc.remote_addr.clone(),
+                    cc.hash.clone(),
+                    cc.padding,
+                )),
             )
             .handle()
             .await;