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,63 +44,79 @@ 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! {
-                info!("Read from tun."); // hex::encode(&buf[..n])
+                _ = listen_dr => {
-                dx.send(buf[..n].to_vec()).unwrap();
+                    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! {
-                info!("Read from socket");
+                _ = listen_sr => {
-                let mut s_cipher = cipher_shared_clone.lock().await;
+                    return;
-                match buf.first() {
-                    Some(h) => {
-                        match h {
-                            0 => {
-                                let handshake = UDPVpnHandshake::deserialize(&(buf[..l].to_vec()));
-                                let mut k = [0u8; 32];
-                                for (&x, p) in handshake.public_key.iter().zip(k.iter_mut()) {
-                                    *p = x;
-                                }
-                                let mut k1 = [0u8; 32];
-                                for (&x, p) in priv_key.iter().zip(k1.iter_mut()) {
-                                    *p = x;
-                                }
-                                *s_cipher = Some(StaticSecret::from(k1)
-                                    .diffie_hellman(&PublicKey::from(k)));
-                            }, // handshake
-                            1 => {
-                                let wrapped_packet = UDPVpnPacket::deserialize(&(buf[..l].to_vec()));
-                                if s_cipher.is_some() {
-                                    let aes = Aes256Gcm::new(s_cipher.as_ref().unwrap().as_bytes().into());
-                                    let nonce = Nonce::clone_from_slice(&wrapped_packet.nonce);
-                                    match aes.decrypt(&nonce, &wrapped_packet.data[..]) {
-                                        Ok(decrypted) => { let _ = tx.send(decrypted); },
-                                        Err(error) => { error!("Decryption error! {:?}", error); }
-                                    }
-                                } else {
-                                    warn!("There is no static_secret");
-                                }
-                            }, // payload
-                            2 => { info!("Got keepalive packet"); },
-                            _ => { error!("Unexpected header value."); }
-                        }
-                    },
-                    None => { error!("There is no header."); }
                 }
-                drop(s_cipher);
+                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() {
+                            Some(h) => {
+                                match h {
+                                    0 => {
+                                        let handshake = UDPVpnHandshake::deserialize(&(buf[..l].to_vec()));
+                                        let mut k = [0u8; 32];
+                                        for (&x, p) in handshake.public_key.iter().zip(k.iter_mut()) {
+                                            *p = x;
+                                        }
+                                        let mut k1 = [0u8; 32];
+                                        for (&x, p) in priv_key.iter().zip(k1.iter_mut()) {
+                                            *p = x;
+                                        }
+                                        *s_cipher = Some(StaticSecret::from(k1)
+                                            .diffie_hellman(&PublicKey::from(k)));
+                                    }, // handshake
+                                    1 => {
+                                        let wrapped_packet = UDPVpnPacket::deserialize(&(buf[..l].to_vec()));
+                                        if s_cipher.is_some() {
+                                            let aes = Aes256Gcm::new(s_cipher.as_ref().unwrap().as_bytes().into());
+                                            let nonce = Nonce::clone_from_slice(&wrapped_packet.nonce);
+                                            match aes.decrypt(&nonce, &wrapped_packet.data[..]) {
+                                                Ok(decrypted) => { let _ = tx.send(decrypted); },
+                                                Err(error) => { error!("Decryption error! {:?}", error); }
+                                            }
+                                        } else {
+                                            warn!("There is no static_secret");
+                                        }
+                                    }, // payload
+                                    2 => { info!("Got keepalive packet"); },
+                                    _ => { error!("Unexpected header value."); }
+                                }
+                            },
+                            None => { error!("There is no header."); }
+                        }
+                        drop(s_cipher);
+                    }
+                }
+            };
         }
     });
 
@@ -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();
+                cancel_sr.send(());
-                dev_read_task.abort();
+                cancel_dr.send(());
                 break;
             }
             rr = rx.recv() => {