From 4fcab1ad8ff115cd632409d5ffefa4ab45f7bb4a Mon Sep 17 00:00:00 2001 From: "Y. T. Chung" Date: Sun, 18 Jan 2015 00:02:32 +0800 Subject: [PATCH] reuse thread --- src/relay/tcprelay/local.rs | 30 ++++++++++----------- src/relay/tcprelay/server.rs | 52 +++++++++++++++--------------------- src/relay/tcprelay/stream.rs | 36 +++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 46 deletions(-) diff --git a/src/relay/tcprelay/local.rs b/src/relay/tcprelay/local.rs index 316a1f5c..88665682 100644 --- a/src/relay/tcprelay/local.rs +++ b/src/relay/tcprelay/local.rs @@ -230,23 +230,21 @@ impl TcpRelayLocal { remote_iv.as_slice(), CryptoMode::Decrypt); let mut decrypt_stream = DecryptedReader::new(remote_stream.clone(), decryptor); - Thread::spawn(move || { - match io::util::copy(&mut decrypt_stream, &mut stream) { - Err(err) => { - match err.kind { - EndOfFile | BrokenPipe => { - debug!("{} relay from local to remote stream: {}", addr, err) - }, - _ => { - error!("{} relay from local to remote stream: {}", addr, err) - } + match io::util::copy(&mut decrypt_stream, &mut stream) { + Err(err) => { + match err.kind { + EndOfFile | BrokenPipe => { + debug!("{} relay from local to remote stream: {}", addr, err) + }, + _ => { + error!("{} relay from local to remote stream: {}", addr, err) } - remote_stream.close_write().or(Ok(())).unwrap(); - stream.close_read().or(Ok(())).unwrap(); - }, - Ok(..) => {}, - } - }); + } + remote_stream.close_write().or(Ok(())).unwrap(); + stream.close_read().or(Ok(())).unwrap(); + }, + Ok(..) => {}, + } }, socks5::Command::TcpBind => { warn!("BIND is not supported"); diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index 3226349e..12622f41 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -121,7 +121,7 @@ impl TcpRelayServer { ); info!("Connecting to {}", addr); - let remote_stream = match addr { + let mut remote_stream = match addr { Address::SocketAddress(ip, port) => { try_result!(TcpStream::connect_timeout(SocketAddr {ip: ip, port: port}, Duration::seconds(30)), prefix: format!("Unable to connect {}:", addr) @@ -157,17 +157,15 @@ impl TcpRelayServer { panic!("Unable to connect {}", addr); }; connector() - - // TcpStream::connect(domainaddr.domain_name.as_slice(), domainaddr.port) - // .ok().expect(format!("Unable to connect {}", domainaddr).as_slice()) } }; let mut remote_stream_cloned = remote_stream.clone(); let addr_cloned = addr.clone(); Thread::spawn(move || { - io::util::copy(&mut decrypt_stream, &mut remote_stream_cloned) - .unwrap_or_else(|err| { + match io::util::copy(&mut decrypt_stream, &mut remote_stream_cloned) { + Ok(..) => {}, + Err(err) => { match err.kind { EndOfFile | BrokenPipe => { debug!("{} relay from local to remote stream: {}", addr_cloned, err) @@ -176,14 +174,12 @@ impl TcpRelayServer { error!("{} relay from local to remote stream: {}", addr_cloned, err) } } - // remote_stream.close_write().or(Ok(())).unwrap(); - // stream.close_read().or(Ok(())).unwrap(); - }) + remote_stream_cloned.close_write().or(Ok(())).unwrap(); + decrypt_stream.get_mut().get_mut().close_read().or(Ok(())).unwrap(); + } + } }); - // Fixed issue #3 - // io::util::copy(&mut bufr, &mut remote_stream).unwrap(); - let iv = encrypt_method.gen_init_vec(); let encryptor = cipher::with_type(encrypt_method, pwd.as_slice(), @@ -192,25 +188,21 @@ impl TcpRelayServer { try_result!(stream.write(iv.as_slice())); let mut buffered_remote_stream = BufferedStream::new(remote_stream.clone()); let mut encrypt_stream = EncryptedWriter::new(stream.clone(), encryptor); - // let mut remote_cipher = cipher.clone(); - // let remote_addr_clone = addr.clone(); - Thread::spawn(move || { - io::util::copy(&mut buffered_remote_stream, &mut encrypt_stream) - .unwrap_or_else(|err| { - match err.kind { - EndOfFile | BrokenPipe => { - debug!("{} relay from remote to local stream: {}", addr, err) - }, - _ => { - error!("{} relay from remote to local stream: {}", addr, err) - } + match io::util::copy(&mut buffered_remote_stream, &mut encrypt_stream) { + Ok(..) => {}, + Err(err) => { + match err.kind { + EndOfFile | BrokenPipe => { + debug!("{} relay from remote to local stream: {}", addr, err) + }, + _ => { + error!("{} relay from remote to local stream: {}", addr, err) } - // stream.close_write().or(Ok(())).unwrap(); - // remote_stream.close_read().or(Ok(())).unwrap(); - }) - }); - - + } + encrypt_stream.get_mut().close_write().or(Ok(())).unwrap(); + buffered_remote_stream.get_mut().close_read().or(Ok(())).unwrap(); + } + } }); } } diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index 550bff24..df639b07 100644 --- a/src/relay/tcprelay/stream.rs +++ b/src/relay/tcprelay/stream.rs @@ -45,6 +45,28 @@ impl DecryptedReader { sent_final: false, } } + + pub fn get_ref(&self) -> &R { + &self.reader + } + + /// Gets a mutable reference to the underlying reader. + /// + /// # Warning + /// + /// It is inadvisable to read directly from or write directly to the + /// underlying reader. + pub fn get_mut(&mut self) -> &mut R { + &mut self.reader + } + + /// Unwraps this `DecryptedReader`, returning the underlying reader. + /// + /// The internal buffer is flushed before returning the reader. Any leftover + /// data in the read buffer is lost. + pub fn into_inner(self) -> R { + self.reader + } } impl Buffer for DecryptedReader { @@ -140,6 +162,20 @@ impl EncryptedWriter { } } } + + pub fn get_ref(&self) -> &W { + &self.writer + } + + /// Gets a mutable reference to the underlying writer. + /// + /// # Warning + /// + /// It is inadvisable to read directly from or write directly to the + /// underlying writer. + pub fn get_mut(&mut self) -> &mut W { + &mut self.writer + } } impl Writer for EncryptedWriter {