From e32c869f00236131eb94a63523f0aa7a2c026bd5 Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Sun, 6 Jun 2021 03:33:45 +0800 Subject: [PATCH] set read timeout for half-open connections (one direction closed) - ref #490 --- Cargo.lock | 23 ++++-- crates/shadowsocks/Cargo.toml | 1 + .../shadowsocks/src/relay/tcprelay/utils.rs | 76 +++++++++++++++---- 3 files changed, 78 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bcbbe8bf..07553f3e 100644 --- a/Cargo.lock +++ b/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" diff --git a/crates/shadowsocks/Cargo.toml b/crates/shadowsocks/Cargo.toml index 034a750d..962feac1 100644 --- a/crates/shadowsocks/Cargo.toml +++ b/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 } diff --git a/crates/shadowsocks/src/relay/tcprelay/utils.rs b/crates/shadowsocks/src/relay/tcprelay/utils.rs index 8f66bcd7..1aeec939 100644 --- a/crates/shadowsocks/src/relay/tcprelay/utils.rs +++ b/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( 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> 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 { + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { // 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))), }