From 34a70a7dc77abd217ea5625cf93083f2feef04bd Mon Sep 17 00:00:00 2001 From: "Y. T. Chung" Date: Mon, 7 Nov 2016 03:19:14 +0800 Subject: [PATCH] memory optimization --- src/relay/tcprelay/http_local.rs | 53 ++++++++------- src/relay/tcprelay/mod.rs | 101 ++++++++++++++--------------- src/relay/tcprelay/server.rs | 27 ++++---- src/relay/tcprelay/socks5_local.rs | 9 ++- src/relay/tcprelay/stream.rs | 4 +- 5 files changed, 92 insertions(+), 102 deletions(-) diff --git a/src/relay/tcprelay/http_local.rs b/src/relay/tcprelay/http_local.rs index 0e3c192f..02eb6236 100644 --- a/src/relay/tcprelay/http_local.rs +++ b/src/relay/tcprelay/http_local.rs @@ -474,13 +474,12 @@ fn handle_connect(handle: Handle, .map(|w| (svr_s, w)) }) .and_then(move |(svr_s, w)| { - super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(move |(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.write_all_encrypted(remains)) - .and_then(move |(svr_w, _)| svr_w.copy_from_encrypted(r)); + let (svr_r, svr_w) = super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr); + let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); + let whalf = svr_w.and_then(move |svr_w| svr_w.write_all_encrypted(remains)) + .and_then(move |(svr_w, _)| svr_w.copy_from_encrypted(r)); - tunnel(cloned_addr, whalf, rhalf) - }) + tunnel(cloned_addr, whalf, rhalf) }); Box::new(fut) @@ -535,28 +534,28 @@ fn handle_http_proxy(handle: Handle, trace!("Proxy server connected"); let cloned_addr = addr.clone(); - super::proxy_server_handshake(svr_s, svr_cfg, addr).and_then(move |(svr_r, svr_w)| { - // Just proxy anything to client - let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); - let whalf = svr_w.and_then(move |svr_w| { - // Send the first request to server - trace!("Going to proxy request: {:?}", req); - trace!("Should keep alive? {}", should_keep_alive); - proxy_request_encrypted((r, svr_w), None, req, remains).and_then(move |(r, svr_w, req_remains)| { - if should_keep_alive { - handle_http_keepalive(r, svr_w, req_remains) - } else { - futures::finished(()).boxed() - } - }) - }); + let (svr_r, svr_w) = super::proxy_server_handshake(svr_s, svr_cfg, addr); - rhalf.join(whalf) - .then(move |_| { - trace!("Relay to {} is finished", cloned_addr); - Ok(()) - }) - }) + // Just proxy anything to client + let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); + let whalf = svr_w.and_then(move |svr_w| { + // Send the first request to server + trace!("Going to proxy request: {:?}", req); + trace!("Should keep alive? {}", should_keep_alive); + proxy_request_encrypted((r, svr_w), None, req, remains).and_then(move |(r, svr_w, req_remains)| { + if should_keep_alive { + handle_http_keepalive(r, svr_w, req_remains) + } else { + futures::finished(()).boxed() + } + }) + }); + + rhalf.join(whalf) + .then(move |_| { + trace!("Relay to {} is finished", cloned_addr); + Ok(()) + }) }); Box::new(fut) diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index 6e44b550..eda23a85 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -36,11 +36,11 @@ use config::{ServerConfig, ServerAddr}; use tokio_core::net::TcpStream; use tokio_core::reactor::Handle; -use tokio_core::io::{read_exact, write_all, flush}; +use tokio_core::io::{read_exact, write_all}; use tokio_core::io::{ReadHalf, WriteHalf}; use tokio_core::io::Io; -use futures::{self, Future, Poll}; +use futures::{Future, Poll}; use net2::TcpBuilder; @@ -76,7 +76,7 @@ pub type EncryptedHalfFut = BoxIoFuture; fn connect_proxy_server(handle: &Handle, svr_cfg: Rc) -> BoxIoFuture { match svr_cfg.addr() { - &ServerAddr::SocketAddr(ref addr) => Box::new(TcpStream::connect(addr, handle)), + &ServerAddr::SocketAddr(ref addr) => TcpStream::connect(addr, handle).boxed(), &ServerAddr::DomainName(ref domain, port) => { let handle = handle.clone(); let fut = DnsResolver::get_instance() @@ -86,7 +86,7 @@ fn connect_proxy_server(handle: &Handle, svr_cfg: Rc) -> BoxIoFutu IpAddr::V4(v4) => SocketAddr::V4(SocketAddrV4::new(v4, port)), IpAddr::V6(v6) => SocketAddr::V6(SocketAddrV6::new(v6, port, 0, 0)), }; - TcpStream::connect(&sockaddr, &handle).boxed() + TcpStream::connect(&sockaddr, &handle) }); Box::new(fut) } @@ -97,73 +97,66 @@ fn connect_proxy_server(handle: &Handle, svr_cfg: Rc) -> BoxIoFutu pub fn proxy_server_handshake(remote_stream: TcpStream, svr_cfg: Rc, relay_addr: Address) - -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> { - let fut = proxy_handshake(remote_stream, svr_cfg).and_then(|(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); + -> (DecryptedHalfFut, EncryptedHalfFut) { + let (r_fut, w_fut) = proxy_handshake(remote_stream, svr_cfg); + let w_fut = w_fut.and_then(move |enc_w| { + trace!("Got encrypt stream and going to send addr: {:?}", + relay_addr); - // 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(|(enc_w, _)| flush(enc_w)) - }); - Ok((r_fut, boxed_future(w_fut))) + // 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)) + .map(|(w, _)| w) }); - Box::new(fut) + + (r_fut, boxed_future(w_fut)) } /// ShadowSocks Client-Server handshake protocol /// Exchange cipher IV and creates stream wrapper -pub fn proxy_handshake(remote_stream: TcpStream, - svr_cfg: Rc) - -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> { - let fut = futures::lazy(move || { - let (r, w) = remote_stream.split(); +pub fn proxy_handshake(remote_stream: TcpStream, svr_cfg: Rc) -> (DecryptedHalfFut, EncryptedHalfFut) { + let (r, w) = remote_stream.split(); - let svr_cfg_cloned = svr_cfg.clone(); + let svr_cfg_cloned = svr_cfg.clone(); - let enc = futures::lazy(move || { - // Encrypt data to remote server + let enc = { + // Encrypt data to remote server - // Send initialize vector to remote and create encryptor + // Send initialize vector to remote and create encryptor - let local_iv = svr_cfg.method().gen_init_vec(); - trace!("Going to send initialize vector: {:?}", local_iv); + 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)| { - let encryptor = cipher::with_type(svr_cfg.method(), - svr_cfg.key(), - &local_iv[..], - CryptoMode::Encrypt); + write_all(w, local_iv).and_then(move |(w, local_iv)| { + let encryptor = cipher::with_type(svr_cfg.method(), + svr_cfg.key(), + &local_iv[..], + CryptoMode::Encrypt); - Ok(EncryptedWriter::new(w, encryptor)) - }) + Ok(EncryptedWriter::new(w, encryptor)) + }) + }; - }); + let dec = { + let svr_cfg = svr_cfg_cloned; - let dec = futures::lazy(move || { - let svr_cfg = svr_cfg_cloned; + // 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)| { + trace!("Got initialize vector {:?}", remote_iv); - // 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)| { - trace!("Got initialize vector {:?}", remote_iv); + let decryptor = cipher::with_type(svr_cfg.method(), + svr_cfg.key(), + &remote_iv[..], + CryptoMode::Decrypt); + let decrypt_stream = DecryptedReader::new(r, decryptor); - let decryptor = cipher::with_type(svr_cfg.method(), - svr_cfg.key(), - &remote_iv[..], - CryptoMode::Decrypt); - let decrypt_stream = DecryptedReader::new(r, decryptor); + Ok(decrypt_stream) + }) + }; - Ok(decrypt_stream) - }) - }); - - Ok((boxed_future(dec), boxed_future(enc))) - }); - Box::new(fut) + (boxed_future(dec), boxed_future(enc)) } /// Copy exactly N bytes by encryption diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index f1c9583e..d4119c58 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -29,7 +29,7 @@ use std::collections::HashSet; use config::{Config, ServerConfig}; use relay::socks5::Address; -use relay::BoxIoFuture; +use relay::{BoxIoFuture, boxed_future}; use relay::dns_resolver::DnsResolver; use futures::{self, Future}; @@ -59,19 +59,18 @@ 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(From::from)) - .map(move |(r, addr)| { - TcpRelayClientPending { - handle: handle, - r: r, - addr: addr, - w: w_fut, - forbidden_ip: forbidden_ip, - } - }) - }); - Box::new(fut) + let (r_fut, w_fut) = proxy_handshake(s, svr_cfg); + let fut = r_fut.and_then(|r| Address::read_from(r).map_err(From::from)) + .map(move |(r, addr)| { + TcpRelayClientPending { + handle: handle, + r: r, + addr: addr, + w: w_fut, + forbidden_ip: forbidden_ip, + } + }); + boxed_future(fut) } } diff --git a/src/relay/tcprelay/socks5_local.rs b/src/relay/tcprelay/socks5_local.rs index 6afe3643..dbbe1792 100644 --- a/src/relay/tcprelay/socks5_local.rs +++ b/src/relay/tcprelay/socks5_local.rs @@ -74,12 +74,11 @@ fn handle_socks5_connect(handle: &Handle, .map(move |w| (svr_s, w)) }) .and_then(move |(svr_s, w)| { - super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(move |(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 (svr_r, svr_w) = super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr); + 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)); - tunnel(cloned_addr, whalf, rhalf) - }) + tunnel(cloned_addr, whalf, rhalf) }); Box::new(fut) diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index bfe6fffb..a45fe3f7 100644 --- a/src/relay/tcprelay/stream.rs +++ b/src/relay/tcprelay/stream.rs @@ -40,7 +40,7 @@ pub struct DecryptedReader sent_final: bool, } -const BUFFER_SIZE: usize = 2048; +const BUFFER_SIZE: usize = 4096; impl DecryptedReader where R: Read @@ -275,7 +275,7 @@ impl Future for EncryptedCopy { type Error = io::Error; fn poll(&mut self) -> Poll { - let mut local_buf = [0u8; 2048]; + let mut local_buf = [0u8; BUFFER_SIZE]; loop { // If our buffer is empty, then we need to read some data to // continue.