Removed Connection, uses try_timeout instead. Revert back to use 8K buffer

src/relay/dns_resolver/trust_dns_resolver.rs

@@ -1,9 +1,7 @@
 //! Asynchronous DNS resolver
 
-use std::{
-    io::{self, ErrorKind},
-    net::SocketAddr,
-};
+use std::io::{self, ErrorKind};
+use std::net::SocketAddr;
 
 use log::{debug, error, trace};
 use tokio;

src/relay/loadbalancing/server/mod.rs

@@ -2,7 +2,8 @@
 
 use std::sync::Arc;
 
-pub use self::{ping::PingBalancer, roundrobin::RoundRobin};
+pub use self::ping::PingBalancer;
+pub use self::roundrobin::RoundRobin;
 
 use crate::config::ServerConfig;
 

src/relay/loadbalancing/server/ping.rs

@@ -1,24 +1,19 @@
-use std::{
-    collections::VecDeque,
-    fmt, io,
-    sync::{
-        atomic::{AtomicU64, AtomicUsize, Ordering},
-        Arc, Mutex,
-    },
-};
+use std::collections::VecDeque;
+use std::fmt;
+use std::io;
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
 
-use crate::{
-    config::ServerConfig,
-    context::SharedContext,
-    relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
-};
+use crate::config::ServerConfig;
+use crate::context::SharedContext;
+use crate::relay::loadbalancing::server::LoadBalancer;
+use crate::relay::socks5::Address;
+use crate::relay::tcprelay::client::ServerClient;
 
 use log::{debug, error, info};
-use tokio::{
-    self,
-    io::{AsyncReadExt, AsyncWriteExt},
-    time::{self, Duration, Instant},
-};
+use tokio;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::time::{self, Duration, Instant};
 
 struct Server {
     config: Arc<ServerConfig>,

src/relay/loadbalancing/server/roundrobin.rs

@@ -2,10 +2,8 @@
 
 use std::sync::Arc;
 
-use crate::{
-    config::{Config, ServerConfig},
-    relay::loadbalancing::server::LoadBalancer,
-};
+use crate::config::{Config, ServerConfig};
+use crate::relay::loadbalancing::server::LoadBalancer;
 
 #[derive(Clone)]
 pub struct RoundRobin {

src/relay/tcprelay/client.rs

@@ -9,13 +9,13 @@ use std::{
 };
 
 use log::trace;
-use tokio::{net::TcpStream, prelude::*};
+use tokio::net::TcpStream;
+use tokio::prelude::*;
 
+use super::CryptoStream;
 use crate::relay::socks5::{
     self, Address, Command, HandshakeRequest, HandshakeResponse, Reply, TcpRequestHeader, TcpResponseHeader,
 };
-
-use super::{CryptoStream, STcpStream};
 use crate::{config::ServerConfig, context::SharedContext};
 
 /// Socks5 proxy client
@@ -124,7 +124,7 @@ impl AsyncWrite for Socks5Client {
 }
 
 pub(crate) struct ServerClient {
-    pub stream: CryptoStream<STcpStream>,
+    pub stream: CryptoStream<TcpStream>,
 }
 
 impl ServerClient {

src/relay/tcprelay/mod.rs

@@ -4,15 +4,7 @@
 // Maybe removed in the future
 #![allow(clippy::unnecessary_mut_passed)]
 
-use std::{
-    io,
-    marker::Unpin,
-    ops::{Deref, DerefMut},
-    pin::Pin,
-    sync::Arc,
-    task::{Context, Poll},
-    time::Duration,
-};
+use std::{io, marker::Unpin, sync::Arc};
 
 use crate::{
     config::{ConfigType, ServerAddr, ServerConfig},
@@ -21,17 +13,11 @@ use crate::{
 };
 
 use bytes::BytesMut;
-use futures::{
-    future::{self, FusedFuture, Pending},
-    ready, select, Future,
-};
+use futures::future::FusedFuture;
+use futures::{select, Future};
 use log::trace;
-use tokio::{
-    io::{ReadHalf, WriteHalf},
-    net::TcpStream,
-    prelude::*,
-    time::{self, Timeout},
-};
+use tokio::net::TcpStream;
+use tokio::prelude::*;
 
 mod aead;
 pub mod client;
@@ -47,145 +33,10 @@ mod utils;
 
 pub use self::crypto_io::CryptoStream;
 
-const BUFFER_SIZE: usize = 32 * 1024; // 32K buffer
-
-pub struct Connection<S> {
-    stream: S,
-    read_timer: Option<Timeout<Pending<()>>>,
-    write_timer: Option<Timeout<Pending<()>>>,
-    timeout: Option<Duration>,
-}
-
-impl<S> Connection<S> {
-    pub fn new(stream: S, timeout: Option<Duration>) -> Connection<S> {
-        Connection {
-            stream,
-            read_timer: None,
-            write_timer: None,
-            timeout,
-        }
-    }
-
-    fn poll_read_timeout(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        loop {
-            if let Some(ref mut timer) = self.read_timer {
-                ready!(Pin::new(timer).poll(cx))?;
-            } else {
-                match self.timeout {
-                    Some(timeout) => self.read_timer = Some(time::timeout(timeout, future::pending())),
-                    None => break,
-                }
-            }
-        }
-        Poll::Ready(Ok(()))
-    }
-
-    fn poll_write_timeout(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        loop {
-            if let Some(ref mut timer) = self.write_timer {
-                ready!(Pin::new(timer).poll(cx))?;
-            } else {
-                match self.timeout {
-                    Some(timeout) => self.write_timer = Some(time::timeout(timeout, future::pending())),
-                    None => break,
-                }
-            }
-        }
-        Poll::Ready(Ok(()))
-    }
-
-    fn cancel_read_timeout(&mut self) {
-        let _ = self.read_timer.take();
-    }
-
-    fn cancel_write_timeout(&mut self) {
-        let _ = self.write_timer.take();
-    }
-}
-
-impl<S> Connection<S>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-{
-    pub fn split(self) -> (ReadHalf<Connection<S>>, WriteHalf<Connection<S>>) {
-        use tokio::io::split;
-        split(self)
-    }
-}
-
-impl<S> Deref for Connection<S> {
-    type Target = S;
-
-    fn deref(&self) -> &Self::Target {
-        &self.stream
-    }
-}
-
-impl<S> DerefMut for Connection<S> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.stream
-    }
-}
-
-impl<S: Unpin> Unpin for Connection<S> {}
-
-impl<S> AsyncRead for Connection<S>
-where
-    S: AsyncRead + Unpin,
-{
-    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
-        match Pin::new(&mut self.stream).poll_read(cx, buf) {
-            Poll::Ready(r) => {
-                self.cancel_read_timeout();
-                Poll::Ready(r)
-            }
-            Poll::Pending => {
-                ready!(self.poll_read_timeout(cx))?;
-                Poll::Pending
-            }
-        }
-    }
-}
-
-impl<S> AsyncWrite for Connection<S>
-where
-    S: AsyncWrite + Unpin,
-{
-    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
-        match Pin::new(&mut self.stream).poll_write(cx, buf) {
-            Poll::Ready(r) => {
-                self.cancel_write_timeout();
-                Poll::Ready(r)
-            }
-            Poll::Pending => {
-                ready!(self.poll_write_timeout(cx))?;
-                Poll::Pending
-            }
-        }
-    }
-
-    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        match Pin::new(&mut self.stream).poll_flush(cx) {
-            Poll::Ready(r) => {
-                self.cancel_write_timeout();
-                Poll::Ready(r)
-            }
-            Poll::Pending => {
-                ready!(self.poll_write_timeout(cx))?;
-                Poll::Pending
-            }
-        }
-    }
-
-    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        Pin::new(&mut self.stream).poll_shutdown(cx)
-    }
-}
-
-pub type STcpStream = Connection<TcpStream>;
+const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer
 
 /// Connect to proxy server with `ServerConfig`
-async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc<ServerConfig>) -> io::Result<STcpStream> {
+async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc<ServerConfig>) -> io::Result<TcpStream> {
     let timeout = svr_cfg.timeout();
 
     let svr_addr = match context.config().config_type {
@@ -195,10 +46,7 @@ async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc<ServerConfig>
 
     trace!("Connecting to proxy {:?}, timeout: {:?}", svr_addr, timeout);
     match svr_addr {
-        ServerAddr::SocketAddr(ref addr) => {
-            let stream = try_timeout(TcpStream::connect(addr), timeout).await?;
-            Ok(STcpStream::new(stream, timeout))
-        }
+        ServerAddr::SocketAddr(ref addr) => try_timeout(TcpStream::connect(addr), timeout).await,
         #[cfg(feature = "trust-dns")]
         ServerAddr::DomainName(ref domain, port) => {
             use crate::relay::dns_resolver::resolve;
@@ -211,7 +59,7 @@ async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc<ServerConfig>
             let mut last_err: Option<io::Error> = None;
             for addr in &vec_ipaddr {
                 match try_timeout(TcpStream::connect(addr), timeout).await {
-                    Ok(s) => return Ok(STcpStream::new(s, timeout)),
+                    Ok(s) => return Ok(s),
                     Err(e) => {
                         error!(
                             "Failed to connect {}:{}, resolved address {}, try another (err: {})",
@@ -231,18 +79,17 @@ async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc<ServerConfig>
         }
         #[cfg(not(feature = "trust-dns"))]
         ServerAddr::DomainName(ref domain, port) => {
-            let stream = try_timeout(TcpStream::connect((domain.as_str(), *port)), timeout).await?;
-            Ok(STcpStream::new(stream, timeout))
+            try_timeout(TcpStream::connect((domain.as_str(), *port)), timeout).await
         }
     }
 }
 
 /// Handshake logic for ShadowSocks Client
 pub async fn proxy_server_handshake(
-    remote_stream: STcpStream,
+    remote_stream: TcpStream,
     svr_cfg: Arc<ServerConfig>,
     relay_addr: &Address,
-) -> io::Result<CryptoStream<STcpStream>> {
+) -> io::Result<CryptoStream<TcpStream>> {
     let mut stream = CryptoStream::new(remote_stream, svr_cfg.clone());
 
     trace!("Got encrypt stream and going to send addr: {:?}", relay_addr);
@@ -250,7 +97,7 @@ pub async fn proxy_server_handshake(
     // Send relay address to remote
     let mut addr_buf = BytesMut::with_capacity(relay_addr.serialized_len());
     relay_addr.write_to_buf(&mut addr_buf);
-    stream.write_all(&addr_buf).await?;
+    try_timeout(stream.write_all(&addr_buf), svr_cfg.timeout()).await?;
 
     Ok(stream)
 }

src/relay/tcprelay/server.rs

@@ -2,25 +2,21 @@
 
 use std::{io, net::SocketAddr, sync::Arc};
 
-use crate::relay::socks5::Address;
-
 use crate::context::SharedContext;
+use crate::relay::socks5::Address;
+use crate::relay::utils::try_timeout;
 
-use futures::{
-    future::{self, Either},
-    stream::{FuturesUnordered, StreamExt},
-};
+use futures::future::{self, Either};
+use futures::stream::{FuturesUnordered, StreamExt};
+use futures::FutureExt;
 use log::{debug, error, info, trace};
-use tokio::{
-    self,
-    net::{TcpListener, TcpStream},
-};
+use tokio;
+use tokio::net::{TcpListener, TcpStream};
 
-use super::{
-    context::{SharedTcpServerContext, TcpServerContext},
-    monitor::TcpMonStream,
-    CryptoStream, STcpStream,
-};
+use super::context::{SharedTcpServerContext, TcpServerContext};
+use super::monitor::TcpMonStream;
+use super::utils::copy_timeout;
+use super::CryptoStream;
 
 #[allow(clippy::cognitive_complexity)]
 async fn handle_client(
@@ -44,24 +40,21 @@ async fn handle_client(
         svr_context.svr_cfg()
     );
 
-    let stream = TcpMonStream::new(
-        svr_context.clone(),
-        STcpStream::new(socket, svr_context.svr_cfg().timeout()),
-    );
+    let stream = TcpMonStream::new(svr_context.clone(), socket);
 
     // Do server-client handshake
     // Perform encryption IV exchange
     let mut stream = CryptoStream::new(stream, svr_context.svr_cfg().clone());
 
     // Read remote Address
-    let remote_addr = match Address::read_from(&mut stream).await {
+    let remote_addr = match try_timeout(Address::read_from(&mut stream), svr_context.svr_cfg().timeout()).await {
         Ok(o) => o,
         Err(err) => {
             error!(
                 "Failed to decode Address, may be wrong method or key, peer {}",
                 peer_addr
             );
-            return Err(From::from(err));
+            return Err(err);
         }
     };
 
@@ -167,15 +160,13 @@ async fn handle_client(
     let (mut cr, mut cw) = stream.split();
     let (mut sr, mut sw) = remote_stream.split();
 
-    use tokio::io::copy;
-
     // CLIENT -> SERVER
-    let rhalf = copy(&mut cr, &mut sw);
+    let rhalf = copy_timeout(&mut cr, &mut sw, svr_context.svr_cfg().timeout());
 
     // CLIENT <- SERVER
-    let whalf = copy(&mut sr, &mut cw);
+    let whalf = copy_timeout(&mut sr, &mut cw, svr_context.svr_cfg().timeout());
 
-    match future::select(rhalf, whalf).await {
+    match future::select(rhalf.boxed(), whalf.boxed()).await {
         Either::Left((Ok(_), _)) => trace!("Relay {} -> {} closed", peer_addr, remote_addr),
         Either::Left((Err(err), _)) => trace!("Relay {} -> {} closed with error {:?}", peer_addr, remote_addr, err),
         Either::Right((Ok(_), _)) => trace!("Relay {} <- {} closed", peer_addr, remote_addr),

src/relay/tcprelay/socks5_local.rs

@@ -3,24 +3,20 @@
 use std::{io, net::SocketAddr, sync::Arc};
 
 use futures::future::{self, Either};
+use futures::FutureExt;
 use log::{debug, error, info, trace, warn};
-use tokio::{
-    self,
-    net::{
-        tcp::{ReadHalf, WriteHalf},
-        TcpListener, TcpStream,
-    },
-    prelude::*,
-};
+use tokio;
+use tokio::net::tcp::{ReadHalf, WriteHalf};
+use tokio::net::{TcpListener, TcpStream};
+use tokio::prelude::*;
 
-use crate::{config::ServerConfig, context::SharedContext};
-
-use crate::relay::{
-    loadbalancing::server::{LoadBalancer, PingBalancer},
-    socks5::{self, Address, HandshakeRequest, HandshakeResponse, TcpRequestHeader, TcpResponseHeader},
-};
+use crate::config::ServerConfig;
+use crate::context::SharedContext;
+use crate::relay::loadbalancing::server::{LoadBalancer, PingBalancer};
+use crate::relay::socks5::{self, Address, HandshakeRequest, HandshakeResponse, TcpRequestHeader, TcpResponseHeader};
 
 use super::ignore_until_end;
+use super::utils::copy_timeout;
 
 #[derive(Debug, Clone)]
 struct UdpConfig {
@@ -72,10 +68,8 @@ async fn handle_socks5_connect<'a>(
     let mut svr_s = super::proxy_server_handshake(svr_s, svr_cfg.clone(), addr).await?;
     let (mut svr_r, mut svr_w) = svr_s.split();
 
-    use tokio::io::copy;
-
-    let rhalf = copy(&mut r, &mut svr_w);
-    let whalf = copy(&mut svr_r, &mut w);
+    let rhalf = copy_timeout(&mut r, &mut svr_w, svr_cfg.timeout());
+    let whalf = copy_timeout(&mut svr_r, &mut w, svr_cfg.timeout());
 
     debug!(
         "CONNECT relay established {} <-> {} ({})",
@@ -84,22 +78,22 @@ async fn handle_socks5_connect<'a>(
         addr
     );
 
-    match future::select(rhalf, whalf).await {
+    match future::select(rhalf.boxed(), whalf.boxed()).await {
         Either::Left((Ok(..), _)) => trace!("CONNECT relay {} -> {} ({}) closed", client_addr, svr_cfg.addr(), addr),
         Either::Left((Err(err), _)) => trace!(
             "CONNECT relay {} -> {} ({}) closed with error {:?}",
             client_addr,
             svr_cfg.addr(),
-            err,
             addr,
+            err,
         ),
         Either::Right((Ok(..), _)) => trace!("CONNECT relay {} <- {} ({}) closed", client_addr, svr_cfg.addr(), addr),
         Either::Right((Err(err), _)) => trace!(
             "CONNECT relay {} <- {} ({}) closed with error {:?}",
             client_addr,
             svr_cfg.addr(),
-            err,
             addr,
+            err,
         ),
     }
 

src/relay/tcprelay/tunnel_local.rs

@@ -5,6 +5,7 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use futures::future::{self, Either};
+use futures::FutureExt;
 use log::{debug, error, info, trace};
 use tokio::net::{TcpListener, TcpStream};
 
@@ -13,6 +14,8 @@ use crate::context::SharedContext;
 use crate::relay::loadbalancing::server::{LoadBalancer, PingBalancer};
 use crate::relay::socks5::Address;
 
+use super::utils::copy_timeout;
+
 /// Established Client Tunnel
 ///
 /// This method must be called after handshaking with client (for example, socks5 handshaking)
@@ -39,10 +42,8 @@ async fn establish_client_tcp_tunnel<'a>(
 
     let (mut r, mut w) = s.split();
 
-    use tokio::io::copy;
-
-    let rhalf = copy(&mut r, &mut svr_w);
-    let whalf = copy(&mut svr_r, &mut w);
+    let rhalf = copy_timeout(&mut r, &mut svr_w, svr_cfg.timeout());
+    let whalf = copy_timeout(&mut svr_r, &mut w, svr_cfg.timeout());
 
     debug!(
         "TUNNEL relay established {} <-> {} ({})",
@@ -51,22 +52,22 @@ async fn establish_client_tcp_tunnel<'a>(
         addr
     );
 
-    match future::select(rhalf, whalf).await {
+    match future::select(rhalf.boxed(), whalf.boxed()).await {
         Either::Left((Ok(..), _)) => trace!("TUNNEL relay {} -> {} ({}) closed", client_addr, svr_cfg.addr(), addr),
         Either::Left((Err(err), _)) => trace!(
             "TUNNEL relay {} -> {} ({}) closed with error {:?}",
             client_addr,
             svr_cfg.addr(),
-            err,
             addr,
+            err,
         ),
         Either::Right((Ok(..), _)) => trace!("TUNNEL relay {} <- {} ({}) closed", client_addr, svr_cfg.addr(), addr),
         Either::Right((Err(err), _)) => trace!(
             "TUNNEL relay {} <- {} ({}) closed with error {:?}",
             client_addr,
             svr_cfg.addr(),
+            addr,
             err,
-            addr
         ),
     }
 

src/relay/tcprelay/utils.rs

@@ -1 +1,32 @@
 //! Utility functions
+
+use std::io;
+use std::marker::Unpin;
+use std::time::Duration;
+
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+
+use super::BUFFER_SIZE;
+use crate::relay::utils::try_timeout;
+
+pub async fn copy_timeout<R, W>(r: &mut R, w: &mut W, timeout: Option<Duration>) -> io::Result<u64>
+where
+    R: AsyncRead + Unpin,
+    W: AsyncWrite + Unpin,
+{
+    let mut buf = [0u8; BUFFER_SIZE];
+
+    let mut amt = 0;
+    loop {
+        let n = try_timeout(r.read(&mut buf), timeout).await?;
+        if n == 0 {
+            try_timeout(w.flush(), timeout).await?;
+            break;
+        }
+        try_timeout(w.write_all(&buf[..n]), timeout).await?;
+
+        amt += n as u64;
+    }
+
+    Ok(amt)
+}