Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tg_unblock"
-version = "0.3.1"
+version = "0.3.0"
 edition = "2021"
 
 [dependencies]

src/ws_proxy.rs

@@ -235,6 +235,7 @@ async fn relay_via_ws(
     let url = ws_url(dc);
     let mut request = url.as_str().into_client_request()?;
 
+    // Required by the Telegram WebSocket transport protocol
     request
         .headers_mut()
         .insert("Sec-WebSocket-Protocol", "binary".parse()?);
@@ -243,56 +244,55 @@ async fn relay_via_ws(
         native_tls::TlsConnector::new().map_err(|e| format!("TLS: {}", e))?,
     );
 
-    let (mut ws, _resp) = tokio_tungstenite::connect_async_tls_with_config(
+    let (ws, _resp) = tokio_tungstenite::connect_async_tls_with_config(
         request, None, false, Some(connector),
     )
     .await?;
 
+    let (mut ws_tx, mut ws_rx) = ws.split();
     let (mut tcp_rx, mut tcp_tx) = tokio::io::split(tcp_stream);
 
     // Send the buffered 64-byte init as the first WebSocket message
-    ws.send(tungstenite::Message::Binary(init.to_vec())).await?;
+    ws_tx
+        .send(tungstenite::Message::Binary(init.to_vec()))
+        .await?;
 
-    // Single loop: handles TCP→WS, WS→TCP, and Ping/Pong in one place.
+    let up = async {
-    // This ensures Pong replies are sent immediately so the server
-    // doesn't kill the connection after a timeout.
         let mut buf = vec![0u8; 32768];
-
         loop {
-        tokio::select! {
+            match tcp_rx.read(&mut buf).await {
-            biased;
+                Ok(0) => break,
-
+                Ok(n) => {
-            ws_msg = ws.next() => {
+                    let msg = tungstenite::Message::Binary(buf[..n].to_vec());
-                match ws_msg {
+                    if ws_tx.send(msg).await.is_err() {
-                    Some(Ok(tungstenite::Message::Binary(data))) => {
-                        if tcp_tx.write_all(data.as_ref()).await.is_err() {
                         break;
                     }
                 }
-                    Some(Ok(tungstenite::Message::Ping(payload))) => {
+                Err(_) => break,
-                        let _ = ws.send(tungstenite::Message::Pong(payload)).await;
             }
-                    Some(Ok(tungstenite::Message::Close(_))) | None => break,
+        }
-                    Some(Err(_)) => break,
+        let _ = ws_tx.close().await;
+    };
+
+    let down = async {
+        while let Some(Ok(msg)) = ws_rx.next().await {
+            match msg {
+                tungstenite::Message::Binary(data) => {
+                    if tcp_tx.write_all(&data).await.is_err() {
+                        break;
+                    }
+                }
+                tungstenite::Message::Close(_) => break,
                 _ => {}
             }
         }
+    };
 
-            n = tcp_rx.read(&mut buf) => {
+    tokio::select! {
-                match n {
+        _ = up => {}
-                    Ok(0) | Err(_) => break,
+        _ = down => {}
-                    Ok(n) => {
-                        let msg = tungstenite::Message::Binary(buf[..n].to_vec());
-                        if ws.send(msg).await.is_err() {
-                            break;
-                        }
-                    }
-                }
-            }
-        }
     }
 
-    let _ = ws.close(None).await;
     Ok(())
 }