Start to forward messages

tcp_relay_client/Cargo.toml

@@ -10,4 +10,6 @@ log = "0.4.17"
 env_logger = "0.9.0"
 reqwest = { version = "0.11", features = ["json"] }
 tokio = { version = "1", features = ["full"] }
-futures = "0.3.24"
+futures = "0.3.24"
+tokio-tungstenite = "0.17.2"
+urlencoding = "2.1.0"

tcp_relay_client/src/main.rs

@@ -50,9 +50,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         let listen_address = format!("{}:{}", args.listen_address, port.port);
 
         let h = tokio::spawn(relay_client(
-            args.token.clone(),
-            port.id,
-            args.listen_address.clone(),
+            format!("{}/ws?id={}&token={}",
+                    args.relay_url, port.id, urlencoding::encode(&args.token))
+                .replace("http", "ws"),
             listen_address,
         ));
         handles.push(h);

tcp_relay_client/src/relay_client.rs

@@ -1,8 +1,10 @@
+use futures::{SinkExt, StreamExt};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+use tokio_tungstenite::tungstenite::Message;
 
-use tokio::net::TcpListener;
-
-pub async fn relay_client(token: String, port_id: usize, server: String, listen_address: String) {
-    log::info!("({}) Start to listen on {}", port_id, listen_address);
+pub async fn relay_client(ws_url: String, listen_address: String) {
+    log::info!("Start to listen on {}", listen_address);
     let listener = match TcpListener::bind(&listen_address).await {
         Ok(l) => l,
         Err(e) => {
@@ -12,13 +14,68 @@ pub async fn relay_client(token: String, port_id: usize, server: String, listen_
     };
 
     loop {
-        let (mut socket, _) = listener.accept().await
+        let (socket, _) = listener.accept().await
             .expect("Failed to accept new connection!");
 
-        let token = token.clone();
-        let listen_address = listen_address.clone();
-        tokio::spawn(async move {
-            log::info!("New connection to {}", &listen_address);
-        });
+        tokio::spawn(relay_connection(ws_url.clone(), socket));
+    }
+}
+
+/// Relay connection
+///
+/// WS read => TCP write
+/// TCP read => WS write
+async fn relay_connection(ws_url: String, socket: TcpStream) {
+    log::debug!("Connecting to {}...", ws_url);
+    let (ws_stream, _) = tokio_tungstenite::connect_async(ws_url)
+        .await.expect("Failed to connect to server relay!");
+
+
+    let (mut tcp_read, mut tcp_write) = socket.into_split();
+    let (mut ws_write, mut ws_read) =
+        ws_stream.split();
+
+    // TCP read -> WS write
+    let future = async move {
+        let mut buff: [u8; 5000] = [0; 5000];
+        loop {
+            match tcp_read.read(&mut buff).await {
+                Ok(s) => {
+                    if let Err(e) = ws_write.send(Message::Binary(Vec::from(&buff[0..s]))).await {
+                        log::error!(
+                            "Failed to write to WS connection! {:?} Exiting TCP read -> WS write loop...",e);
+                    }
+                }
+                Err(e) => {
+                    log::error!(
+                        "Failed to read from TCP connection! {:?} Exitin TCP read -> WS write loop...",
+                        e
+                    );
+                    break;
+                }
+            }
+        }
+    };
+    tokio::spawn(future);
+
+    // WS read -> TCP write
+    while let Some(m) = ws_read.next().await {
+        match m {
+            Err(e) => {
+                log::error!("Failed to read from WebSocket. Breaking read loop... {:?}", e);
+                break;
+            }
+            Ok(Message::Binary(b)) => {
+                if let Err(e) = tcp_write.write_all(&b).await {
+                    log::error!("Failed to forward message to websocket. Closing reading end... {:?}", e);
+                    break;
+                };
+            }
+            Ok(Message::Close(_r)) => {
+                log::info!("Server asked to close this WebSocket connection");
+                break;
+            }
+            Ok(m) => log::info!("{:?}", m)
+        }
     }
 }