diff --git a/src/crypto/crypto.rs b/src/crypto/crypto.rs index 780ab93e..7a66bcf2 100644 --- a/src/crypto/crypto.rs +++ b/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, + nounce: Vec, } - - 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 } } diff --git a/src/crypto/mod.rs b/src/crypto/mod.rs index 8e8112f9..117b208a 100644 --- a/src/crypto/mod.rs +++ b/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}; diff --git a/src/relay/tcprelay/aead.rs b/src/relay/tcprelay/aead.rs index 4adce7b8..9f5b2ef7 100644 --- a/src/relay/tcprelay/aead.rs +++ b/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 DecryptedReader 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::().unwrap() as usize + }; + + self.read_step = ReadingStep::DataAndTag(len); } self.buffer.clear(); } @@ -167,8 +171,14 @@ impl DecryptedReader 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 BufRead for DecryptedReader } 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 EncryptedWrite for EncryptedWriter 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(()) } } \ No newline at end of file diff --git a/src/relay/tcprelay/crypto_io.rs b/src/relay/tcprelay/crypto_io.rs index 94afcb56..b02aaa17 100644 --- a/src/relay/tcprelay/crypto_io.rs +++ b/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; /// Flush the writer @@ -21,7 +21,9 @@ pub trait EncryptedWrite: Sized { fn encrypt(&mut self, data: &[u8], buf: &mut Vec) -> io::Result<()>; /// Encrypt data in `buf` and write all to the writer - fn write_all>(self, buf: B) -> EncryptedWriteAll { + fn write_all>(self, buf: B) -> EncryptedWriteAll + where Self: Sized + { EncryptedWriteAll::Writing { writer: self, buf: buf, @@ -32,7 +34,9 @@ pub trait EncryptedWrite: Sized { } /// Copies all data from `r` - fn copy(self, r: R) -> EncryptedCopy { + fn copy(self, r: R) -> EncryptedCopy + 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(self, r: R, timeout: Duration, handle: Handle) -> EncryptedCopyTimeout { + fn copy_timeout(self, r: R, timeout: Duration, handle: Handle) -> EncryptedCopyTimeout + where Self: Sized + { EncryptedCopyTimeout::new(r, self, timeout, handle) } /// Copies all data from `r` with optional timeout - fn copy_timeout_opt(self, r: R, timeout: Option, handle: Handle) -> EncryptedCopyOpt { + fn copy_timeout_opt(self, r: R, timeout: Option, handle: Handle) -> EncryptedCopyOpt + where Self: Sized + { match timeout { Some(t) => EncryptedCopyOpt::CopyTimeout(self.copy_timeout(r, t, handle)), None => EncryptedCopyOpt::Copy(self.copy(r)), diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index a99993dd..1b0a757b 100644 --- a/src/relay/tcprelay/mod.rs +++ b/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` and `Box` +type TcpReadHalf = ReadHalf; +type TcpWriteHalf = WriteHalf; /// `ReadHalf `of `TcpStream` with decryption -pub type DecryptedHalf = DecryptedReader>; +pub enum DecryptedHalf { + Stream(StreamDecryptedReader), + Aead(AeadDecryptedReader), +} + +impl Read for DecryptedHalf { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + 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> for DecryptedHalf { + fn from(r: StreamDecryptedReader) -> DecryptedHalf { + DecryptedHalf::Stream(r) + } +} + +impl From> for DecryptedHalf { + fn from(r: AeadDecryptedReader) -> DecryptedHalf { + DecryptedHalf::Aead(r) + } +} + /// `WriteHalf` of `TcpStream` with encryption -pub type EncryptedHalf = EncryptedWriter>; +pub enum EncryptedHalf { + Stream(StreamEncryptedWriter), + Aead(AeadEncryptedWriter), +} + +impl EncryptedWrite for EncryptedHalf { + fn write_raw(&mut self, data: &[u8]) -> io::Result { + 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) -> io::Result<()> { + match *self { + EncryptedHalf::Stream(ref mut e) => e.encrypt(data, buf), + EncryptedHalf::Aead(ref mut e) => e.encrypt(data, buf), + } + } +} + +impl From> for EncryptedHalf { + fn from(d: StreamEncryptedWriter) -> EncryptedHalf { + EncryptedHalf::Stream(d) + } +} + +impl From> for EncryptedHalf { + fn from(d: AeadEncryptedWriter) -> EncryptedHalf { + EncryptedHalf::Aead(d) + } +} /// Boxed future of `DecryptedHalf` pub type DecryptedHalfFut = BoxIoFuture; @@ -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)) + } + } }) };