modified: src/client.rs

src/client.rs

@@ -1,6 +1,6 @@
 //use crossbeam_channel::unbounded;
 use socket2::SockAddr;
-use tokio::{net::UdpSocket, sync::{Mutex, mpsc}, io::{BufReader, BufWriter, AsyncWriteExt, AsyncReadExt}, fs::File};
+use tokio::{net::UdpSocket, sync::{Mutex, mpsc, oneshot}, io::{BufReader, BufWriter, AsyncWriteExt, AsyncReadExt}, fs::File};
 use tokio_util::sync::CancellationToken;
 use base64::prelude::*;
 use log::{error, info, warn};
@@ -44,24 +44,38 @@ 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, mut listen_dr) = oneshot::channel();
     let dev_read_task = tokio::spawn(async move {
         let mut buf = vec![0; 1400]; // mtu
         loop {
-            if let Ok(n) = dev_reader.read(&mut buf).await {
+            tokio::select! {
+                _ = listen_dr => {
+                    return;
+                }
+                rr = dev_reader.read(&mut buf) {
+                    if let Ok(n) = rr {
                         info!("Read from tun."); // hex::encode(&buf[..n])
                         dx.send(buf[..n].to_vec()).unwrap();
                     }
                 }
+            };
+        }
     });
 
     let priv_key = BASE64_STANDARD.decode(client_config.client.private_key).unwrap();
     
     let cipher_shared_clone = cipher_shared.clone();
+    let (cancel_sr, mut listen_sr) = oneshot::channel();
     let sock_read_task = tokio::spawn(async move {
         let mut buf = vec![0; 4096];
 
         loop {
-            if let Ok(l) = sock_rec.recv(&mut buf).await {
+            tokio::select! {
+                _ = listen_sr => {
+                    return;
+                }
+                rr = sock_rec.recv(&mut buf) => {
+                    if let Ok(l) = rr {
                         info!("Read from socket");
                         let mut s_cipher = cipher_shared_clone.lock().await;
                         match buf.first() {
@@ -102,6 +116,8 @@ pub async fn client_mode(client_config: ClientConfiguration, fd: i32, close_toke
                         drop(s_cipher);
                     }
                 }
+            };
+        }
     });
 
     let pkey = BASE64_STANDARD.decode(client_config.client.public_key).unwrap();
@@ -118,8 +134,8 @@ pub async fn client_mode(client_config: ClientConfiguration, fd: i32, close_toke
         tokio::select! {
             _ = close_token.cancelled() => {
                 info!("Cancellation token has been thrown");
-                sock_read_task.abort();
-                dev_read_task.abort();
+                cancel_sr.send(());
+                cancel_dr.send(());
                 break;
             }
             rr = rx.recv() => {