fulfills basic aead protocol

src/crypto/crypto.rs

@@ -21,6 +21,8 @@
 
 //! Cipher defined with Rust-Crypto
 
+use std::mem;
+
 use rust_crypto::symmetriccipher::SynchronousStreamCipher;
 use rust_crypto::chacha20::ChaCha20;
 use rust_crypto::salsa20::Salsa20;
@@ -76,13 +78,23 @@ pub enum CryptoAeadCryptoVariant {
 
 pub struct CryptoAeadCrypto {
     cipher: CryptoAeadCryptoVariant,
+    cipher_type: CipherType,
+    key: Vec<u8>,
+    nounce: Vec<u8>,
 }
 
-
-
 impl CryptoAeadCrypto {
     pub fn new(t: CipherType, key: &[u8], nounce: &[u8]) -> CryptoAeadCrypto {
-        let var = match t {
+        CryptoAeadCrypto {
+            cipher: CryptoAeadCrypto::new_variant(t, key, nounce),
+            cipher_type: t,
+            key: key.to_owned(),
+            nounce: nounce.to_owned(),
+        }
+    }
+
+    fn new_variant(t: CipherType, key: &[u8], nounce: &[u8]) -> CryptoAeadCryptoVariant {
+        match t {
             CipherType::Aes128Gcm => {
                 CryptoAeadCryptoVariant::AesGcm(AesGcm::new(KeySize::KeySize128, key, nounce, &[]))
             }
@@ -94,9 +106,12 @@ impl CryptoAeadCrypto {
             }
 
             _ => panic!("Unsupported {:?}", t),
-        };
+        }
+    }
 
-        CryptoAeadCrypto { cipher: var }
+    fn reset(&mut self) {
+        let var = CryptoAeadCrypto::new_variant(self.cipher_type, &self.key, &self.nounce);
+        mem::replace(&mut self.cipher, var);
     }
 }
 
@@ -104,12 +119,16 @@ impl AeadEncryptor for CryptoAeadCrypto {
     fn encrypt(&mut self, input: &[u8], output: &mut [u8], tag: &mut [u8]) {
         use rust_crypto::aead::AeadEncryptor;
 
-        let CryptoAeadCrypto { ref mut cipher, .. } = *self;
-        match *cipher {
-            CryptoAeadCryptoVariant::AesGcm(ref mut gcm) => {
-                gcm.encrypt(input, output, tag);
+        {
+            let CryptoAeadCrypto { ref mut cipher, .. } = *self;
+            match *cipher {
+                CryptoAeadCryptoVariant::AesGcm(ref mut gcm) => {
+                    gcm.encrypt(input, output, tag);
+                }
             }
         }
+
+        self.reset();
     }
 }
 
@@ -117,16 +136,22 @@ impl AeadDecryptor for CryptoAeadCrypto {
     fn decrypt(&mut self, input: &[u8], output: &mut [u8], tag: &[u8]) -> CipherResult<()> {
         use rust_crypto::aead::AeadDecryptor;
 
-        let CryptoAeadCrypto { ref mut cipher, .. } = *self;
-        match *cipher {
-            CryptoAeadCryptoVariant::AesGcm(ref mut gcm) => {
-                if !gcm.decrypt(input, output, tag) {
-                    Err(Error::AeadDecryptFailed)
-                } else {
-                    Ok(())
+        let r = {
+            let CryptoAeadCrypto { ref mut cipher, .. } = *self;
+            match *cipher {
+                CryptoAeadCryptoVariant::AesGcm(ref mut gcm) => {
+                    if !gcm.decrypt(input, output, tag) {
+                        Err(Error::AeadDecryptFailed)
+                    } else {
+                        Ok(())
+                    }
                 }
             }
-        }
+        };
+
+        self.reset();
+
+        r
     }
 }
 

src/crypto/mod.rs

@@ -26,7 +26,7 @@ use std::convert::From;
 
 use openssl::symm;
 
-pub use self::cipher::{CipherType, CipherResult};
+pub use self::cipher::{CipherType, CipherCategory, CipherResult};
 pub use self::stream::{StreamCipher, StreamCipherVariant, new_stream};
 pub use self::aead::{AeadEncryptor, AeadDecryptor, new_aead_encryptor, new_aead_decryptor};
 

src/relay/tcprelay/aead.rs

@@ -36,7 +36,7 @@ use std::io::{self, Read, Write, BufRead, Cursor};
 use std::cmp;
 use std::u16;
 
-use byteorder::{BigEndian, WriteBytesExt};
+use byteorder::{BigEndian, WriteBytesExt, ReadBytesExt};
 
 use crypto::{self, CipherType, AeadEncryptor, AeadDecryptor};
 
@@ -134,8 +134,12 @@ impl<R> DecryptedReader<R>
                 let mut len_buf = [0u8; 2];
                 self.cipher.decrypt(data, &mut len_buf, tag)?;
 
-                let len = ((len_buf[0] as u16) << 8) | (len_buf[1] as u16);
-                self.read_step = ReadingStep::DataAndTag(u16::from_be(len) as usize);
+                let len = {
+                    let mut cur = Cursor::new(&mut len_buf);
+                    cur.read_u16::<BigEndian>().unwrap() as usize
+                };
+
+                self.read_step = ReadingStep::DataAndTag(len);
             }
             self.buffer.clear();
         }
@@ -167,8 +171,14 @@ impl<R> DecryptedReader<R>
         while !self.sent_final {
             match self.read_step {
                 ReadingStep::DataLength => self.read_length()?,
-                ReadingStep::DataAndTag(dlen) => self.read_data(dlen)?,
-                ReadingStep::DataDone => self.read_step = ReadingStep::DataLength,
+                ReadingStep::DataAndTag(dlen) => {
+                    self.read_data(dlen)?;
+                    break; // Read finished! Break out
+                }
+                ReadingStep::DataDone => {
+                    self.read_step = ReadingStep::DataLength;
+                    self.data.clear();
+                }
             }
         }
         Ok(())
@@ -194,7 +204,7 @@ impl<R> BufRead for DecryptedReader<R>
     }
 
     fn consume(&mut self, amt: usize) {
-        self.pos = cmp::min(self.pos + amt, self.buffer.len());
+        self.pos = cmp::min(self.pos + amt, self.data.len());
     }
 }
 
@@ -270,8 +280,10 @@ impl<W> EncryptedWrite for EncryptedWriter<W>
         buf.reserve(data.len() + self.tag_size);
         buf.resize(orig_buf_len + data.len(), 0);
         self.cipher.encrypt(data, &mut buf[orig_buf_len..], &mut tag_buf);
+
         buf.append(&mut tag_buf);
 
+
         Ok(())
     }
 }

src/relay/tcprelay/crypto_io.rs

@@ -12,7 +12,7 @@ use super::BUFFER_SIZE;
 
 pub trait DecryptedRead: BufRead {}
 
-pub trait EncryptedWrite: Sized {
+pub trait EncryptedWrite {
     /// Writes raw bytes directly to the writer directly
     fn write_raw(&mut self, data: &[u8]) -> io::Result<usize>;
     /// Flush the writer
@@ -21,7 +21,9 @@ pub trait EncryptedWrite: Sized {
     fn encrypt(&mut self, data: &[u8], buf: &mut Vec<u8>) -> io::Result<()>;
 
     /// Encrypt data in `buf` and write all to the writer
-    fn write_all<B: AsRef<[u8]>>(self, buf: B) -> EncryptedWriteAll<Self, B> {
+    fn write_all<B: AsRef<[u8]>>(self, buf: B) -> EncryptedWriteAll<Self, B>
+        where Self: Sized
+    {
         EncryptedWriteAll::Writing {
             writer: self,
             buf: buf,
@@ -32,7 +34,9 @@ pub trait EncryptedWrite: Sized {
     }
 
     /// Copies all data from `r`
-    fn copy<R: Read>(self, r: R) -> EncryptedCopy<R, Self> {
+    fn copy<R: Read>(self, r: R) -> EncryptedCopy<R, Self>
+        where Self: Sized
+    {
         EncryptedCopy {
             reader: r,
             writer: self,
@@ -45,12 +49,16 @@ pub trait EncryptedWrite: Sized {
     }
 
     /// Copies all data from `r` with timeout
-    fn copy_timeout<R: Read>(self, r: R, timeout: Duration, handle: Handle) -> EncryptedCopyTimeout<R, Self> {
+    fn copy_timeout<R: Read>(self, r: R, timeout: Duration, handle: Handle) -> EncryptedCopyTimeout<R, Self>
+        where Self: Sized
+    {
         EncryptedCopyTimeout::new(r, self, timeout, handle)
     }
 
     /// Copies all data from `r` with optional timeout
-    fn copy_timeout_opt<R: Read>(self, r: R, timeout: Option<Duration>, handle: Handle) -> EncryptedCopyOpt<R, Self> {
+    fn copy_timeout_opt<R: Read>(self, r: R, timeout: Option<Duration>, handle: Handle) -> EncryptedCopyOpt<R, Self>
+        where Self: Sized
+    {
         match timeout {
             Some(t) => EncryptedCopyOpt::CopyTimeout(self.copy_timeout(r, t, handle)),
             None => EncryptedCopyOpt::Copy(self.copy(r)),

src/relay/tcprelay/mod.rs

@@ -21,14 +21,14 @@
 
 //! Relay for TCP implementation
 
-use std::io::{self, Read, Write};
+use std::io::{self, Read, Write, BufRead};
 use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
 use std::rc::Rc;
 use std::mem;
 use std::time::Duration;
 
 use crypto;
-use crypto::CryptoMode;
+use crypto::{CryptoMode, CipherCategory};
 use relay::socks5::Address;
 use relay::{BoxIoFuture, boxed_future};
 use relay::dns_resolver::DnsResolver;
@@ -46,9 +46,11 @@ use net2::TcpBuilder;
 
 use ip::IpAddr;
 
-use self::stream::{EncryptedWriter, DecryptedReader};
 pub use self::crypto_io::{DecryptedRead, EncryptedWrite};
 
+use self::stream::{DecryptedReader as StreamDecryptedReader, EncryptedWriter as StreamEncryptedWriter};
+use self::aead::{DecryptedReader as AeadDecryptedReader, EncryptedWriter as AeadEncryptedWriter};
+
 pub mod local;
 mod socks5_local;
 pub mod server;
@@ -68,12 +70,94 @@ pub enum TunnelDirection {
     Server2Client,
 }
 
-// TODO: `DecryptedHalf` and `EncryptedHalf` should be `Box<DecryptedRead>` and `Box<EncryptedWrite>`
+type TcpReadHalf = ReadHalf<TcpStream>;
+type TcpWriteHalf = WriteHalf<TcpStream>;
 
 /// `ReadHalf `of `TcpStream` with decryption
-pub type DecryptedHalf = DecryptedReader<ReadHalf<TcpStream>>;
+pub enum DecryptedHalf {
+    Stream(StreamDecryptedReader<TcpReadHalf>),
+    Aead(AeadDecryptedReader<TcpReadHalf>),
+}
+
+impl Read for DecryptedHalf {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        match *self {
+            DecryptedHalf::Stream(ref mut d) => d.read(buf),
+            DecryptedHalf::Aead(ref mut d) => d.read(buf),
+        }
+    }
+}
+
+impl DecryptedRead for DecryptedHalf {}
+
+impl BufRead for DecryptedHalf {
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+        match *self {
+            DecryptedHalf::Stream(ref mut d) => d.fill_buf(),
+            DecryptedHalf::Aead(ref mut d) => d.fill_buf(),
+        }
+    }
+
+    fn consume(&mut self, amt: usize) {
+        match *self {
+            DecryptedHalf::Stream(ref mut d) => d.consume(amt),
+            DecryptedHalf::Aead(ref mut d) => d.consume(amt),
+        }
+    }
+}
+
+impl From<StreamDecryptedReader<TcpReadHalf>> for DecryptedHalf {
+    fn from(r: StreamDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
+        DecryptedHalf::Stream(r)
+    }
+}
+
+impl From<AeadDecryptedReader<TcpReadHalf>> for DecryptedHalf {
+    fn from(r: AeadDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
+        DecryptedHalf::Aead(r)
+    }
+}
+
 /// `WriteHalf` of `TcpStream` with encryption
-pub type EncryptedHalf = EncryptedWriter<WriteHalf<TcpStream>>;
+pub enum EncryptedHalf {
+    Stream(StreamEncryptedWriter<TcpWriteHalf>),
+    Aead(AeadEncryptedWriter<TcpWriteHalf>),
+}
+
+impl EncryptedWrite for EncryptedHalf {
+    fn write_raw(&mut self, data: &[u8]) -> io::Result<usize> {
+        match *self {
+            EncryptedHalf::Stream(ref mut e) => e.write_raw(data),
+            EncryptedHalf::Aead(ref mut e) => e.write_raw(data),
+        }
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        match *self {
+            EncryptedHalf::Stream(ref mut e) => e.flush(),
+            EncryptedHalf::Aead(ref mut e) => e.flush(),
+        }
+    }
+
+    fn encrypt(&mut self, data: &[u8], buf: &mut Vec<u8>) -> io::Result<()> {
+        match *self {
+            EncryptedHalf::Stream(ref mut e) => e.encrypt(data, buf),
+            EncryptedHalf::Aead(ref mut e) => e.encrypt(data, buf),
+        }
+    }
+}
+
+impl From<StreamEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
+    fn from(d: StreamEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
+        EncryptedHalf::Stream(d)
+    }
+}
+
+impl From<AeadEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
+    fn from(d: AeadEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
+        EncryptedHalf::Aead(d)
+    }
+}
 
 /// Boxed future of `DecryptedHalf`
 pub type DecryptedHalfFut = BoxIoFuture<DecryptedHalf>;
@@ -116,7 +200,10 @@ 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(move |buf| try_timeout(enc_w.write_all(buf), timeout, &handle))
+                .and_then(move |buf| {
+                    trace!("Sending address buffer as {:?}", buf);
+                    try_timeout(enc_w.write_all(buf), timeout, &handle)
+                })
                 .map(|(w, _)| w)
         });
 
@@ -146,14 +233,20 @@ pub fn proxy_handshake(remote_stream: TcpStream,
             trace!("Going to send initialize vector: {:?}", local_iv);
 
             try_timeout(write_all(w, local_iv), timeout, &handle).and_then(move |(w, local_iv)| {
-                // TODO: If crypto type is Aead, returns `aead::EncryptedWriter` instead
+                match svr_cfg.method().category() {
+                    CipherCategory::Stream => {
+                        let encryptor = crypto::new_stream(svr_cfg.method(),
+                                                           svr_cfg.key(),
+                                                           &local_iv[..],
+                                                           CryptoMode::Encrypt);
 
-                let encryptor = crypto::new_stream(svr_cfg.method(),
-                                                   svr_cfg.key(),
-                                                   &local_iv[..],
-                                                   CryptoMode::Encrypt);
-
-                Ok(EncryptedWriter::new(w, encryptor))
+                        Ok(From::from(StreamEncryptedWriter::new(w, encryptor)))
+                    }
+                    CipherCategory::Aead => {
+                        let wtr = AeadEncryptedWriter::new(w, svr_cfg.method(), svr_cfg.key(), &local_iv[..]);
+                        Ok(From::from(wtr))
+                    }
+                }
             })
         };
 
@@ -167,13 +260,21 @@ pub fn proxy_handshake(remote_stream: TcpStream,
 
                 trace!("Got initialize vector {:?}", remote_iv);
 
-                let decryptor = crypto::new_stream(svr_cfg.method(),
-                                                   svr_cfg.key(),
-                                                   &remote_iv[..],
-                                                   CryptoMode::Decrypt);
-                let decrypt_stream = DecryptedReader::new(r, decryptor);
+                match svr_cfg.method().category() {
+                    CipherCategory::Stream => {
+                        let decryptor = crypto::new_stream(svr_cfg.method(),
+                                                           svr_cfg.key(),
+                                                           &remote_iv[..],
+                                                           CryptoMode::Decrypt);
+                        let decrypt_stream = StreamDecryptedReader::new(r, decryptor);
 
-                Ok(decrypt_stream)
+                        Ok(From::from(decrypt_stream))
+                    }
+                    CipherCategory::Aead => {
+                        let dr = AeadDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv[..]);
+                        Ok(From::from(dr))
+                    }
+                }
             })
         };