diff --git a/Cargo.lock b/Cargo.lock index 79eac7c..6f8eba3 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 944d669..c192279 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/main.rs b/src/main.rs index cbe29e2..f6d5581 100644 --- a/src/main.rs +++ b/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( - vec![ - TermLogger::new(LevelFilter::Trace, Config::default(), TerminalMode::Mixed, ColorChoice::Auto), - WriteLogger::new(LevelFilter::Debug, Config::default(), File::create("server.log").unwrap()), - ] - )?; + CombinedLogger::init(vec![ + TermLogger::new( + LevelFilter::Trace, + 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 diff --git a/src/tcp/handle_connection.rs b/src/tcp/handle_connection.rs index 4125cba..b7f5b79 100644 --- a/src/tcp/handle_connection.rs +++ b/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; diff --git a/src/tcp/handle_websocket.rs b/src/tcp/handle_websocket.rs index 8c2674a..2e6898e 100644 --- a/src/tcp/handle_websocket.rs +++ b/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 (mut ws_write, mut ws_read) = ws_stream.split(); + let (tcp_read, tcp_write) = tcp_stream.into_split(); + let (ws_write, ws_read) = ws_stream.split(); // tcp read -> ws write - tokio::spawn(async move { - // 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 - let msg = Message::Binary(recv_vec); - - // write buffer to websocket - ws_write.send(msg).await.unwrap(); - } - } - }); + tokio::spawn(tcp_to_ws(tcp_read, ws_write)); // 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, 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 + 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>, +) -> anyhow::Result<()> { while let Some(msg) = ws_read.next().await { // handle error in the message let msg = msg?;