Fix: keep alive

Author: kotiln

core/Cargo.toml

@@ -16,3 +16,4 @@ uuid = { version = "1.3.4", features = ["v4"] }
 quinn = "0.10.2"
 rustls = { version = "0.21.6", default-features = false, features = ["quic", "dangerous_configuration"] }
 rcgen = "0.11.1"
+macro-log = "0.2.0"

core/src/client.rs

@@ -1,6 +1,6 @@
-use std::{io::{Error, ErrorKind}, future::Future};
+use std::{io::{Error, ErrorKind}, future::Future, time::Duration};
 
-use tokio::{net::TcpStream, task::JoinHandle};
+use tokio::{net::TcpStream, task::JoinHandle, time::timeout};
 
 use crate::{
     cmd::{read_cmd, write_cmd, Command},
@@ -39,7 +39,11 @@ impl Client {
         R: Future<Output = ()>,
     {
         loop {
-            let cmd: Command = read_cmd(&mut self.stream, "").await;
+            let Ok(cmd) = timeout(Duration::from_millis(10000), read_cmd(&mut self.stream, "")).await
+                else {
+                    e!("连接不活跃！");
+                    break;
+                };
             let local_service = local_service.clone();
             wtf!(&cmd);
             match cmd {
@@ -97,8 +101,12 @@ impl Client {
                     e!("会话异常：{}", e);
                     break;
                 }
-                Command::Nothing => {
-                    let _ = write_cmd(&mut self.stream, Command::Nothing, "").await;
+                Command::KeepAlive => {
+                    let Ok(_) = write_cmd(&mut self.stream, Command::KeepAlive, "").await
+                        else {
+                            e!("连接中断！");
+                            break;
+                        };
                 },
                 _ => continue,
             };

core/src/cmd.rs

@@ -30,6 +30,7 @@ pub enum Command {
         nat_type: NatType,
         udp_addr: String, // Udp公网地址
     },
+    KeepAlive,
     Nothing,
     Success,
     Failure {
@@ -142,6 +143,7 @@ pub async fn read_cmd(tcp: &mut TcpStream, password: &str) -> Command {
         Some("failure") => Command::Failure {
             reason: cmds.into_iter().collect::<Vec<&str>>().join(" "),
         },
+        Some("keepalive") => Command::KeepAlive,
         _ => Command::Nothing,
     }
 }
@@ -168,6 +170,7 @@ pub async fn write_cmd(tcp: &mut TcpStream, cmd: Command, password: &str) -> std
         Command::Failure { reason } => {
             format!("failure {reason}")
         }
+        Command::KeepAlive => "keepalive".into(),
         _ => "".into(),
     };
     let len_password = password.len() as u8;

core/src/server.rs

@@ -3,7 +3,7 @@ use std::{borrow::BorrowMut, sync::Arc, collections::HashMap, io::ErrorKind, tim
 use once_cell::sync::Lazy;
 use tokio::{
     net::{TcpListener, TcpStream},
-    sync::Mutex, task::JoinHandle, time::sleep,
+    sync::Mutex, task::JoinHandle, time::{sleep, timeout},
 };
 
 use crate::{
@@ -138,20 +138,20 @@ impl Server {
                                         // Bug: 代理人已经死掉还能写入命令
                                         let mut agent = agent.lock().await;
                                         let agent = agent.borrow_mut();
-                                        match write_cmd(agent, Command::Nothing, "").await {
+                                        match write_cmd(agent, Command::KeepAlive, "").await {
                                             Err(_) => {
-                                                i!("PORT({port}) -> Agent {agent_addr} offline, release the port!");
+                                                e!("PORT({port}) -> Agent {agent_addr} offline, release the port!");
                                                 task.abort();
                                                 break;
                                             }
                                             _ => ()
                                         }
-                                        match read_cmd(agent, "").await {
-                                            Command::Nothing => {
+                                        match timeout(Duration::from_millis(2000), read_cmd(agent, "")).await {
+                                            Ok(Command::KeepAlive) => {
                                                 // i!("AGENT({agent_addr}) -> Living"); 
                                             }
                                             _ => {
-                                                i!("PORT({port}) -> Agent {agent_addr} no response, release the port!");
+                                                e!("PORT({port}) -> Agent {agent_addr} no response, release the port!");
                                                 task.abort();
                                                 break;
                                             }