diff --git a/src/crypto/cipher.rs b/src/crypto/cipher.rs index 84a1e538..f98a1a51 100644 --- a/src/crypto/cipher.rs +++ b/src/crypto/cipher.rs @@ -39,8 +39,8 @@ use crypto::digest::{self, DigestType}; /// The `update` method could be called multiple times, and the `finalize` method will /// encrypt the last block pub trait Cipher { - fn update(&mut self, data: &[u8]) -> CipherResult>; - fn finalize(&mut self) -> CipherResult>; + fn update(&mut self, data: &[u8], out: &mut Vec) -> CipherResult<()>; + fn finalize(&mut self, out: &mut Vec) -> CipherResult<()>; } pub type CipherResult = Result; @@ -482,8 +482,10 @@ mod test_cipher { let mut decryptor = with_type(CipherType::Aes128Cfb, &key[0..], &iv[0..], CryptoMode::Decrypt); let message = "HELLO WORLD"; - let encrypted_msg = encryptor.update(message.as_bytes()).unwrap(); - let decrypted_msg = decryptor.update(&encrypted_msg[..]).unwrap(); + let mut encrypted_msg = Vec::new(); + encryptor.update(message.as_bytes(), &mut encrypted_msg).unwrap(); + let mut decrypted_msg = Vec::new(); + decryptor.update(&encrypted_msg[..], &mut decrypted_msg).unwrap(); assert!(message.as_bytes() == &decrypted_msg[..]); } diff --git a/src/crypto/openssl.rs b/src/crypto/openssl.rs index 04a83283..6169974b 100644 --- a/src/crypto/openssl.rs +++ b/src/crypto/openssl.rs @@ -337,15 +337,15 @@ impl OpenSSLCrypto { } } - pub fn update(&self, data: &[u8]) -> CipherResult> { + pub fn update(&self, data: &[u8], out: &mut Vec) -> CipherResult<()> { let pdata: *const u8 = data.as_ptr(); let datalen: libc::c_int = data.len() as libc::c_int; - let mut out = Vec::with_capacity(data.len() + self.cipher_type.block_size()); - unsafe { out.set_len(data.len() + self.cipher_type.block_size()); } + out.reserve(data.len() + self.cipher_type.block_size()); let mut len: libc::c_int = 0; - let pres: *mut u8 = out.as_mut_ptr(); + let orig_len = out.len(); + let pres: *mut u8 = out[orig_len..].as_mut_ptr(); unsafe { if ffi::EVP_CipherUpdate(self.evp_ctx, @@ -360,14 +360,14 @@ impl OpenSSLCrypto { } } - out.truncate(len as usize); - Ok(out) + unsafe { out.set_len(orig_len + len as usize) } + Ok(()) } - pub fn finalize(&self) -> CipherResult> { + pub fn finalize(&self, out: &mut Vec) -> CipherResult<()> { let mut len: libc::c_int = 0; - let mut out = Vec::with_capacity(self.cipher_type.block_size()); - unsafe { out.set_len(self.cipher_type.block_size()) } + let orig_len = out.len(); + out.reserve(self.cipher_type.block_size()); unsafe { if ffi::EVP_CipherFinal(self.evp_ctx, out.as_mut_ptr(), &mut len) != 1 { @@ -379,8 +379,8 @@ impl OpenSSLCrypto { } } - out.truncate(len as usize); - Ok(out) + unsafe { out.set_len(orig_len + len as usize) } + Ok(()) } } @@ -433,10 +433,12 @@ impl Drop for OpenSSLCrypto { /// let mut enc = OpenSSLCipher::new(cipher::CipherType::Aes128Cfb, &key[0..], &iv[0..], CryptoMode::Encrypt); /// let mut dec = OpenSSLCipher::new(cipher::CipherType::Aes128Cfb, &key[0..], &iv[0..], CryptoMode::Decrypt); /// let message = "hello world"; -/// let encrypted_message = enc.update(message.as_bytes()).unwrap(); -/// let decrypted_message = dec.update(&encrypted_message[]).unwrap(); +/// let mut encrypted_message = Vec::new(); +/// enc.update(message.as_bytes(), &mut encrypted_message).unwrap(); +/// let mut decrypted_message = Vec::new(); +/// dec.update(&encrypted_message[..], &mut decrypted_message).unwrap(); /// -/// assert!(&decrypted_message[] == message.as_bytes()); +/// assert!(&decrypted_message[..] == message.as_bytes()); /// ``` #[derive(Clone)] pub struct OpenSSLCipher { @@ -452,12 +454,12 @@ impl OpenSSLCipher { } impl Cipher for OpenSSLCipher { - fn update(&mut self, data: &[u8]) -> CipherResult> { - self.worker.update(data) + fn update(&mut self, data: &[u8], out: &mut Vec) -> CipherResult<()> { + self.worker.update(data, out) } - fn finalize(&mut self) -> CipherResult> { - self.worker.finalize() + fn finalize(&mut self, out: &mut Vec) -> CipherResult<()> { + self.worker.finalize(out) } } @@ -511,13 +513,15 @@ mod test_openssl { let mut enc = OpenSSLCipher::new(*t, &k[..], &iv[..], CryptoMode::Encrypt); - let mut encrypted_msg = enc.update(message.as_bytes()).unwrap(); - encrypted_msg.push_all(&enc.finalize().unwrap()[..]); + let mut encrypted_msg = Vec::new(); + enc.update(message.as_bytes(), &mut encrypted_msg).unwrap(); + enc.finalize(&mut encrypted_msg).unwrap(); println!("ENC {:?}", encrypted_msg); let mut dec = OpenSSLCipher::new(*t, &k[..], &iv[..], CryptoMode::Decrypt); - let mut decrypted_msg = dec.update(&encrypted_msg[..]).unwrap(); - decrypted_msg.push_all(&dec.finalize().unwrap()[..]); + let mut decrypted_msg = Vec::new(); + dec.update(&encrypted_msg[..], &mut decrypted_msg).unwrap(); + dec.finalize(&mut decrypted_msg).unwrap(); println!("DEC {:?}", &decrypted_msg[..]); assert_eq!(message.as_bytes(), &decrypted_msg[..]); @@ -526,14 +530,14 @@ mod test_openssl { #[bench] fn bench_openssl_default_cipher_encrypt(b: &mut test::Bencher) { - use std::rand::random; + use rand::random; let msg_size: usize = 0xffff; let mut test_data = Vec::new(); - for _ in range::(0, 100) { - let msg = range(0, msg_size).map(|_| random::()).collect::>(); - let key = range(1, random::() % 63).map(|_| random::()).collect::>(); + for _ in 0..100 { + let msg = (0..msg_size).map(|_| random::()).collect::>(); + let key = (1..random::() % 63).map(|_| random::()).collect::>(); let k = cipher::CipherType::Aes256Cfb.bytes_to_key(&key[..]); let v = cipher::CipherType::Aes256Cfb.gen_init_vec(); @@ -545,25 +549,27 @@ mod test_openssl { let mut enc = OpenSSLCipher::new(cipher::CipherType::Aes256Cfb, &key[..], &iv[..], CryptoMode::Encrypt); - enc.update(&msg[..]).unwrap(); + let mut out = Vec::new(); + enc.update(&msg[..], &mut out).unwrap(); }); b.bytes = msg_size as u64; } #[bench] fn bench_openssl_default_cipher_decrypt(b: &mut test::Bencher) { - use std::rand::random; + use rand::random; let msg_size: usize = 0xffff; let mut test_data = Vec::new(); - for _ in range::(0, 100) { - let msg = range(0, msg_size).map(|_| random::()).collect::>(); - let key = range(1, random::() % 63).map(|_| random::()).collect::>(); + for _ in 0..100 { + let msg = (0..msg_size).map(|_| random::()).collect::>(); + let key = (1..random::() % 63).map(|_| random::()).collect::>(); let k = cipher::CipherType::Aes256Cfb.bytes_to_key(&key[..]); let v = cipher::CipherType::Aes256Cfb.gen_init_vec(); let mut cipher = OpenSSLCipher::new(cipher::CipherType::Aes256Cfb, &k[..], &v[..], CryptoMode::Encrypt); - let encrypted_msg = cipher.update(&msg[..]).unwrap(); + let mut encrypted_msg = Vec::new(); + cipher.update(&msg[..], &mut encrypted_msg).unwrap(); test_data.push((k, v, encrypted_msg)); } @@ -571,7 +577,8 @@ mod test_openssl { let (ref key, ref iv, ref encrypted_msg) = test_data[random::() % test_data.len()]; let mut cipher = OpenSSLCipher::new(cipher::CipherType::Aes256Cfb, &key[..], &iv[..], CryptoMode::Decrypt); - cipher.update(&encrypted_msg[..]).unwrap(); + let mut out = Vec::new(); + cipher.update(&encrypted_msg[..], &mut out).unwrap(); }); b.bytes = msg_size as u64; } diff --git a/src/crypto/rc4_md5.rs b/src/crypto/rc4_md5.rs index bc6b3aa3..0c75d5d9 100644 --- a/src/crypto/rc4_md5.rs +++ b/src/crypto/rc4_md5.rs @@ -45,12 +45,12 @@ impl Rc4Md5Cipher { } impl Cipher for Rc4Md5Cipher { - fn update(&mut self, data: &[u8]) -> CipherResult> { - self.crypto.update(data) + fn update(&mut self, data: &[u8], out: &mut Vec) -> CipherResult<()> { + self.crypto.update(data, out) } - fn finalize(&mut self) -> CipherResult> { - self.crypto.finalize() + fn finalize(&mut self, out: &mut Vec) -> CipherResult<()> { + self.crypto.finalize(out) } } diff --git a/src/crypto/sodium.rs b/src/crypto/sodium.rs index a5401545..5b6adab6 100644 --- a/src/crypto/sodium.rs +++ b/src/crypto/sodium.rs @@ -60,7 +60,7 @@ impl SodiumCipher { } impl Cipher for SodiumCipher { - fn update(&mut self, data: &[u8]) -> CipherResult> { + fn update(&mut self, data: &[u8], out: &mut Vec) -> CipherResult<()> { let padding_len = self.counter % BLOCK_SIZE; let mut pad; let padded_data = @@ -95,11 +95,13 @@ impl Cipher for SodiumCipher { self.counter += data.len(); - Ok(self.buf[padding_len..padding_len + data.len()].to_vec()) + out.push_all(&self.buf[padding_len..padding_len + data.len()]); + + Ok(()) } - fn finalize(&mut self) -> CipherResult> { - Ok(Vec::new()) + fn finalize(&mut self, _: &mut Vec) -> CipherResult<()> { + Ok(()) } } @@ -118,10 +120,12 @@ mod test_sodium { let iv = ct.gen_init_vec(); let mut enc = SodiumCipher::new(ct, &key[..], &iv[..]); - let encrypted_msg = enc.update(message).unwrap(); + let mut encrypted_msg = Vec::new(); + enc.update(message, &mut encrypted_msg).unwrap(); let mut dec = SodiumCipher::new(ct, &key[..], &iv[..]); - let decrypted_msg = dec.update(&encrypted_msg[0..]).unwrap(); + let mut decrypted_msg = Vec::new(); + dec.update(&encrypted_msg[0..], &mut decrypted_msg).unwrap(); assert_eq!(message, &decrypted_msg[..]); } diff --git a/src/crypto/table.rs b/src/crypto/table.rs index b9035202..54e007a1 100644 --- a/src/crypto/table.rs +++ b/src/crypto/table.rs @@ -69,19 +69,22 @@ impl TableCipher { } } - fn process(&mut self, data: &[u8]) -> CipherResult> { - let r = data.iter().map(|d| self.table[*d as usize]).collect(); - Ok(r) + fn process(&mut self, data: &[u8], out: &mut Vec) -> CipherResult<()> { + out.reserve(data.len()); + for d in data.iter() { + out.push(self.table[*d as usize]); + } + Ok(()) } } impl Cipher for TableCipher { - fn update(&mut self, data: &[u8]) -> CipherResult> { - self.process(data) + fn update(&mut self, data: &[u8], out: &mut Vec) -> CipherResult<()> { + self.process(data, out) } - fn finalize(&mut self) -> CipherResult> { - Ok(Vec::new()) + fn finalize(&mut self, _: &mut Vec) -> CipherResult<()> { + Ok(()) } } @@ -92,8 +95,10 @@ fn test_table_cipher() { let mut enc = TableCipher::new(key.as_bytes(), CryptoMode::Encrypt); let mut dec = TableCipher::new(key.as_bytes(), CryptoMode::Decrypt); - let encrypted_msg = enc.update(message.as_bytes()).unwrap(); - let decrypted_msg = dec.update(&encrypted_msg[..]).unwrap(); + let mut encrypted_msg = Vec::new(); + enc.update(message.as_bytes(), &mut encrypted_msg).unwrap(); + let mut decrypted_msg = Vec::new(); + dec.update(&encrypted_msg[..], &mut decrypted_msg).unwrap(); assert_eq!(&decrypted_msg[..], message.as_bytes()); } diff --git a/src/relay/local.rs b/src/relay/local.rs index c0a9c803..f3ee4a1d 100644 --- a/src/relay/local.rs +++ b/src/relay/local.rs @@ -34,18 +34,15 @@ use config::Config; /// UDP Associate and Bind commands are not supported currently. /// /// ```no_run -/// use std::io::net::ip::SocketAddr; +/// use std::net::SocketAddr; /// /// use shadowsocks::relay::Relay; /// use shadowsocks::relay::RelayLocal; -/// use shadowsocks::config::{Config, ClientConfig, ServerConfig}; +/// use shadowsocks::config::{Config, ServerConfig}; /// use shadowsocks::crypto::cipher::CipherType; /// /// let mut config = Config::new(); -/// config.local = Some(ClientConfig { -/// ip: "127.0.0.1".parse().unwrap(), -/// port: 1080 -/// }); +/// config.local = Some(SocketAddr::new("127.0.0.1".parse().unwrap(), 1080)); /// config.server = vec![ServerConfig { /// addr: "127.0.0.1".to_string(), /// port: 8388, diff --git a/src/relay/server.rs b/src/relay/server.rs index b03f8269..ddd1bc73 100644 --- a/src/relay/server.rs +++ b/src/relay/server.rs @@ -34,7 +34,7 @@ use config::Config; /// UDP Associate and Bind commands are not supported currently. /// /// ```no_run -/// use std::io::net::ip::SocketAddr; +/// use std::net::SocketAddr; /// /// use shadowsocks::relay::Relay; /// use shadowsocks::relay::RelayServer; diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index f73cdd90..ca1c8147 100644 --- a/src/relay/tcprelay/stream.rs +++ b/src/relay/tcprelay/stream.rs @@ -78,18 +78,18 @@ impl BufRead for DecryptedReader { match self.reader.read(&mut incoming) { Ok(0) => { // EOF - self.buffer = try!(self.cipher - .finalize() - .map_err(|err| io::Error::new(io::ErrorKind::Other, - err.desc, - err.detail))); + try!(self.cipher + .finalize(&mut self.buffer) + .map_err(|err| io::Error::new(io::ErrorKind::Other, + err.desc, + err.detail))); } Ok(l) => { - self.buffer = try!(self.cipher - .update(&incoming[0..l]) - .map_err(|err| io::Error::new(io::ErrorKind::Other, - err.desc, - err.detail))); + try!(self.cipher + .update(&incoming[0..l], &mut self.buffer) + .map_err(|err| io::Error::new(io::ErrorKind::Other, + err.desc, + err.detail))); }, Err(err) => { return Err(err); @@ -124,6 +124,7 @@ impl Read for DecryptedReader { pub struct EncryptedWriter { writer: W, cipher: Box, + buffer: Vec, } impl EncryptedWriter { @@ -131,13 +132,16 @@ impl EncryptedWriter { EncryptedWriter { writer: w, cipher: cipher, + buffer: Vec::new(), } } pub fn finalize(&mut self) -> io::Result<()> { - match self.cipher.finalize() { - Ok(fin) => { - self.writer.write_all(&fin[..]) + match self.cipher.finalize(&mut self.buffer) { + Ok(..) => { + try!(self.writer.write_all(&self.buffer[..])); + self.buffer.clear(); + Ok(()) }, Err(err) => { Err(io::Error::new( @@ -166,10 +170,13 @@ impl EncryptedWriter { impl Write for EncryptedWriter { fn write(&mut self, buf: &[u8]) -> io::Result { - match self.cipher.update(buf) { - Ok(ret) => { - match self.writer.write_all(ret.as_slice()) { - Ok(..) => Ok(buf.len()), + match self.cipher.update(buf, &mut self.buffer) { + Ok(..) => { + match self.writer.write_all(&self.buffer[..]) { + Ok(..) => { + self.buffer.clear(); + Ok(buf.len()) + }, Err(err) => return Err(err), } }, diff --git a/src/relay/udprelay/local.rs b/src/relay/udprelay/local.rs index cc57ad0a..51fb6588 100644 --- a/src/relay/udprelay/local.rs +++ b/src/relay/udprelay/local.rs @@ -204,8 +204,8 @@ fn handle_request(socket: UdpSocket, request.write_to(&mut wbuf).unwrap(); io::copy(&mut bufr, &mut wbuf).unwrap(); - iv.push_all(&encryptor.update(&wbuf[..]).unwrap()[..]); - iv.push_all(&encryptor.finalize().unwrap()[..]); + encryptor.update(&wbuf[..], &mut iv).unwrap(); + encryptor.finalize(&mut iv).unwrap(); socket.send_to(&iv[..], &server_addr) .ok().expect("Error occurs while sending to remote"); @@ -222,8 +222,9 @@ fn handle_response(socket: UdpSocket, &key[..], &response_message[0..config.method.block_size()], CryptoMode::Decrypt); - let mut decrypted_data = decryptor.update(&response_message[config.method.block_size()..]).unwrap(); - decrypted_data.push_all(&decryptor.finalize().unwrap()[..]); + let mut decrypted_data = Vec::new(); + decryptor.update(&response_message[config.method.block_size()..], &mut decrypted_data).unwrap(); + decryptor.finalize(&mut decrypted_data).unwrap(); let mut bufr = BufReader::new(&decrypted_data[..]); diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index 7f41e9a4..58700d66 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -86,8 +86,8 @@ impl UdpRelayServer { &key[..], &iv[..], CryptoMode::Encrypt); - iv.push_all(&encryptor.update(&response_buf[..]).unwrap()[..]); - iv.push_all(&encryptor.finalize().unwrap()[..]); + encryptor.update(&response_buf[..], &mut iv).unwrap(); + encryptor.finalize(&mut iv).unwrap(); captured_socket .send_to(&iv[..], &client_addr) @@ -111,8 +111,9 @@ impl UdpRelayServer { &key[..], &data[0..method.block_size()], CryptoMode::Decrypt); - let mut decrypted_data = decryptor.update(&data[method.block_size()..]).unwrap(); - decrypted_data.push_all(&decryptor.finalize().unwrap()[..]); + let mut decrypted_data = Vec::new(); + decryptor.update(&data[method.block_size()..], &mut decrypted_data).unwrap(); + &decryptor.finalize(&mut decrypted_data).unwrap(); let mut bufr = BufReader::new(&decrypted_data[..]); let header = socks5::UdpAssociateHeader::read_from(bufr.by_ref()).unwrap();