From b87703d2858187a3f0651dd731cacd6e8f79e107 Mon Sep 17 00:00:00 2001
From: alterdekim <alterwain@protonmail.com>
Date: Sat, 12 Oct 2024 17:10:44 +0300
Subject: [PATCH] 	modified:   src/client.rs

---
 src/client.rs | 77 ++++++++++++++++++++++++++++-----------------------
 1 file changed, 42 insertions(+), 35 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 471f0c8..4e236b3 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -119,43 +119,50 @@ pub async fn client_mode(client_config: ClientConfiguration, fd: i32, env: &JNIE
 
     let s_cipher = cipher_shared.clone();
     loop {
-        if close_token.is_cancelled() {
-            sock_read_task.abort();
-            dev_read_task.abort();
-            return;
-        }
-        if let Ok(bytes) = rx.try_recv() {
-            ltx.send("Write to tun.".as_bytes().to_vec());
-            if let Err(e) = dev_writer.write_all(&bytes).await {
-                ltx.send(format!("Writing error: {:?}", e).as_bytes().to_vec());
+        tokio::select! {
+            _ = close_token.cancelled() => {
+                sock_read_task.abort();
+                dev_read_task.abort();
+                break;
             }
-            if let Err(e) = dev_writer.flush().await {
-                ltx.send(format!("Flushing error: {:?}", e).as_bytes().to_vec());
-            }
-        }
-        if let Ok(bytes) = lrx.try_recv() {
-            FridaLib::traceFromNative(&env, String::from_utf8_lossy(&bytes).to_string());
-        }
-        if let Ok(bytes) = mx.try_recv() {
-            let s_c = s_cipher.lock().await;
-            
-            if s_c.is_some() {
-                let aes = Aes256Gcm::new(s_c.as_ref().unwrap().as_bytes().into());
-                let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
-                let ciphered_data = aes.encrypt(&nonce, &bytes[..]);
-                
-                if let Ok(ciphered_d) = ciphered_data {
-                    let vpn_packet = UDPVpnPacket{ data: ciphered_d, nonce: nonce.to_vec()};
-                    let serialized_data = vpn_packet.serialize();
-                    ltx.send("Write to socket".as_bytes().to_vec());
-                    sock_snd.send(&serialized_data).await.unwrap();
-                } else {
-                    ltx.send("Socket encryption failed.".as_bytes().to_vec());
+            rr = rx.recv() => {
+                if let Ok(bytes) = rr {
+                    ltx.send("Write to tun.".as_bytes().to_vec());
+                    if let Err(e) = dev_writer.write_all(&bytes).await {
+                        ltx.send(format!("Writing error: {:?}", e).as_bytes().to_vec());
+                    }
+                    if let Err(e) = dev_writer.flush().await {
+                        ltx.send(format!("Flushing error: {:?}", e).as_bytes().to_vec());
+                    }
                 }
-            } else {
-                ltx.send("There is no shared_secret in main loop".as_bytes().to_vec());
             }
-        }
-        //ltx.send("We've got there!".as_bytes().to_vec());
+            rr1 = lrx.recv() {
+                if let Ok(bytes) = rr1 {
+                    FridaLib::traceFromNative(&env, String::from_utf8_lossy(&bytes).to_string());
+                }
+            }
+            rr2 = mx.recv() {
+                if let Ok(bytes) = rr2 {
+                    let s_c = s_cipher.lock().await;
+                    
+                    if s_c.is_some() {
+                        let aes = Aes256Gcm::new(s_c.as_ref().unwrap().as_bytes().into());
+                        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
+                        let ciphered_data = aes.encrypt(&nonce, &bytes[..]);
+                        
+                        if let Ok(ciphered_d) = ciphered_data {
+                            let vpn_packet = UDPVpnPacket{ data: ciphered_d, nonce: nonce.to_vec()};
+                            let serialized_data = vpn_packet.serialize();
+                            ltx.send("Write to socket".as_bytes().to_vec());
+                            sock_snd.send(&serialized_data).await.unwrap();
+                        } else {
+                            ltx.send("Socket encryption failed.".as_bytes().to_vec());
+                        }
+                    } else {
+                        ltx.send("There is no shared_secret in main loop".as_bytes().to_vec());
+                    }
+                }
+            }
+        };
     }
 }
\ No newline at end of file