From f99f6b4629e845f97068d9e43bac4de01cc1718d Mon Sep 17 00:00:00 2001
From: alterdekim <alterwain@protonmail.com>
Date: Sun, 11 Aug 2024 02:51:06 +0300
Subject: [PATCH]  Changes to be committed: 	modified:   src/server.rs

---
 src/server.rs | 134 ++++++++++++++++++++++++++------------------------
 1 file changed, 69 insertions(+), 65 deletions(-)

diff --git a/src/server.rs b/src/server.rs
index a08cd6e..49fc182 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,5 +1,6 @@
 use tokio::{net::UdpSocket, sync::{mpsc, Mutex}};
-use std::{borrow::{Borrow, BorrowMut}, future::IntoFuture, io::{self, Read, Write}, net::{SocketAddr, Ipv4Addr, IpAddr}, sync::Arc, thread, time};
+use tokio::task::JoinSet;
+use std::{borrow::{Borrow, BorrowMut}, future::IntoFuture, io::{self, Read, Write}, net::{SocketAddr, Ipv4Addr, IpAddr}, sync::{Arc}, thread, time};
 use std::process::Command;
 use clap::{App, Arg};
 use env_logger::Builder;
@@ -9,7 +10,6 @@ use serde_derive::Serialize;
 use serde_derive::Deserialize;
 //use packet::{builder::Builder, icmp, ip, Packet};
 use std::collections::HashMap;
-use futures::future;
 
 pub async fn server_mode() {
     info!("Starting server...");
@@ -26,10 +26,10 @@ pub async fn server_mode() {
 
     let tun_device = Arc::new(Mutex::new(tun::create(&config).unwrap()));
 
-    let sock = Arc::new(match UdpSocket::bind("192.168.0.5:8879".parse::<SocketAddr>().unwrap()).await {
+    let sock = Arc::new(Mutex::new(match UdpSocket::bind("192.168.0.5:8879".parse::<SocketAddr>().unwrap()).await {
         Ok(s) => s,
         Err(_error) => panic!("Cannot bind to address")
-    });
+    }));
 
     let clients = Arc::new(Mutex::new(HashMap::new()));
 
@@ -60,71 +60,75 @@ pub async fn server_mode() {
     let sock_clone = sock.clone();
     let clients_clone = clients.clone();
     let tun_device_clone_second = tun_device.clone();
+    
+    let mut set = JoinSet::new();
 
-    let tasks = vec![
-        tokio::spawn(async move {
-            let mut buf = [0; 1024];
-            let sock_main_instance = sock_main;
-            loop {
-                let mut tun = tun_device_clone_second.lock().await;
-                let clients_main_instance = clients_main.lock().await;
-                let len = match tun.read(&mut buf) {
-                    Ok(l) => l,
-                    Err(error) => {
-                        error!("Problem with reading from tun: {error:?}");
-                        0
-                    },
-                };
-                
-                if len <= 0 { continue; }
-        
-                info!("{:?} bytes received from tun", len);
+    set.spawn(async move {
+        let mut buf = [0; 1024];
+        loop {
+            let mut tun = tun_device_clone_second.lock().await;
+            let len = match tun.read(&mut buf) {
+                Ok(l) => l,
+                Err(error) => {
+                    error!("Problem with reading from tun: {error:?}");
+                    0
+                },
+            };
             
-                match clients_main_instance.get(&"10.8.0.2") {
-                    Some(&addr) => {
-                        let len = match sock_main_instance.send_to(&buf, addr).await {
-                            Ok(l) => l,
-                            Err(error) => {error!("Problem with writing to tun: {error:?}");
-                            0},
-                        };
-                        info!("{:?} bytes sent to socket", len);
-                    },
-                    None => error!("There is no client..."),
+            if len <= 0 { continue; }
+    
+            info!("{:?} bytes received from tun", len);
+            let sock_main_instance = sock_main.lock().await;
+            let clients_main_instance = clients_main.lock().await;
+            match clients_main_instance.get(&"10.8.0.2") {
+                Some(&addr) => {
+                    sock_main_instance.send_to(&buf, addr);
+                    info!("bytes sent to socket");
+                },
+                None => error!("There is no client..."),
+            }
+        }
+    });
+
+    set.spawn(async move {
+        let mut buf = [0; 1024];
+        loop {
+            let sock = sock_clone.lock().await;
+            let (len, addr) = match sock.recv_from(&mut buf).await {
+                Err(error) => {
+                    error!("Problem with reading from socket: {error:?}");
+                    (0, SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0))
+                },
+                Ok(l) => l,
+            };
+
+            if len <= 0 { continue; }
+            
+            let mut tun = tun_device_clone.lock().await;
+            let mut clients = clients_clone.lock().await;
+
+            clients.insert("10.8.0.2", addr);
+            info!("{:?} bytes received from {:?}", len, addr);
+            
+            let len = match tun.write(&buf) {
+                Ok(l) => l,
+                Err(error) => {
+                    error!("Problem with writing to tun: {error:?}");
+                    0
                 }
-            }
-        }),
+            };
+
+            info!("{:?} bytes sent to tun", len);
+        }
+    });
+
+   /* let tasks = vec![
+        tokio::spawn(),
         
-        tokio::spawn(async move {
-            let mut buf = [0; 1024];
-            let sock = sock_clone;
-            loop {
-                let mut tun = tun_device_clone.lock().await;
-                let mut clients = clients_clone.lock().await;
-                let (len, addr) = match sock.recv_from(&mut buf).await {
-                    Err(error) => {
-                        error!("Problem with reading from socket: {error:?}");
-                        (0, SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0))
-                    },
-                    Ok(l) => l,
-                };
-
-                if len <= 0 { continue; }
-                
-                clients.insert("10.8.0.2", addr);
-                info!("{:?} bytes received from {:?}", len, addr);
-                
-                let len = match tun.write(&buf) {
-                    Ok(l) => l,
-                    Err(error) => {
-                        error!("Problem with writing to tun: {error:?}");
-                        0
-                    }
-                };
-
-                info!("{:?} bytes sent to tun", len);
-            }
-        })
+        tokio::spawn()
     ];
 
-    futures::future::join_all(tasks).await;
+    futures::future::join_all(tasks).await;*/
+
+    while let Some(res) = set.join_next().await {}
 }
\ No newline at end of file