From 5a4c0b1b7d880108fc050a9079c4816cfad71a00 Mon Sep 17 00:00:00 2001
From: yinheli <me@yinheli.com>
Date: Wed, 31 Aug 2022 19:55:33 +0800
Subject: [PATCH] feat: fast check before relay

---
 src/gateway/relay_tcp.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/gateway/relay_tcp.rs b/src/gateway/relay_tcp.rs
index 92a6c39..ee938f3 100644
--- a/src/gateway/relay_tcp.rs
+++ b/src/gateway/relay_tcp.rs
@@ -1,6 +1,7 @@
 use std::{
     net::IpAddr,
     sync::{Arc, Mutex},
+    time::Duration,
 };
 
 use log::warn;
@@ -14,6 +15,7 @@ use tokio::{
         tcp::{ReadHalf, WriteHalf},
         TcpListener,
     },
+    time,
 };
 
 use crate::{
@@ -184,6 +186,8 @@ fn find_target(setting: ArcSetting, session: Session) -> Option<(String, String,
 }
 
 async fn copy<'a>(r: &mut ReadHalf<'a>, w: &mut WriteHalf<'a>) -> Result<u64, io::Error> {
+    time::timeout(Duration::from_secs(1), r.readable()).await??;
+    time::timeout(Duration::from_secs(1), w.writable()).await??;
     let n = io::copy(r, w).await?;
     w.shutdown().await?;
     Ok(n)