From ae787715d528ac7992d43afdab52a7d1b82f1bc5 Mon Sep 17 00:00:00 2001
From: "Y. T. Chung" <zonyitoo@gmail.com>
Date: Sat, 18 Feb 2017 16:49:36 +0800
Subject: [PATCH] construct cipher inside stream wrapper

---
 src/relay/tcprelay/mod.rs    | 19 ++++---------------
 src/relay/tcprelay/stream.rs | 10 +++++-----
 2 files changed, 9 insertions(+), 20 deletions(-)

diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs
index d63b18d6..c23a2d5b 100644
--- a/src/relay/tcprelay/mod.rs
+++ b/src/relay/tcprelay/mod.rs
@@ -27,8 +27,7 @@ use std::rc::Rc;
 use std::mem;
 use std::time::Duration;
 
-use crypto;
-use crypto::{CryptoMode, CipherCategory};
+use crypto::CipherCategory;
 use relay::socks5::Address;
 use relay::{BoxIoFuture, boxed_future};
 use relay::dns_resolver::DnsResolver;
@@ -248,12 +247,7 @@ pub fn proxy_handshake(remote_stream: TcpStream,
                 match svr_cfg.method().category() {
                     CipherCategory::Stream => {
                         let local_iv = prev_buf;
-                        let encryptor = crypto::new_stream(svr_cfg.method(),
-                                                           svr_cfg.key(),
-                                                           &local_iv[..],
-                                                           CryptoMode::Encrypt);
-
-                        Ok(From::from(StreamEncryptedWriter::new(w, encryptor)))
+                        Ok(From::from(StreamEncryptedWriter::new(w, svr_cfg.method(), svr_cfg.key(), &local_iv)))
                     }
                     CipherCategory::Aead => {
                         let local_salt = prev_buf;
@@ -282,16 +276,11 @@ pub fn proxy_handshake(remote_stream: TcpStream,
                 match svr_cfg.method().category() {
                     CipherCategory::Stream => {
                         trace!("Got initialize vector  {:?}", remote_iv);
-
-                        let decryptor = crypto::new_stream(svr_cfg.method(),
-                                                           svr_cfg.key(),
-                                                           &remote_iv[..],
-                                                           CryptoMode::Decrypt);
-                        let decrypt_stream = StreamDecryptedReader::new(r, decryptor);
-
+                        let decrypt_stream = StreamDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv);
                         Ok(From::from(decrypt_stream))
                     }
                     CipherCategory::Aead => {
+                        trace!("Got salt {:?}", remote_iv);
                         let dr = AeadDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv[..]);
                         Ok(From::from(dr))
                     }
diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs
index d6265f64..cd0963f9 100644
--- a/src/relay/tcprelay/stream.rs
+++ b/src/relay/tcprelay/stream.rs
@@ -22,7 +22,7 @@
 use std::io::{self, Read, BufRead, Write};
 use std::cmp;
 
-use crypto::{StreamCipher, StreamCipherVariant};
+use crypto::{CipherType, StreamCipher, StreamCipherVariant, CryptoMode, new_stream};
 
 use super::BUFFER_SIZE;
 use super::{EncryptedWrite, DecryptedRead};
@@ -41,11 +41,11 @@ pub struct DecryptedReader<R>
 impl<R> DecryptedReader<R>
     where R: Read
 {
-    pub fn new(r: R, cipher: StreamCipherVariant) -> DecryptedReader<R> {
+    pub fn new(r: R, t: CipherType, key: &[u8], iv: &[u8]) -> DecryptedReader<R> {
         DecryptedReader {
             reader: r,
             buffer: Vec::new(),
-            cipher: cipher,
+            cipher: new_stream(t, key, iv, CryptoMode::Decrypt),
             pos: 0,
             sent_final: false,
         }
@@ -139,10 +139,10 @@ impl<W> EncryptedWriter<W>
     where W: Write
 {
     /// Creates a new EncryptedWriter
-    pub fn new(w: W, cipher: StreamCipherVariant) -> EncryptedWriter<W> {
+    pub fn new(w: W, t: CipherType, key: &[u8], iv: &[u8]) -> EncryptedWriter<W> {
         EncryptedWriter {
             writer: w,
-            cipher: cipher,
+            cipher: new_stream(t, key, iv, CryptoMode::Encrypt),
         }
     }