update

Cargo.lock

@@ -2,6 +2,21 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "addr2line"
+version = "0.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9ecd88a8c8378ca913a680cd98f0f13ac67383d35993f86c90a70e3f137816b"
+dependencies = [
+ "gimli",
+]
+
+[[package]]
+name = "adler"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
+
 [[package]]
 name = "anyhow"
 version = "1.0.58"
@@ -36,12 +51,37 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
+[[package]]
+name = "backtrace"
+version = "0.3.65"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "11a17d453482a265fd5f8479f2a3f405566e6ca627837aaddb85af8b1ab8ef61"
+dependencies = [
+ "addr2line",
+ "cc",
+ "cfg-if",
+ "libc",
+ "miniz_oxide",
+ "object",
+ "rustc-demangle",
+]
+
 [[package]]
 name = "base64"
 version = "0.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "904dfeac50f3cdaba28fc6f57fdcddb75f49ed61346676a78c4ffe55877802fd"
 
+[[package]]
+name = "better-panic"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6fa9e1d11a268684cbd90ed36370d7577afb6c62d912ddff5c15fc34343e5036"
+dependencies = [
+ "backtrace",
+ "console",
+]
+
 [[package]]
 name = "bitflags"
 version = "1.3.2"
@@ -69,6 +109,12 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c4872d67bab6358e59559027aa3b9157c53d9358c51423c17554809a8858e0f8"
 
+[[package]]
+name = "cc"
+version = "1.0.73"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2fff2a6927b3bb87f9595d67196a70493f627687a71d87a0d692242c33f58c11"
+
 [[package]]
 name = "cfg-if"
 version = "1.0.0"
@@ -114,6 +160,19 @@ dependencies = [
  "os_str_bytes",
 ]
 
+[[package]]
+name = "console"
+version = "0.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a28b32d32ca44b70c3e4acd7db1babf555fa026e385fb95f18028f88848b3c31"
+dependencies = [
+ "encode_unicode",
+ "libc",
+ "once_cell",
+ "terminal_size",
+ "winapi",
+]
+
 [[package]]
 name = "cpufeatures"
 version = "0.2.2"
@@ -143,6 +202,12 @@ dependencies = [
  "crypto-common",
 ]
 
+[[package]]
+name = "encode_unicode"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f"
+
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -212,6 +277,12 @@ dependencies = [
  "wasi",
 ]
 
+[[package]]
+name = "gimli"
+version = "0.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78cc372d058dcf6d5ecd98510e7fbc9e5aec4d21de70f65fea8fecebcd881bd4"
+
 [[package]]
 name = "hashbrown"
 version = "0.12.1"
@@ -314,6 +385,15 @@ version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
 
+[[package]]
+name = "miniz_oxide"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f5c75688da582b8ffc1f1799e9db273f32133c49e048f614d22ec3256773ccc"
+dependencies = [
+ "adler",
+]
+
 [[package]]
 name = "mio"
 version = "0.8.4"
@@ -345,6 +425,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "object"
+version = "0.28.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e42c982f2d955fac81dd7e1d0e1426a7d702acd9c98d19ab01083a6a0328c424"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "once_cell"
 version = "1.12.0"
@@ -461,12 +550,19 @@ dependencies = [
  "getrandom",
 ]
 
+[[package]]
+name = "rustc-demangle"
+version = "0.1.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ef03e0a2b150c7a90d01faf6254c9c48a41e95fb2a8c2ac1c6f0d2b9aefc342"
+
 [[package]]
 name = "servers"
 version = "0.1.0"
 dependencies = [
  "anyhow",
  "async-trait",
+ "better-panic",
  "clap",
  "futures-util",
  "libloading",
@@ -541,6 +637,16 @@ dependencies = [
  "winapi-util",
 ]
 
+[[package]]
+name = "terminal_size"
+version = "0.1.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "633c1a546cee861a1a6d0dc69ebeca693bf4296661ba7852b9d21d159e0506df"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
 [[package]]
 name = "textwrap"
 version = "0.15.0"

Cargo.toml

@@ -19,6 +19,7 @@ codegen-units = 1
 [dependencies]
 anyhow = "1.0.58"
 async-trait = "0.1.56"
+better-panic = "0.3.0"
 libloading = "0.7.3"
 simplelog = "0.12.0"
 tokio-tungstenite = "0.17.1"

src/main.rs

@@ -1,12 +1,12 @@
-use std::{net::TcpListener, fs::File};
+use std::{fs::File, net::TcpListener};
 
 use clap::Parser;
-use log::{LevelFilter, info};
+use log::{error, info, LevelFilter};
 use servers::{
     plugins::loader,
     tcp::{handle_connection, handle_websocket, Client},
 };
-use simplelog::{CombinedLogger, TermLogger, WriteLogger, Config, TerminalMode, ColorChoice};
+use simplelog::{ColorChoice, CombinedLogger, Config, TermLogger, TerminalMode, WriteLogger};
 
 #[derive(Parser)]
 #[clap(
@@ -50,13 +50,22 @@ struct Cli {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    // init better panic
+    better_panic::install();
     // init logger
-    CombinedLogger::init(
+    CombinedLogger::init(vec![
-        vec![
+        TermLogger::new(
-            TermLogger::new(LevelFilter::Trace, Config::default(), TerminalMode::Mixed, ColorChoice::Auto),
+            LevelFilter::Trace,
-            WriteLogger::new(LevelFilter::Debug, Config::default(), File::create("server.log").unwrap()),
+            Config::default(),
-        ]
+            TerminalMode::Mixed,
-    )?;
+            ColorChoice::Auto,
+        ),
+        WriteLogger::new(
+            LevelFilter::Debug,
+            Config::default(),
+            File::create("server.log").unwrap(),
+        ),
+    ])?;
 
     // parse cli args
     let cli = Cli::parse();
@@ -89,9 +98,17 @@ async fn start_tcp_server(host: String, port: String) -> anyhow::Result<()> {
     // Accepts a new incoming connection from this listener.
     while let Ok((stream, _address)) = listener.accept() {
         let client = Client::new(stream);
+        let plugin_manager = plugin_manager.clone();
 
         // handle client connection in new thread
-        tokio::spawn(handle_connection(client, plugin_manager.clone()));
+        tokio::spawn(async move {
+            let ip = client.stream.peer_addr().unwrap();
+
+            match handle_connection(client, plugin_manager).await {
+                Ok(_) => (),
+                Err(err) => error!("Client {}, {}", ip, err),
+            }
+        });
     }
 
     // server for a unexpectedly reason be terminated
@@ -108,7 +125,14 @@ async fn start_ws_server(host: String, port: String, tcp_port: String) -> anyhow
     // Accepts a new incoming connection from this listener.
     while let Ok((stream, _address)) = listener.accept().await {
         let tcp_port = tcp_port.clone();
-        tokio::spawn(handle_websocket(stream, tcp_port));
+        tokio::spawn(async {
+            let ip = stream.peer_addr().unwrap();
+
+            match handle_websocket(stream, tcp_port).await {
+                Ok(_) => (),
+                Err(err) => error!("Client {}, {}", ip, err),
+            }
+        });
     }
 
     // server for a unexpectedly reason be terminated

src/tcp/handle_connection.rs

@@ -1,6 +1,6 @@
 use std::io::Write;
 
-use log::{error, trace, info};
+use log::{error, info, trace};
 
 use crate::plugins::PluginManagerType;
 

src/tcp/handle_websocket.rs

@@ -1,11 +1,18 @@
 #![allow(clippy::unused_io_amount)]
 
-use futures_util::{SinkExt, StreamExt};
+use futures_util::{
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt,
+};
 use log::info;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
-    net::TcpStream,
+    net::{
+        tcp::{OwnedReadHalf, OwnedWriteHalf},
+        TcpStream,
+    },
 };
+use tokio_tungstenite::WebSocketStream;
 use tungstenite::Message;
 
 use super::MAX_PACKET_LEN;
@@ -21,33 +28,49 @@ pub async fn handle_websocket(stream: TcpStream, tcp_port: String) -> anyhow::Re
     let tcp_stream = TcpStream::connect(format!("0.0.0.0:{}", tcp_port)).await?;
 
     // split streams
-    let (mut tcp_read, mut tcp_write) = tcp_stream.into_split();
+    let (tcp_read, tcp_write) = tcp_stream.into_split();
-    let (mut ws_write, mut ws_read) = ws_stream.split();
+    let (ws_write, ws_read) = ws_stream.split();
 
     // tcp read -> ws write
-    tokio::spawn(async move {
+    tokio::spawn(tcp_to_ws(tcp_read, ws_write));
-        // allocate an empty buffer
-        let mut buf = [0; MAX_PACKET_LEN];
-
-        loop {
-            // read buffer from tcp
-            let len = tcp_read.read(&mut buf).await.unwrap();
-
-            if len > 0 {
-                // select only used bytes from the buffer
-                let recv_buf = &buf[0..len];
-                // covert &[u8] buffer to a vector
-                let recv_vec = recv_buf.to_vec();
-                // create a `Message` type from buffer Vec<u8>
-                let msg = Message::Binary(recv_vec);
-
-                // write buffer to websocket
-                ws_write.send(msg).await.unwrap();
-            }
-        }
-    });
 
     // ws read -> tcp write
+    ws_to_tcp(tcp_write, ws_read).await?;
+
+    Ok(())
+}
+
+/// Tcp read -> WebSocket write
+async fn tcp_to_ws(
+    mut tcp_read: OwnedReadHalf,
+    mut ws_write: SplitSink<WebSocketStream<TcpStream>, Message>,
+) -> anyhow::Result<()> {
+    // allocate an empty buffer
+    let mut buf = [0; MAX_PACKET_LEN];
+
+    loop {
+        // read buffer from tcp
+        let len = tcp_read.read(&mut buf).await?;
+
+        if len > 0 {
+            // select only used bytes from the buffer
+            let recv_buf = &buf[0..len];
+            // covert &[u8] buffer to a vector
+            let recv_vec = recv_buf.to_vec();
+            // create a `Message` type from buffer Vec<u8>
+            let msg = Message::Binary(recv_vec);
+
+            // write buffer to websocket
+            ws_write.send(msg).await?;
+        }
+    }
+}
+
+/// WebSocket read -> Tcp write
+async fn ws_to_tcp(
+    mut tcp_write: OwnedWriteHalf,
+    mut ws_read: SplitStream<WebSocketStream<TcpStream>>,
+) -> anyhow::Result<()> {
     while let Some(msg) = ws_read.next().await {
         // handle error in the message
         let msg = msg?;