feat: add tcp_timeout

Cargo.toml

@@ -36,3 +36,4 @@ hex = "0.4.3"
 redis = { version = "0.20.0", features = ["tokio-comp"] }
 tokio-openssl = "0.6.1"
 openssl = { version = "0.10.34", features = ["vendored"] }
+pin-project-lite = "0.2.6"

src/relay.rs

@@ -3,6 +3,7 @@ use crate::common::UdpStream;
 use crate::inbound::{Inbound, InboundAccept, InboundRequest};
 use crate::outbound::Outbound;
 use crate::utils::count_stream::CountStream;
+use crate::utils::timeout_stream::TimeoutStream;
 use anyhow::{anyhow, bail, Context, Error, Result};
 use futures::{SinkExt, StreamExt, TryStreamExt};
 use log::error;
@@ -20,6 +21,7 @@ pub struct Relay<I, O> {
     inbound: Arc<I>,
     outbound: Arc<O>,
     tcp_nodelay: bool,
+    pub tcp_timeout: Option<Duration>,
 }
 
 impl<I, O> Relay<I, O>
@@ -78,6 +80,7 @@ where
             inbound: Arc::new(inbound),
             outbound: Arc::new(outbound),
             tcp_nodelay,
+            tcp_timeout: None,
         }
     }
     pub async fn serve_trojan(&self, auth: Arc<dyn Auth>) -> Result<()> {
@@ -91,6 +94,7 @@ where
                 .context("Set TCP_NODELAY failed")?;
 
             let (stream, sender) = CountStream::new2(stream);
+            let stream = TimeoutStream::new(stream, self.tcp_timeout);
 
             let inbound = self.inbound.clone();
             let outbound = self.outbound.clone();

src/server.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use crate::{
     auth::{Auth, AuthHub},
@@ -23,7 +23,8 @@ pub async fn start(config: Config) -> Result<()> {
     let inbound = TrojanInbound::new(auth_hub.clone(), tls_context, config.trojan).await?;
     let outbound = DirectOutbound::new();
 
-    let relay = Relay::new(listener, inbound, outbound, config.tls.tcp_nodelay);
+    let mut relay = Relay::new(listener, inbound, outbound, config.tls.tcp_nodelay);
+    relay.tcp_timeout = Some(Duration::from_secs(600));
 
     info!("Service started.");
 

src/utils/mod.rs

@@ -2,3 +2,4 @@ pub mod config;
 pub mod count_stream;
 pub mod logger;
 pub mod peekable_stream;
+pub mod timeout_stream;

src/utils/timeout_stream.rs

@@ -0,0 +1,92 @@
+use futures::ready;
+use pin_project_lite::pin_project;
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+use tokio::{
+    io::{self, AsyncRead, AsyncWrite},
+    time::{sleep, Duration, Instant, Sleep},
+};
+
+pin_project! {
+    #[derive(Debug)]
+    pub struct TimeoutStream<S> {
+        #[pin]
+        s: S,
+        duration: Option<Duration>,
+        sleep: Pin<Box<Sleep>>,
+    }
+}
+
+impl<S> TimeoutStream<S> {
+    pub fn new(s: S, duration: Option<Duration>) -> TimeoutStream<S>
+    where
+        S: AsyncRead + AsyncWrite,
+    {
+        TimeoutStream {
+            s,
+            duration,
+            sleep: Box::pin(sleep(duration.unwrap_or(Duration::from_secs(1)))),
+        }
+    }
+}
+
+impl<S> AsyncRead for TimeoutStream<S>
+where
+    S: AsyncRead + Unpin,
+{
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut io::ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let this = self.project();
+        match this.duration {
+            Some(duration) => match this.s.poll_read(cx, buf) {
+                Poll::Ready(r) => {
+                    this.sleep.as_mut().reset(Instant::now() + *duration);
+                    Poll::Ready(r)
+                }
+                Poll::Pending => {
+                    ready!(this.sleep.as_mut().poll(cx));
+                    Poll::Ready(Err(io::ErrorKind::TimedOut.into()))
+                }
+            },
+            None => this.s.poll_read(cx, buf),
+        }
+    }
+}
+
+impl<S> AsyncWrite for TimeoutStream<S>
+where
+    S: AsyncWrite + Unpin,
+{
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, io::Error>> {
+        let this = self.project();
+        match this.duration {
+            Some(duration) => match this.s.poll_write(cx, buf) {
+                Poll::Ready(r) => {
+                    this.sleep.as_mut().reset(Instant::now() + *duration);
+                    Poll::Ready(r)
+                }
+                Poll::Pending => {
+                    ready!(this.sleep.as_mut().poll(cx));
+                    Poll::Ready(Err(io::ErrorKind::TimedOut.into()))
+                }
+            },
+            None => this.s.poll_write(cx, buf),
+        }
+    }
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        self.project().s.poll_flush(cx)
+    }
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        self.project().s.poll_shutdown(cx)
+    }
+}