From f433df23fa62018b69415232864c669956ad58be Mon Sep 17 00:00:00 2001
From: alterdekim <alterwain@protonmail.com>
Date: Sun, 13 Oct 2024 21:06:56 +0300
Subject: [PATCH] 	modified:   src/client.rs

---
 src/client.rs | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 16f8b47..fa3fc09 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -27,7 +27,8 @@ use crate::jni::FridaLib;
 pub async fn client_mode(client_config: ClientConfiguration, fd: i32, close_token: CancellationToken) {
     info!("Starting client...");
 
-    let subtasks_quit: std::sync::Mutex<tokio_util::sync::CancellationToken> = std::sync::Mutex::new(tokio_util::sync::CancellationToken::new());
+    let dr_cancel: CancellationToken = CancellationToken::new();
+    let sr_cancel: CancellationToken = CancellationToken::new();
 
     let sock = UdpSocket::bind("0.0.0.0:25565").await.unwrap();
     sock.connect(&client_config.server.endpoint).await.unwrap();
@@ -46,13 +47,11 @@ pub async fn client_mode(client_config: ClientConfiguration, fd: i32, close_toke
 
     let cipher_shared: Arc<Mutex<Option<x25519_dalek::SharedSecret>>> = Arc::new(Mutex::new(None));
 
-    let cancel_dr = subtasks_quit.clone();
     let dev_read_task = tokio::spawn(async move {
         let mut buf = vec![0; 1400]; // mtu
         loop {
-            let sf = cancel_dr.lock().unwrap();
             tokio::select! {
-                _ = sf.cancelled() => {
+                _ = cancel_dr.cancelled() => {
                     info!("Cancellation token has been thrown dev_read_task");
                     return;
                 }
@@ -63,21 +62,18 @@ pub async fn client_mode(client_config: ClientConfiguration, fd: i32, close_toke
                     }
                 }
             };
-            drop(sf);
         }
     });
 
     let priv_key = BASE64_STANDARD.decode(client_config.client.private_key).unwrap();
     
     let cipher_shared_clone = cipher_shared.clone();
-    let cancel_sr = subtasks_quit.clone();
     let sock_read_task = tokio::spawn(async move {
         let mut buf = vec![0; 4096];
 
         loop {
-            let sf = cancel_sr.lock().unwrap();
             tokio::select! {
-                _ = sf.cancelled() => {
+                _ = sr_cancel.cancelled() => {
                     info!("Cancellation token has been thrown sock_read_task");
                     return;
                 }
@@ -124,7 +120,6 @@ pub async fn client_mode(client_config: ClientConfiguration, fd: i32, close_toke
                     }
                 }
             };
-            drop(sf);
         }
     });
 
@@ -142,9 +137,8 @@ pub async fn client_mode(client_config: ClientConfiguration, fd: i32, close_toke
         tokio::select! {
             _ = close_token.cancelled() => {
                 info!("Cancellation token has been thrown");
-                if let Ok(mut l) = subtasks_quit.lock() {
-                    l.cancel();
-                }
+                sr_cancel.cancel();
+                dr_cancel.cancel();
                 dev_read_task.await;
                 sock_read_task.await;
                 return;