From a698bbb14f9e44ab393746cbe800e300f8086a91 Mon Sep 17 00:00:00 2001
From: groche97 <groche97@gmail.com>
Date: Mon, 26 Sep 2022 23:03:16 +0200
Subject: [PATCH] improve scalability and security

---
 src/main.rs         | 49 +++++++++++++++++++++++++++------------------
 src/protocol/mod.rs |  5 +++++
 2 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 1fccffb..0dede85 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,5 @@
 use std::net::{TcpListener, TcpStream};
+use std::sync::{Arc, RwLock};
 use std::io::prelude::*;
 use crate::client::guard;
 use std::thread;
@@ -10,24 +11,16 @@ mod conf;
 mod protocol;
 
 fn main() {
-    let listener = TcpListener::bind("127.0.0.1:25567").unwrap();
-    let mut buf: [u8; 256] = [1; 256];
-    let servers = conf::Servers::new();
-    let mut guard = guard::Guard::new();
+    let listener = TcpListener::bind("0.0.0.0:25565").unwrap();
+    let servers = Arc::new(RwLock::new(conf::Servers::new()));
+    let guard = Arc::new(RwLock::new(guard::Guard::new()));
     for stream in listener.incoming() {
-        if guard.can_add(){
+        if guard.read().unwrap().can_add(){
             match stream {
-                Ok(mut stream) => {
-                    stream.set_read_timeout(Some(Duration::from_millis(5000)));
-                    //stream.set_write_timeout(Some(Duration::from_millis(5000)));
-                    let leng = match stream.read(&mut buf) {
-                        Ok(l) => l,
-                        Err(_e) => break,
-                    };
-                    let mut hs = protocol::HandShake::new(&mut buf[.. leng]);
-                    if hs.get_raw()[0] < 200 { //Filtra los ping, solo controlamos los handshakes
-                        conect_server(&servers, hs, stream, &mut guard);
-                    }
+                Ok(stream) => {
+                    let g = guard.clone();
+                    let s = servers.clone();
+                    thread::spawn(|| read_connection(stream, s , g));
                 },
 
                 Err(_e) => println!("{}",_e),
@@ -36,18 +29,34 @@ fn main() {
     }
 }
 
-fn conect_server(servers: &conf::Servers,
+fn read_connection(mut stream: TcpStream,
+        servers: Arc<RwLock<conf::Servers>>,
+        guard: Arc<RwLock<guard::Guard>> ) {
+    let mut buf: [u8; 256] = [1; 256];
+    stream.set_read_timeout(Some(Duration::from_millis(5000)));
+    let leng = match stream.read(&mut buf) {
+        Ok(l) => l,
+        Err(_e) => return,
+    };
+    let hs = protocol::HandShake::new(&mut buf[.. leng]);
+    if hs.is_handshake() { //Filtra los ping, solo controlamos los handshakes
+            conect_server(servers, hs, stream, guard);
+    }
+
+}
+
+fn conect_server(servers: Arc<RwLock<conf::Servers>>,
         mut hs: protocol::HandShake,
         stream: TcpStream,
-        guard: &mut guard::Guard){
+        guard: Arc<RwLock<guard::Guard>>){
 
-    match servers.get_server(&hs.get_host_name()) {
+    match servers.read().unwrap().get_server(&hs.get_host_name()) {
         Some(s) => {
             hs.replace_port(s.1);
             let mut sstream = TcpStream::connect(s.0 + ":" + &s.1.to_string()).unwrap();
             sstream.write(hs.get_raw());
             let c1 = client::Client::new(stream,sstream, hs);
-            guard.add_thread(c1.start_proxy());
+            guard.write().unwrap().add_thread(c1.start_proxy());
         },
         None => println!("No server found for {}", hs.get_host_name())
     }
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 2de975f..7d7947d 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -36,5 +36,10 @@ impl<'a> HandShake<'a>{
     pub fn get_raw(&self) -> &[u8]{
         self.datagram
     }
+
+    pub fn is_handshake(&self) -> bool {
+        (self.datagram[0] == self.len_pack) &&
+            (self.len_dom + 7 == self.len_pack)
+    }
 }