implements TCP timeouts

src/config.rs

@@ -219,6 +219,11 @@ impl ServerConfig {
     pub fn method(&self) -> CipherType {
         self.method
     }
+
+    /// Get timeout
+    pub fn timeout(&self) -> &Option<Duration> {
+        &self.timeout
+    }
 }
 
 impl ServerConfig {
@@ -257,7 +262,6 @@ pub struct Config {
     pub server: Vec<ServerConfig>,
     pub local: Option<ClientConfig>,
     pub enable_udp: bool,
-    pub timeout: Option<Duration>,
     pub forbidden_ip: HashSet<IpAddr>,
 }
 
@@ -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<String, Value>, require_local_info: bool) -> Result<Config, Error> {
         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()

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<DecryptedHalf>;
 pub type EncryptedHalfFut = BoxIoFuture<EncryptedHalf>;
 
 fn connect_proxy_server(handle: &Handle, svr_cfg: Rc<ServerConfig>) -> BoxIoFuture<TcpStream> {
+    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<ServerConfig>) -> BoxIoFutu
 /// Handshake logic for ShadowSocks Client
 pub fn proxy_server_handshake(remote_stream: TcpStream,
                               svr_cfg: Rc<ServerConfig>,
-                              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<ServerConfig>)
+                       svr_cfg: Rc<ServerConfig>,
+                       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<T, F>(fut: F, dur: Option<Duration>, handle: &Handle) -> BoxIoFuture<T>
+    where F: Future<Item = T, Error = io::Error> + 'static,
+          T: 'static
+{
+    match dur {
+        Some(dur) => io_timeout(fut, dur, handle),
+        None => boxed_future(fut),
+    }
+}
+
+fn io_timeout<T, F>(fut: F, dur: Duration, handle: &Handle) -> BoxIoFuture<T>
+    where F: Future<Item = T, Error = io::Error> + '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, W>(r: R, w: W, timeout: Option<Duration>, handle: Handle) -> BoxIoFuture<u64>
+    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)
 }

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<TcpRelayClientPending> {
         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<HashSet<IpAddr>>,
+    timeout: Option<Duration>,
 }
 
 impl TcpRelayClientPending {
     /// Resolve Address to SocketAddr
-    fn resolve_address(addr: Address) -> BoxIoFuture<SocketAddr> {
+    fn resolve_address(handle: Handle, addr: Address, timeout: Option<Duration>) -> BoxIoFuture<SocketAddr> {
         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<HashSet<IpAddr>>) -> BoxIoFuture<SocketAddr> {
-        let fut = TcpRelayClientPending::resolve_address(addr).and_then(move |addr| {
+    fn resolve_remote(handle: Handle,
+                      addr: Address,
+                      forbidden_ip: Rc<HashSet<IpAddr>>,
+                      timeout: Option<Duration>)
+                      -> BoxIoFuture<SocketAddr> {
+        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<HashSet<IpAddr>>) -> BoxIoFuture<TcpStream> {
+    fn connect_remote(handle: Handle,
+                      addr: Address,
+                      forbidden_ip: Rc<HashSet<IpAddr>>,
+                      timeout: Option<Duration>)
+                      -> BoxIoFuture<TcpStream> {
         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<TcpRelayClientConnected> {
         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<TcpStream>, WriteHalf<TcpStream>),
     client: (DecryptedHalf, EncryptedHalfFut),
     addr: Address,
+    timeout: Option<Duration>,
+    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)))
     }
 }
 

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)
             })

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<R>
@@ -40,8 +46,6 @@ pub struct DecryptedReader<R>
     sent_final: bool,
 }
 
-const BUFFER_SIZE: usize = 4096;
-
 impl<R> DecryptedReader<R>
     where R: Read
 {
@@ -131,14 +135,14 @@ impl<R> Read for DecryptedReader<R>
 
 /// Writer wrapper that will encrypt data automatically
 pub struct EncryptedWriter<W>
-    where W: Write
+    where W: Write + 'static
 {
     writer: W,
     cipher: CipherVariant,
 }
 
 impl<W> EncryptedWriter<W>
-    where W: Write
+    where W: Write + 'static
 {
     /// Creates a new EncryptedWriter
     pub fn new(w: W, cipher: CipherVariant) -> EncryptedWriter<W> {
@@ -196,10 +200,39 @@ impl<W> EncryptedWriter<W>
             buf: Vec::new(),
         }
     }
+
+    /// Copy all data from reader with timeout
+    pub fn copy_from_encrypted_timeout<R>(mut self, r: R, timeout: Option<Duration>, handle: Handle) -> BoxIoFuture<u64>
+        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<W> Write for EncryptedWriter<W>
-    where W: Write
+    where W: Write + 'static
 {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         self.writer.write(buf)
@@ -212,7 +245,7 @@ impl<W> Write for EncryptedWriter<W>
 
 /// write_all and encrypt data
 pub enum EncryptedWriteAll<W, B>
-    where W: Write,
+    where W: Write + 'static,
           B: AsRef<[u8]>
 {
     Writing {
@@ -260,7 +293,7 @@ impl<W, B> Future for EncryptedWriteAll<W, B>
 }
 
 /// Encrypted copy
-pub struct EncryptedCopy<R: Read, W: Write> {
+pub struct EncryptedCopy<R: Read, W: Write + 'static> {
     reader: R,
     writer: EncryptedWriter<W>,
     read_done: bool,