set read timeout for half-open connections (one direction closed)

- ref #490

Cargo.lock

@@ -918,9 +918,9 @@ dependencies = [
 
 [[package]]
 name = "notify"
-version = "5.0.0-pre.9"
+version = "5.0.0-pre.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b89869d77edd64db917d7903abeadc166f93686b342c56cc0ca51acb68441d09"
+checksum = "51f18203a26893ca1d3526cf58084025d5639f91c44f8b70ab3b724f60e819a0"
 dependencies = [
  "bitflags",
  "crossbeam-channel",
@@ -1382,9 +1382,9 @@ dependencies = [
 
 [[package]]
 name = "security-framework"
-version = "2.2.0"
+version = "2.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3670b1d2fdf6084d192bc71ead7aabe6c06aa2ea3fbd9cc3ac111fa5c2b1bd84"
+checksum = "b239a3d5db51252f6f48f42172c65317f37202f4a21021bf5f9d40a408f4592c"
 dependencies = [
  "bitflags",
  "core-foundation",
@@ -1395,9 +1395,9 @@ dependencies = [
 
 [[package]]
 name = "security-framework-sys"
-version = "2.2.0"
+version = "2.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3676258fd3cfe2c9a0ec99ce3038798d847ce3e4bb17746373eb9f0f1ac16339"
+checksum = "7e4effb91b4b8b6fb7732e670b6cee160278ff8e6bf485c7805d9e319d76e284"
 dependencies = [
  "core-foundation-sys",
  "libc",
@@ -1508,6 +1508,7 @@ dependencies = [
  "spin 0.9.0",
  "thiserror",
  "tokio",
+ "tokio-io-timeout",
  "trust-dns-resolver",
  "url",
  "winapi",
@@ -1820,6 +1821,16 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "tokio-io-timeout"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "90c49f106be240de154571dd31fbe48acb10ba6c6dd6f6517ad603abffa42de9"
+dependencies = [
+ "pin-project-lite",
+ "tokio",
+]
+
 [[package]]
 name = "tokio-macros"
 version = "1.2.0"

crates/shadowsocks/Cargo.toml

@@ -56,6 +56,7 @@ async-trait = "0.1"
 mio = "0.7"
 socket2 = { version = "0.4", features = ["all"] }
 tokio = { version = "1.5", features = ["io-util", "macros", "net", "parking_lot", "process", "rt", "sync", "time"] }
+tokio-io-timeout = "1.1"
 
 trust-dns-resolver = { version = "0.20", optional = true }
 arc-swap = { version = "1.3", optional = true }

crates/shadowsocks/src/relay/tcprelay/utils.rs

@@ -8,10 +8,13 @@ use std::{
     io,
     pin::Pin,
     task::{Context, Poll},
+    time::Duration,
 };
 
 use futures::ready;
+use pin_project::pin_project;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio_io_timeout::TimeoutReader;
 
 use crate::crypto::v1::{CipherCategory, CipherKind};
 
@@ -175,9 +178,12 @@ enum TransferState {
     Done(u64),
 }
 
+#[pin_project(project = CopyBidirectionalProj)]
 struct CopyBidirectional<'a, A: ?Sized, B: ?Sized> {
-    a: &'a mut A,
-    b: &'a mut B,
+    #[pin]
+    a: TimeoutReader<&'a mut A>,
+    #[pin]
+    b: TimeoutReader<&'a mut B>,
     a_to_b: TransferState,
     b_to_a: TransferState,
 }
@@ -185,16 +191,13 @@ struct CopyBidirectional<'a, A: ?Sized, B: ?Sized> {
 fn transfer_one_direction<A, B>(
     cx: &mut Context<'_>,
     state: &mut TransferState,
-    r: &mut A,
-    w: &mut B,
+    mut r: Pin<&mut TimeoutReader<&mut A>>,
+    mut w: Pin<&mut TimeoutReader<&mut B>>,
 ) -> Poll<io::Result<u64>>
 where
     A: AsyncRead + AsyncWrite + Unpin + ?Sized,
     B: AsyncRead + AsyncWrite + Unpin + ?Sized,
 {
-    let mut r = Pin::new(r);
-    let mut w = Pin::new(w);
-
     loop {
         match state {
             TransferState::Running(buf) => {
@@ -218,19 +221,60 @@ where
 {
     type Output = io::Result<(u64, u64)>;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         // Unpack self into mut refs to each field to avoid borrow check issues.
-        let CopyBidirectional { a, b, a_to_b, b_to_a } = &mut *self;
+        let CopyBidirectionalProj {
+            mut a,
+            mut b,
+            a_to_b,
+            b_to_a,
+        } = self.project();
 
-        let a_to_b = transfer_one_direction(cx, a_to_b, &mut *a, &mut *b)?;
-        let b_to_a = transfer_one_direction(cx, b_to_a, &mut *b, &mut *a)?;
+        let poll_a_to_b = transfer_one_direction(cx, a_to_b, a.as_mut(), b.as_mut())?;
+        let poll_b_to_a = transfer_one_direction(cx, b_to_a, b.as_mut(), a.as_mut())?;
 
         // It is not a problem if ready! returns early because transfer_one_direction for the
         // other direction will keep returning TransferState::Done(count) in future calls to poll
-        let a_to_b = ready!(a_to_b);
-        let b_to_a = ready!(b_to_a);
 
-        Poll::Ready(Ok((a_to_b, b_to_a)))
+        // When one direction have already finished, FIN have already sent to the writer.
+        // But the FIN may have been lost.
+        // This timer is for purging those "Half Open" connections.
+        const READ_TIMEOUT_WHEN_ONE_SHUTDOWN: Duration = Duration::from_secs(5);
+
+        // Check if one end have already finished
+        match (poll_a_to_b, poll_b_to_a) {
+            (Poll::Ready(a_to_b), Poll::Ready(b_to_a)) => Poll::Ready(Ok((a_to_b, b_to_a))),
+
+            (Poll::Ready(a_to_b), Poll::Pending) => {
+                // a -> b finished, then FIN have already sent to b, setting a read timeout on b
+                if let None = b.timeout() {
+                    b.as_mut().set_timeout_pinned(Some(READ_TIMEOUT_WHEN_ONE_SHUTDOWN));
+
+                    // poll again to ensure Waker have already registered to the timer
+                    let b_to_a = ready!(transfer_one_direction(cx, b_to_a, b.as_mut(), a.as_mut())?);
+
+                    Poll::Ready(Ok((a_to_b, b_to_a)))
+                } else {
+                    Poll::Pending
+                }
+            }
+
+            (Poll::Pending, Poll::Ready(b_to_a)) => {
+                // b -> a finished, then FIN have already sent to a, setting a read timeout on a
+                if let None = a.timeout() {
+                    a.as_mut().set_timeout_pinned(Some(READ_TIMEOUT_WHEN_ONE_SHUTDOWN));
+
+                    // poll again to ensure Waker have already registered to the timer
+                    let a_to_b = ready!(transfer_one_direction(cx, a_to_b, a.as_mut(), b.as_mut())?);
+
+                    Poll::Ready(Ok((a_to_b, b_to_a)))
+                } else {
+                    Poll::Pending
+                }
+            }
+
+            (Poll::Pending, Poll::Pending) => Poll::Pending,
+        }
     }
 }
 
@@ -271,8 +315,8 @@ where
     P: AsyncRead + AsyncWrite + Unpin + ?Sized,
 {
     CopyBidirectional {
-        a: encrypted,
-        b: plain,
+        a: TimeoutReader::new(encrypted),
+        b: TimeoutReader::new(plain),
         a_to_b: TransferState::Running(CopyBuffer::new(encrypted_read_buffer_size(method))),
         b_to_a: TransferState::Running(CopyBuffer::new(plain_read_buffer_size(method))),
     }