From 82f698e5206a97d8813fd1d37b596df5eb8ef49d Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Sun, 1 Dec 2019 01:33:33 +0800 Subject: [PATCH] Removed Connection, uses try_timeout instead. Revert back to use 8K buffer --- src/relay/dns_resolver/trust_dns_resolver.rs | 6 +- src/relay/loadbalancing/server/mod.rs | 3 +- src/relay/loadbalancing/server/ping.rs | 31 ++-- src/relay/loadbalancing/server/roundrobin.rs | 6 +- src/relay/tcprelay/client.rs | 8 +- src/relay/tcprelay/mod.rs | 179 ++----------------- src/relay/tcprelay/server.rs | 43 ++--- src/relay/tcprelay/socks5_local.rs | 36 ++-- src/relay/tcprelay/tunnel_local.rs | 15 +- src/relay/tcprelay/utils.rs | 31 ++++ 10 files changed, 107 insertions(+), 251 deletions(-) diff --git a/src/relay/dns_resolver/trust_dns_resolver.rs b/src/relay/dns_resolver/trust_dns_resolver.rs index 2742eadb..a12f6bad 100644 --- a/src/relay/dns_resolver/trust_dns_resolver.rs +++ b/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; diff --git a/src/relay/loadbalancing/server/mod.rs b/src/relay/loadbalancing/server/mod.rs index 859ae0d6..a492b7cc 100644 --- a/src/relay/loadbalancing/server/mod.rs +++ b/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; diff --git a/src/relay/loadbalancing/server/ping.rs b/src/relay/loadbalancing/server/ping.rs index 2b178cfe..54cb72dd 100644 --- a/src/relay/loadbalancing/server/ping.rs +++ b/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, diff --git a/src/relay/loadbalancing/server/roundrobin.rs b/src/relay/loadbalancing/server/roundrobin.rs index 9c3bf201..e7427366 100644 --- a/src/relay/loadbalancing/server/roundrobin.rs +++ b/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 { diff --git a/src/relay/tcprelay/client.rs b/src/relay/tcprelay/client.rs index ab421e56..7fc6b36d 100644 --- a/src/relay/tcprelay/client.rs +++ b/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, + pub stream: CryptoStream, } impl ServerClient { diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index 268d6443..786fe92a 100644 --- a/src/relay/tcprelay/mod.rs +++ b/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 { - stream: S, - read_timer: Option>>, - write_timer: Option>>, - timeout: Option, -} - -impl Connection { - pub fn new(stream: S, timeout: Option) -> Connection { - Connection { - stream, - read_timer: None, - write_timer: None, - timeout, - } - } - - fn poll_read_timeout(&mut self, cx: &mut Context<'_>) -> Poll> { - 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> { - 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 Connection -where - S: AsyncRead + AsyncWrite + Unpin, -{ - pub fn split(self) -> (ReadHalf>, WriteHalf>) { - use tokio::io::split; - split(self) - } -} - -impl Deref for Connection { - type Target = S; - - fn deref(&self) -> &Self::Target { - &self.stream - } -} - -impl DerefMut for Connection { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.stream - } -} - -impl Unpin for Connection {} - -impl AsyncRead for Connection -where - S: AsyncRead + Unpin, -{ - fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll> { - 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 AsyncWrite for Connection -where - S: AsyncWrite + Unpin, -{ - fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { - 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> { - 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> { - Pin::new(&mut self.stream).poll_shutdown(cx) - } -} - -pub type STcpStream = Connection; +const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer /// Connect to proxy server with `ServerConfig` -async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc) -> io::Result { +async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc) -> io::Result { 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 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 let mut last_err: Option = 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 } #[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, relay_addr: &Address, -) -> io::Result> { +) -> io::Result> { 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) } diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index 4b8fe611..d54825b9 100644 --- a/src/relay/tcprelay/server.rs +++ b/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), diff --git a/src/relay/tcprelay/socks5_local.rs b/src/relay/tcprelay/socks5_local.rs index 7600a64b..c01c2d13 100644 --- a/src/relay/tcprelay/socks5_local.rs +++ b/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, ), } diff --git a/src/relay/tcprelay/tunnel_local.rs b/src/relay/tcprelay/tunnel_local.rs index 9c3b2b54..e719e23a 100644 --- a/src/relay/tcprelay/tunnel_local.rs +++ b/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 ), } diff --git a/src/relay/tcprelay/utils.rs b/src/relay/tcprelay/utils.rs index ab0c6515..123267ac 100644 --- a/src/relay/tcprelay/utils.rs +++ b/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: &mut R, w: &mut W, timeout: Option) -> io::Result +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) +}