reuse thread

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");

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();
+                    }
+                }
             });
         }
     }

src/relay/tcprelay/stream.rs

@@ -45,6 +45,28 @@ impl<R: Reader> DecryptedReader<R> {
             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<R: Reader> Buffer for DecryptedReader<R> {
@@ -140,6 +162,20 @@ impl<W: Writer> EncryptedWriter<W> {
             }
         }
     }
+
+    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<W: Writer> Writer for EncryptedWriter<W> {