From 0d6948c97b15d9bcaa0467f74fa073718598d870 Mon Sep 17 00:00:00 2001
From: YinMo19 <me@yinmo19.top>
Date: Wed, 23 Apr 2025 09:58:20 +0800
Subject: [PATCH] [test] fix bug: twice enc. add ping pong log.

---
 src/websocket.rs | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/websocket.rs b/src/websocket.rs
index 2790899..73fe102 100644
--- a/src/websocket.rs
+++ b/src/websocket.rs
@@ -72,7 +72,18 @@ impl WsFramedStream {
         } else {
             log::info!("{:?}", url_str);
 
-            let (stream, _) = connect_async(url_str.into_client_request().unwrap()).await?;
+            let mut request = url_str
+                .into_client_request()
+                .map_err(|e| Error::new(ErrorKind::Other, e))?;
+
+            // 添加必要协议头
+            // request.headers_mut().insert(
+            //     "Sec-WebSocket-Protocol",
+            //     tungstenite::http::HeaderValue::from_static("rustdesk"),
+            // );
+
+            let (stream, _) =
+                timeout(Duration::from_millis(ms_timeout), connect_async(request)).await??;
 
             let addr = match stream.get_ref() {
                 MaybeTlsStream::Plain(tcp) => tcp.peer_addr()?,
@@ -141,10 +152,6 @@ impl WsFramedStream {
 
     #[inline]
     pub async fn send_raw(&mut self, msg: Vec<u8>) -> ResultType<()> {
-        let mut msg = msg;
-        if let Some(key) = self.encrypt.as_mut() {
-            msg = key.enc(&msg);
-        }
         self.send_bytes(bytes::Bytes::from(msg)).await
     }
 
@@ -189,6 +196,10 @@ impl WsFramedStream {
                     }
                     continue;
                 }
+                Ok(WsMessage::Pong(_)) => {
+                    log::debug!("Received pong");
+                    continue;
+                }
                 Ok(WsMessage::Close(_)) => return None,
                 Ok(_) => continue,
                 Err(e) => return Some(Err(Error::new(ErrorKind::Other, e))),