From 6798bf3780a8c69b23489bb54ca076f7179c46ef Mon Sep 17 00:00:00 2001
From: YinMo19 <me@yinmo19.top>
Date: Wed, 23 Apr 2025 08:33:01 +0800
Subject: [PATCH] [test] tcp success. Websocket test. (temp ver.)

---
 src/socket_client.rs | 22 +++++++++-------------
 src/websocket.rs     | 13 ++++++++++++-
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/src/socket_client.rs b/src/socket_client.rs
index b26e8f4..926e70e 100644
--- a/src/socket_client.rs
+++ b/src/socket_client.rs
@@ -116,18 +116,15 @@ pub async fn connect_tcp_local<
 ) -> ResultType<Stream> {
     let target_str = target.to_string();
 
-    // 根据目标地址协议决定连接方式
-    if target_str.starts_with("ws://") || target_str.starts_with("wss://") {
-        // WebSocket 连接逻辑
-        Ok(Stream::WebSocket(websocket::WsFramedStream::new(
-            target_str,
-            local,
-            None,
-            ms_timeout,
-        )
-        .await?))
+    // if target_str.starts_with("ws://") || target_str.starts_with("wss://") {
+    //     Ok(Stream::WebSocket(
+    //         websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?,
+    //     ))
+    if true {
+        Ok(Stream::WebSocket(
+            websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?,
+        ))
     } else {
-        // TCP 连接逻辑
         if let Some(conf) = Config::get_socks() {
             return Ok(Stream::Tcp(
                 FramedStream::connect(target, local, &conf, ms_timeout).await?,
@@ -139,8 +136,7 @@ pub async fn connect_tcp_local<
                 if local_addr.is_ipv6() && target_addr.is_ipv4() {
                     let resolved_target = query_nip_io(target_addr).await?;
                     return Ok(Stream::Tcp(
-                        FramedStream::new(resolved_target, Some(local_addr), ms_timeout)
-                            .await?,
+                        FramedStream::new(resolved_target, Some(local_addr), ms_timeout).await?,
                     ));
                 }
             }
diff --git a/src/websocket.rs b/src/websocket.rs
index 274791f..e560cb6 100644
--- a/src/websocket.rs
+++ b/src/websocket.rs
@@ -93,7 +93,7 @@ impl WsFramedStream {
 
     pub async fn from_tcp_stream(stream: TcpStream, addr: SocketAddr) -> ResultType<Self> {
         let ws_stream =
-            WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Server, None)
+            WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Client, None)
                 .await;
 
         Ok(Self {
@@ -169,6 +169,7 @@ impl WsFramedStream {
 
     #[inline]
     pub async fn next(&mut self) -> Option<Result<BytesMut, Error>> {
+        log::info!("test");
         loop {
             match self.stream.next().await? {
                 Ok(WsMessage::Binary(data)) => {
@@ -180,6 +181,16 @@ impl WsFramedStream {
                     }
                     return Some(Ok(bytes));
                 }
+                Ok(WsMessage::Ping(ping)) => {
+                    if let Err(e) = self.stream.send(WsMessage::Pong(ping)).await {
+                        return Some(Err(Error::new(
+                            ErrorKind::Other,
+                            format!("Failed to send pong: {}", e),
+                        )));
+                    }
+                    continue;
+                }
+                Ok(WsMessage::Close(_)) => return None,
                 Ok(_) => continue,
                 Err(e) => return Some(Err(Error::new(ErrorKind::Other, e))),
             }