diff --git a/src/config.rs b/src/config.rs index cdeeac73..7ad3f452 100644 --- a/src/config.rs +++ b/src/config.rs @@ -219,6 +219,11 @@ impl ServerConfig { pub fn method(&self) -> CipherType { self.method } + + /// Get timeout + pub fn timeout(&self) -> &Option { + &self.timeout + } } impl ServerConfig { @@ -257,7 +262,6 @@ pub struct Config { pub server: Vec, pub local: Option, pub enable_udp: bool, - pub timeout: Option, pub forbidden_ip: HashSet, } @@ -340,7 +344,6 @@ impl Config { server: Vec::new(), local: None, enable_udp: false, - timeout: None, forbidden_ip: HashSet::new(), } } @@ -426,15 +429,6 @@ impl Config { fn parse_json_object(o: &Map, require_local_info: bool) -> Result { let mut config = Config::new(); - config.timeout = match o.get("timeout") { - Some(t_str) => { - let val = try!(t_str.as_u64() - .ok_or(Error::new(ErrorKind::Malformed, "`timeout` should be an integer", None))); - Some(Duration::from_secs(val)) - } - None => None, - }; - if o.contains_key("servers") { let server_list = try!(o.get("servers") .unwrap() diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index 634c54c6..a684b6a9 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -21,10 +21,11 @@ //! TcpRelay implementation -use std::io::{self, Read}; +use std::io::{self, Read, Write}; use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6}; use std::rc::Rc; use std::mem; +use std::time::Duration; use crypto::cipher; use crypto::CryptoMode; @@ -34,8 +35,8 @@ use relay::dns_resolver::DnsResolver; use config::{ServerConfig, ServerAddr}; use tokio_core::net::TcpStream; -use tokio_core::reactor::Handle; -use tokio_core::io::{read_exact, write_all}; +use tokio_core::reactor::{Handle, Timeout}; +use tokio_core::io::{read_exact, write_all, read}; use tokio_core::io::{ReadHalf, WriteHalf}; use tokio_core::io::Io; @@ -53,6 +54,8 @@ pub mod server; mod stream; pub mod client; +const BUFFER_SIZE: usize = 4096; + /// Directions in the tunnel #[derive(Debug, Copy, Clone)] pub enum TunnelDirection { @@ -73,18 +76,22 @@ pub type DecryptedHalfFut = BoxIoFuture; pub type EncryptedHalfFut = BoxIoFuture; fn connect_proxy_server(handle: &Handle, svr_cfg: Rc) -> BoxIoFuture { + let timeout = svr_cfg.timeout().clone(); + trace!("Connecting to proxy {:?}, timeout: {:?}", + svr_cfg.addr(), + timeout); match svr_cfg.addr() { - &ServerAddr::SocketAddr(ref addr) => TcpStream::connect(addr, handle).boxed(), + &ServerAddr::SocketAddr(ref addr) => try_timeout(TcpStream::connect(addr, handle), timeout, handle), &ServerAddr::DomainName(ref domain, port) => { let handle = handle.clone(); - let fut = DnsResolver::resolve(&domain[..]).and_then(move |sockaddr| { + let fut = try_timeout(DnsResolver::resolve(&domain[..]), timeout, &handle).and_then(move |sockaddr| { let sockaddr = match sockaddr { IpAddr::V4(v4) => SocketAddr::V4(SocketAddrV4::new(v4, port)), IpAddr::V6(v6) => SocketAddr::V6(SocketAddrV6::new(v6, port, 0, 0)), }; - TcpStream::connect(&sockaddr, &handle) + try_timeout(TcpStream::connect(&sockaddr, &handle), timeout, &handle) }); - Box::new(fut) + boxed_future(fut) } } } @@ -92,9 +99,11 @@ fn connect_proxy_server(handle: &Handle, svr_cfg: Rc) -> BoxIoFutu /// Handshake logic for ShadowSocks Client pub fn proxy_server_handshake(remote_stream: TcpStream, svr_cfg: Rc, - relay_addr: Address) + relay_addr: Address, + handle: Handle) -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> { - let fut = proxy_handshake(remote_stream, svr_cfg).and_then(|(r_fut, w_fut)| {; + let timeout = svr_cfg.timeout().clone(); + let fut = proxy_handshake(remote_stream, svr_cfg, handle.clone()).and_then(move |(r_fut, w_fut)| {; let w_fut = w_fut.and_then(move |enc_w| { trace!("Got encrypt stream and going to send addr: {:?}", relay_addr); @@ -102,7 +111,7 @@ pub fn proxy_server_handshake(remote_stream: TcpStream, // Send relay address to remote let local_buf = Vec::new(); relay_addr.write_to(local_buf) - .and_then(|buf| enc_w.write_all_encrypted(buf)) + .and_then(move |buf| try_timeout(enc_w.write_all_encrypted(buf), timeout, &handle)) .map(|(w, _)| w) }); @@ -114,10 +123,12 @@ pub fn proxy_server_handshake(remote_stream: TcpStream, /// ShadowSocks Client-Server handshake protocol /// Exchange cipher IV and creates stream wrapper pub fn proxy_handshake(remote_stream: TcpStream, - svr_cfg: Rc) + svr_cfg: Rc, + handle: Handle) -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> { - let fut = futures::lazy(|| Ok(remote_stream.split())).and_then(|(r, w)| { + let fut = futures::lazy(|| Ok(remote_stream.split())).and_then(move |(r, w)| { + let timeout = svr_cfg.timeout().clone(); let svr_cfg_cloned = svr_cfg.clone(); @@ -129,7 +140,7 @@ pub fn proxy_handshake(remote_stream: TcpStream, let local_iv = svr_cfg.method().gen_init_vec(); trace!("Going to send initialize vector: {:?}", local_iv); - write_all(w, local_iv).and_then(move |(w, local_iv)| { + try_timeout(write_all(w, local_iv), timeout.clone(), &handle).and_then(move |(w, local_iv)| { let encryptor = cipher::with_type(svr_cfg.method(), svr_cfg.key(), &local_iv[..], @@ -144,7 +155,7 @@ pub fn proxy_handshake(remote_stream: TcpStream, // Decrypt data from remote server let iv_len = svr_cfg.method().iv_size(); - read_exact(r, vec![0u8; iv_len]).and_then(move |(r, remote_iv)| { + try_timeout(read_exact(r, vec![0u8; iv_len]), timeout, &handle).and_then(move |(r, remote_iv)| { trace!("Got initialize vector {:?}", remote_iv); let decryptor = cipher::with_type(svr_cfg.method(), @@ -266,4 +277,48 @@ fn reuse_port(builder: &TcpBuilder) -> io::Result<&TcpBuilder> { #[cfg(windows)] fn reuse_port(builder: &TcpBuilder) -> io::Result<&TcpBuilder> { Ok(builder) +} + +fn try_timeout(fut: F, dur: Option, handle: &Handle) -> BoxIoFuture + where F: Future + 'static, + T: 'static +{ + match dur { + Some(dur) => io_timeout(fut, dur, handle), + None => boxed_future(fut), + } +} + +fn io_timeout(fut: F, dur: Duration, handle: &Handle) -> BoxIoFuture + where F: Future + 'static, + T: 'static +{ + let fut = fut.select(Timeout::new(dur, handle) + .unwrap() // It must be succeeded! + .and_then(|_| Err(io::Error::new(io::ErrorKind::TimedOut, "timeout")))) + .then(|res| { + match res { + Ok((t, _)) => Ok(t), + Err((err, _)) => Err(err), + } + }); + boxed_future(fut) +} + +fn copy_timeout(r: R, w: W, timeout: Option, handle: Handle) -> BoxIoFuture + where R: Read + 'static, + W: Write + 'static +{ + let fut = try_timeout(read(r, vec![0u8; BUFFER_SIZE]), timeout.clone(), &handle) + .and_then(move |(r, mut buf, n)| { + if n == 0 { + boxed_future(futures::finished(n as u64)) + } else { + buf.resize(n, 0); + let fut = try_timeout(write_all(w, buf), timeout.clone(), &handle) + .and_then(move |(w, _)| copy_timeout(r, w, timeout, handle).map(move |x| x + n as u64)); + boxed_future(fut) + } + }); + boxed_future(fut) } \ No newline at end of file diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index 315b6efa..53024ca3 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -25,6 +25,7 @@ use std::io; use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6}; use std::rc::Rc; use std::collections::HashSet; +use std::time::Duration; use config::{Config, ServerConfig}; @@ -38,13 +39,12 @@ use futures::stream::Stream; use tokio_core::reactor::Handle; use tokio_core::net::{TcpStream, TcpListener}; use tokio_core::io::{Io, ReadHalf, WriteHalf}; -use tokio_core::io::copy; use net2::TcpBuilder; use ip::IpAddr; -use super::{tunnel, proxy_handshake, DecryptedHalf, EncryptedHalfFut}; +use super::{tunnel, proxy_handshake, DecryptedHalf, EncryptedHalfFut, try_timeout, copy_timeout}; /// Context for doing handshake with client pub struct TcpRelayClientHandshake { @@ -59,12 +59,17 @@ impl TcpRelayClientHandshake { pub fn handshake(self) -> BoxIoFuture { let TcpRelayClientHandshake { handle, forbidden_ip, s, svr_cfg } = self; - let fut = proxy_handshake(s, svr_cfg).and_then(|(r_fut, w_fut)| { - r_fut.and_then(|r| { - Address::read_from(r).map_err(|_| { - io::Error::new(io::ErrorKind::Other, - "failed to decode Address, may be wrong method or key") - }) + let timeout = svr_cfg.timeout().clone(); + let cloned_handle = handle.clone(); + let fut = proxy_handshake(s, svr_cfg, handle.clone()).and_then(move |(r_fut, w_fut)| { + r_fut.and_then(move |r| { + let handle = cloned_handle; + try_timeout(Address::read_from(r).map_err(|_| { + io::Error::new(io::ErrorKind::Other, + "failed to decode Address, may be wrong method or key") + }), + timeout, + &handle) }) .map(move |(r, addr)| { TcpRelayClientPending { @@ -73,6 +78,7 @@ impl TcpRelayClientHandshake { addr: addr, w: w_fut, forbidden_ip: forbidden_ip, + timeout: timeout, } }) }); @@ -87,15 +93,16 @@ pub struct TcpRelayClientPending { addr: Address, w: EncryptedHalfFut, forbidden_ip: Rc>, + timeout: Option, } impl TcpRelayClientPending { /// Resolve Address to SocketAddr - fn resolve_address(addr: Address) -> BoxIoFuture { + fn resolve_address(handle: Handle, addr: Address, timeout: Option) -> BoxIoFuture { match addr { Address::SocketAddress(addr) => Box::new(futures::finished(addr)), Address::DomainNameAddress(dname, port) => { - let fut = DnsResolver::resolve(&dname[..]).and_then(move |ipaddr| { + let fut = try_timeout(DnsResolver::resolve(&dname[..]), timeout, &handle).and_then(move |ipaddr| { Ok(match ipaddr { IpAddr::V4(v4) => SocketAddr::V4(SocketAddrV4::new(v4, port)), IpAddr::V6(v6) => SocketAddr::V6(SocketAddrV6::new(v6, port, 0, 0)), @@ -108,8 +115,12 @@ impl TcpRelayClientPending { /// Resolve remote address to SocketAddr /// Report failure if the SocketAddr is forbidden by `forbidden_ip` - fn resolve_remote(addr: Address, forbidden_ip: Rc>) -> BoxIoFuture { - let fut = TcpRelayClientPending::resolve_address(addr).and_then(move |addr| { + fn resolve_remote(handle: Handle, + addr: Address, + forbidden_ip: Rc>, + timeout: Option) + -> BoxIoFuture { + let fut = TcpRelayClientPending::resolve_address(handle, addr, timeout).and_then(move |addr| { trace!("Resolved address as {}", addr); let ipaddr = match addr.clone() { SocketAddr::V4(v4) => IpAddr::V4(v4.ip().clone()), @@ -128,23 +139,32 @@ impl TcpRelayClientPending { } /// Connect to the remote server - fn connect_remote(handle: Handle, addr: Address, forbidden_ip: Rc>) -> BoxIoFuture { + fn connect_remote(handle: Handle, + addr: Address, + forbidden_ip: Rc>, + timeout: Option) + -> BoxIoFuture { info!("Connecting to remote {}", addr); - Box::new(TcpRelayClientPending::resolve_remote(addr, forbidden_ip) - .and_then(move |addr| TcpStream::connect(&addr, &handle))) + Box::new(TcpRelayClientPending::resolve_remote(handle.clone(), addr, forbidden_ip, timeout.clone()) + .and_then(move |addr| try_timeout(TcpStream::connect(&addr, &handle), timeout, &handle))) } /// Connect to the remote server pub fn connect(self) -> BoxIoFuture { let addr = self.addr.clone(); let client_pair = (self.r, self.w); - let fut = TcpRelayClientPending::connect_remote(self.handle, self.addr, self.forbidden_ip).map(|stream| { - TcpRelayClientConnected { - server: stream.split(), - client: client_pair, - addr: addr, - } - }); + let timeout = self.timeout.clone(); + let handle = self.handle.clone(); + let fut = TcpRelayClientPending::connect_remote(self.handle, self.addr, self.forbidden_ip, self.timeout) + .map(move |stream| { + TcpRelayClientConnected { + server: stream.split(), + client: client_pair, + addr: addr, + timeout: timeout, + handle: handle, + } + }); Box::new(fut) } } @@ -154,6 +174,8 @@ pub struct TcpRelayClientConnected { server: (ReadHalf, WriteHalf), client: (DecryptedHalf, EncryptedHalfFut), addr: Address, + timeout: Option, + handle: Handle, } impl TcpRelayClientConnected { @@ -161,9 +183,11 @@ impl TcpRelayClientConnected { pub fn tunnel(self) -> BoxIoFuture<()> { let (svr_r, svr_w) = self.server; let (r, w_fut) = self.client; + let timeout = self.timeout.clone(); + let handle = self.handle.clone(); tunnel(self.addr, - copy(r, svr_w), - w_fut.and_then(|w| w.copy_from_encrypted(svr_r))) + copy_timeout(r, svr_w, self.timeout, self.handle), + w_fut.and_then(move |w| w.copy_from_encrypted_timeout(svr_r, timeout, handle))) } } diff --git a/src/relay/tcprelay/socks5_local.rs b/src/relay/tcprelay/socks5_local.rs index 8770509d..8db383bf 100644 --- a/src/relay/tcprelay/socks5_local.rs +++ b/src/relay/tcprelay/socks5_local.rs @@ -32,7 +32,7 @@ use tokio_core::net::{TcpStream, TcpListener}; use tokio_core::reactor::Handle; use tokio_core::io::Io; use tokio_core::io::{ReadHalf, WriteHalf}; -use tokio_core::io::{flush, copy}; +use tokio_core::io::flush; use net2::TcpBuilder; @@ -44,7 +44,7 @@ use relay::loadbalancing::server::RoundRobin; use relay::loadbalancing::server::LoadBalancer; use relay::{BoxIoFuture, boxed_future}; -use super::{tunnel, ignore_until_end}; +use super::{tunnel, ignore_until_end, try_timeout}; #[derive(Debug, Clone)] struct UdpConfig { @@ -60,6 +60,9 @@ fn handle_socks5_connect(handle: &Handle, -> BoxIoFuture<()> { let cloned_addr = addr.clone(); let cloned_svr_cfg = svr_cfg.clone(); + let cloned_handle = handle.clone(); + let cloned_handle2 = handle.clone(); + let timeout = svr_cfg.timeout().clone(); let fut = super::connect_proxy_server(handle, svr_cfg) .and_then(move |svr_s| { trace!("Proxy server connected"); @@ -68,15 +71,23 @@ fn handle_socks5_connect(handle: &Handle, let header = TcpResponseHeader::new(socks5::Reply::Succeeded, Address::SocketAddress(client_addr)); trace!("Send header: {:?}", header); + let handle = cloned_handle; - header.write_to(w) - .and_then(flush) - .map(move |w| (svr_s, w)) + let fut = try_timeout(header.write_to(w), timeout.clone(), &handle); + let fut = try_timeout(fut.and_then(flush), timeout, &handle); + fut.map(move |w| (svr_s, w)) }) .and_then(move |(svr_s, w)| { - super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(|(svr_r, svr_w)| { - let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); - let whalf = svr_w.and_then(move |svr_w| svr_w.copy_from_encrypted(r)); + let handle = cloned_handle2; + let svr_cfg = cloned_svr_cfg; + let timeout = svr_cfg.timeout().clone(); + super::proxy_server_handshake(svr_s, svr_cfg, addr, handle.clone()).and_then(move |(svr_r, svr_w)| { + let cloned_handle = handle.clone(); + let cloned_timeout = timeout.clone(); + + let rhalf = svr_r.and_then(move |svr_r| super::copy_timeout(svr_r, w, timeout, handle)); + let whalf = + svr_w.and_then(move |svr_w| svr_w.copy_from_encrypted_timeout(r, cloned_timeout, cloned_handle)); tunnel(cloned_addr, whalf, rhalf) }) diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index a45fe3f7..3e95edff 100644 --- a/src/relay/tcprelay/stream.rs +++ b/src/relay/tcprelay/stream.rs @@ -24,10 +24,16 @@ use std::io::{self, Read, BufRead, Write}; use std::cmp; use std::mem; +use std::time::Duration; use crypto::{Cipher, CipherVariant}; -use futures::{Future, Poll}; +use futures::{self, Future, Poll}; + +use tokio_core::io::{read, write_all}; +use tokio_core::reactor::Handle; + +use super::{BUFFER_SIZE, BoxIoFuture, try_timeout, boxed_future}; /// Reader wrapper that will decrypt data automatically pub struct DecryptedReader @@ -40,8 +46,6 @@ pub struct DecryptedReader sent_final: bool, } -const BUFFER_SIZE: usize = 4096; - impl DecryptedReader where R: Read { @@ -131,14 +135,14 @@ impl Read for DecryptedReader /// Writer wrapper that will encrypt data automatically pub struct EncryptedWriter - where W: Write + where W: Write + 'static { writer: W, cipher: CipherVariant, } impl EncryptedWriter - where W: Write + where W: Write + 'static { /// Creates a new EncryptedWriter pub fn new(w: W, cipher: CipherVariant) -> EncryptedWriter { @@ -196,10 +200,39 @@ impl EncryptedWriter buf: Vec::new(), } } + + /// Copy all data from reader with timeout + pub fn copy_from_encrypted_timeout(mut self, r: R, timeout: Option, handle: Handle) -> BoxIoFuture + where R: Read + 'static + { + let buf = try_timeout(read(r, vec![0u8; BUFFER_SIZE]), timeout.clone(), &handle) + .and_then(move |(r, buf, n)| { + let mut enc_buf = Vec::new(); + let cont = if n == 0 { + try!(self.cipher_finalize(&mut enc_buf)); + false + } else { + try!(self.cipher_update(&buf[..n], &mut enc_buf)); + true + }; + + Ok((r, self, enc_buf, n, cont)) + }) + .and_then(move |(r, w, buf, n, cont)| { + write_all(w, buf).and_then(move |(w, _)| { + if cont { + boxed_future(w.copy_from_encrypted_timeout(r, timeout, handle).map(move |x| x + n as u64)) + } else { + boxed_future(futures::finished(n as u64)) + } + }) + }); + boxed_future(buf) + } } impl Write for EncryptedWriter - where W: Write + where W: Write + 'static { fn write(&mut self, buf: &[u8]) -> io::Result { self.writer.write(buf) @@ -212,7 +245,7 @@ impl Write for EncryptedWriter /// write_all and encrypt data pub enum EncryptedWriteAll - where W: Write, + where W: Write + 'static, B: AsRef<[u8]> { Writing { @@ -260,7 +293,7 @@ impl Future for EncryptedWriteAll } /// Encrypted copy -pub struct EncryptedCopy { +pub struct EncryptedCopy { reader: R, writer: EncryptedWriter, read_done: bool,