From f1231aa55a8de7ca9783681ea8413deb5b612366 Mon Sep 17 00:00:00 2001
From: cyejing <chenyejing@didiglobal.com>
Date: Fri, 28 Jun 2024 01:42:31 +0800
Subject: [PATCH] feat: ws add padding

---
 Cargo.lock                           |  4 +--
 Cargo.toml                           |  2 +-
 Makefile                             |  1 +
 shuttle-station/Cargo.toml           |  2 +-
 shuttle-station/src/dial.rs          | 37 +++++++++++++++++++---------
 shuttle-station/src/proto/padding.rs | 10 +++++---
 src/client.rs                        |  6 ++++-
 7 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ef920da..e0c34b1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1172,7 +1172,7 @@ dependencies = [
 
 [[package]]
 name = "shuttle"
-version = "0.6.9"
+version = "0.6.10"
 dependencies = [
  "anyhow",
  "atoi",
@@ -1202,7 +1202,7 @@ dependencies = [
 
 [[package]]
 name = "shuttle-station"
-version = "0.6.9"
+version = "0.6.10"
 dependencies = [
  "anyhow",
  "async-trait",
diff --git a/Cargo.toml b/Cargo.toml
index 18c8643..df8a1dd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "shuttle"
-version = "0.6.9"
+version = "0.6.10"
 edition = "2021"
 publish = false
 
diff --git a/Makefile b/Makefile
index 36f40ea..aff9023 100644
--- a/Makefile
+++ b/Makefile
@@ -21,4 +21,5 @@ check:
 
 publish:
 	cargo publish --registry crates-io --manifest-path shuttle-station/Cargo.toml
+	cargo doc
 
diff --git a/shuttle-station/Cargo.toml b/shuttle-station/Cargo.toml
index e23565a..0c4dd2a 100644
--- a/shuttle-station/Cargo.toml
+++ b/shuttle-station/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "shuttle-station"
-version = "0.6.9"
+version = "0.6.10"
 edition = "2021"
 authors = ["Born <jing23501@gmail.com>"]
 description = "shuttle-station"
diff --git a/shuttle-station/src/dial.rs b/shuttle-station/src/dial.rs
index b140d8c..56c9d94 100644
--- a/shuttle-station/src/dial.rs
+++ b/shuttle-station/src/dial.rs
@@ -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))
     }
 }
diff --git a/shuttle-station/src/proto/padding.rs b/shuttle-station/src/proto/padding.rs
index 54cae38..6220d78 100644
--- a/shuttle-station/src/proto/padding.rs
+++ b/shuttle-station/src/proto/padding.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)
     }
 }
diff --git a/src/client.rs b/src/client.rs
index 29946dd..64a0fbe 100644
--- a/src/client.rs
+++ b/src/client.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;