From 6f4f75466d7ed7ee4af5cec4246f2b552c44d42b Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Wed, 7 Aug 2019 01:43:13 +0800 Subject: [PATCH] Make BoxStreamCipher as a Boxed dyn StreamCipher trait object --- src/crypto/aead.rs | 4 +- src/crypto/cipher.rs | 2 +- src/crypto/dummy.rs | 4 +- src/crypto/openssl.rs | 12 ++-- src/crypto/rc4_md5.rs | 4 +- src/crypto/sodium.rs | 6 +- src/crypto/stream.rs | 113 ++++---------------------------- src/crypto/table.rs | 8 +-- src/relay/tcprelay/stream.rs | 2 +- src/relay/udprelay/crypto_io.rs | 2 +- 10 files changed, 33 insertions(+), 124 deletions(-) diff --git a/src/crypto/aead.rs b/src/crypto/aead.rs index d45459fe..f3c6c60b 100644 --- a/src/crypto/aead.rs +++ b/src/crypto/aead.rs @@ -34,10 +34,10 @@ pub trait AeadDecryptor { } /// Variant `AeadDecryptor` -pub type BoxAeadDecryptor = Box; +pub type BoxAeadDecryptor = Box; /// Variant `AeadEncryptor` -pub type BoxAeadEncryptor = Box; +pub type BoxAeadEncryptor = Box; /// Generate a specific AEAD cipher encryptor pub fn new_aead_encryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadEncryptor { diff --git a/src/crypto/cipher.rs b/src/crypto/cipher.rs index 4d8e2e17..78c448e6 100644 --- a/src/crypto/cipher.rs +++ b/src/crypto/cipher.rs @@ -790,7 +790,7 @@ impl Display for CipherType { #[cfg(test)] mod test_cipher { - use crate::crypto::{new_stream, CipherType, CryptoMode, StreamCipher}; + use crate::crypto::{new_stream, CipherType, CryptoMode}; #[test] fn test_get_cipher() { diff --git a/src/crypto/dummy.rs b/src/crypto/dummy.rs index 3c44b953..a8687a8e 100644 --- a/src/crypto/dummy.rs +++ b/src/crypto/dummy.rs @@ -10,12 +10,12 @@ use bytes::BufMut; pub struct DummyCipher; impl StreamCipher for DummyCipher { - fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { out.put_slice(data); Ok(()) } - fn finalize(&mut self, _: &mut B) -> CipherResult<()> { + fn finalize(&mut self, _: &mut BufMut) -> CipherResult<()> { Ok(()) } diff --git a/src/crypto/openssl.rs b/src/crypto/openssl.rs index 0158ace1..0b3897cd 100644 --- a/src/crypto/openssl.rs +++ b/src/crypto/openssl.rs @@ -109,7 +109,7 @@ impl OpenSSLCrypto { } /// Update data - pub fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + pub fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { let least_reserved = data.len() + self.cipher.block_size(); let mut buf = BytesMut::with_capacity(least_reserved); // NOTE: len() is 0 now! unsafe { @@ -117,12 +117,12 @@ impl OpenSSLCrypto { } let length = self.inner.update(data, &mut *buf)?; buf.truncate(length); - out.put(buf); + out.put_slice(&buf); Ok(()) } /// Generate the final block - pub fn finalize(&mut self, out: &mut B) -> CipherResult<()> { + pub fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()> { let least_reserved = self.cipher.block_size(); let mut buf = BytesMut::with_capacity(least_reserved); // NOTE: len() is 0 now! unsafe { @@ -131,7 +131,7 @@ impl OpenSSLCrypto { let length = self.inner.finalize(&mut *buf)?; buf.truncate(length); - out.put(buf); + out.put_slice(&buf); Ok(()) } @@ -185,11 +185,11 @@ impl OpenSSLCipher { unsafe impl Send for OpenSSLCipher {} impl StreamCipher for OpenSSLCipher { - fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { self.worker.update(data, out) } - fn finalize(&mut self, out: &mut B) -> CipherResult<()> { + fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()> { self.worker.finalize(out) } diff --git a/src/crypto/rc4_md5.rs b/src/crypto/rc4_md5.rs index 3ccb5098..7367fd8f 100644 --- a/src/crypto/rc4_md5.rs +++ b/src/crypto/rc4_md5.rs @@ -28,11 +28,11 @@ impl Rc4Md5Cipher { } impl StreamCipher for Rc4Md5Cipher { - fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { self.crypto.update(data, out) } - fn finalize(&mut self, out: &mut B) -> CipherResult<()> { + fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()> { self.crypto.finalize(out) } diff --git a/src/crypto/sodium.rs b/src/crypto/sodium.rs index 1872813a..79bebb7c 100644 --- a/src/crypto/sodium.rs +++ b/src/crypto/sodium.rs @@ -56,7 +56,7 @@ impl SodiumStreamCipher { self.counter % SODIUM_BLOCK_SIZE } - fn process(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + fn process(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { let padding = self.padding_len(); let mut plain_text = vec![0u8; data.len() + padding]; @@ -140,11 +140,11 @@ fn crypto_stream_xor_ic( } impl StreamCipher for SodiumStreamCipher { - fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { self.process(data, out) } - fn finalize(&mut self, _: &mut B) -> CipherResult<()> { + fn finalize(&mut self, _: &mut BufMut) -> CipherResult<()> { Ok(()) } diff --git a/src/crypto/stream.rs b/src/crypto/stream.rs index ef29b099..ebf6b8fc 100644 --- a/src/crypto/stream.rs +++ b/src/crypto/stream.rs @@ -1,7 +1,5 @@ //! Stream ciphers -use std::ops::{Deref, DerefMut}; - #[cfg(feature = "openssl")] use crate::crypto::openssl; #[cfg(feature = "rc4")] @@ -20,99 +18,13 @@ use bytes::BufMut; /// The `update` method could be called multiple times, and the `finalize` method will /// encrypt the last block pub trait StreamCipher { - fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()>; - fn finalize(&mut self, out: &mut B) -> CipherResult<()>; + fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()>; + fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()>; fn buffer_size(&self, data: &[u8]) -> usize; } -macro_rules! define_stream_ciphers { - ($( $(#[$attr:meta])* pub $name:ident => $cipher:ty, )+) => { - /// Variant cipher which contains all possible ciphers - pub enum StreamCipherVariant { - $( - $(#[$attr])* - $name($cipher), - )+ - } - - impl StreamCipherVariant { - /// Creates from an actual cipher - pub fn new(cipher: C) -> StreamCipherVariant - where StreamCipherVariant: From - { - From::from(cipher) - } - } - - impl StreamCipher for StreamCipherVariant { - fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { - match *self { - $( - $(#[$attr])* - StreamCipherVariant::$name(ref mut cipher) => cipher.update(data, out), - )+ - } - } - - fn finalize(&mut self, out: &mut B) -> CipherResult<()> { - match *self { - $( - $(#[$attr])* - StreamCipherVariant::$name(ref mut cipher) => cipher.finalize(out), - )+ - } - } - - fn buffer_size(&self, data: &[u8]) -> usize { - match *self { - $( - $(#[$attr])* - StreamCipherVariant::$name(ref cipher) => cipher.buffer_size(data), - )+ - } - } - } - - $( - $(#[$attr])* - impl From<$cipher> for StreamCipherVariant { - fn from(cipher: $cipher) -> StreamCipherVariant { - StreamCipherVariant::$name(cipher) - } - } - )+ - } -} - /// Variant cipher which contains all possible stream ciphers -pub struct BoxStreamCipher { - cipher: Box, -} - -impl Deref for BoxStreamCipher { - type Target = StreamCipherVariant; - - fn deref(&self) -> &StreamCipherVariant { - &*self.cipher - } -} - -impl DerefMut for BoxStreamCipher { - fn deref_mut(&mut self) -> &mut StreamCipherVariant { - &mut *self.cipher - } -} - -define_stream_ciphers! { - pub TableCipher => table::TableCipher, - pub DummyCipher => dummy::DummyCipher, - #[cfg(feature = "rc4")] - pub Rc4Md5Cipher => rc4_md5::Rc4Md5Cipher, - #[cfg(feature = "openssl")] - pub OpenSSLCipher => openssl::OpenSSLCipher, - #[cfg(feature = "sodium")] - pub SodiumStreamCipher => sodium::SodiumStreamCipher, -} +pub type BoxStreamCipher = Box; /// Generate a specific Cipher with key and initialize vector #[allow(unused_variables)] @@ -122,17 +34,17 @@ pub fn new_stream(t: CipherType, key: &[u8], iv: &[u8], mode: CryptoMode) -> Box "only allow initializing with stream cipher" ); - let cipher = match t { - CipherType::Table => StreamCipherVariant::new(table::TableCipher::new(key, mode)), - CipherType::Plain => StreamCipherVariant::new(dummy::DummyCipher), + match t { + CipherType::Table => Box::new(table::TableCipher::new(key, mode)), + CipherType::Plain => Box::new(dummy::DummyCipher), #[cfg(feature = "sodium")] CipherType::ChaCha20 | CipherType::Salsa20 | CipherType::XSalsa20 | CipherType::ChaCha20Ietf => { - StreamCipherVariant::new(sodium::SodiumStreamCipher::new(t, key, iv)) + Box::new(sodium::SodiumStreamCipher::new(t, key, iv)) } #[cfg(feature = "rc4")] - CipherType::Rc4Md5 => StreamCipherVariant::new(rc4_md5::Rc4Md5Cipher::new(key, iv, mode)), + CipherType::Rc4Md5 => Box::new(rc4_md5::Rc4Md5Cipher::new(key, iv, mode)), #[cfg(feature = "aes-cfb")] CipherType::Aes128Cfb @@ -146,11 +58,11 @@ pub fn new_stream(t: CipherType, key: &[u8], iv: &[u8], mode: CryptoMode) -> Box | CipherType::Aes256Cfb | CipherType::Aes256Cfb1 | CipherType::Aes256Cfb8 - | CipherType::Aes256Cfb128 => StreamCipherVariant::new(openssl::OpenSSLCipher::new(t, key, iv, mode)), + | CipherType::Aes256Cfb128 => Box::new(openssl::OpenSSLCipher::new(t, key, iv, mode)), #[cfg(feature = "aes-ctr")] CipherType::Aes128Ctr | CipherType::Aes192Ctr | CipherType::Aes256Ctr => { - StreamCipherVariant::new(openssl::OpenSSLCipher::new(t, key, iv, mode)) + Box::new(openssl::OpenSSLCipher::new(t, key, iv, mode)) } #[cfg(feature = "camellia-cfb")] @@ -165,11 +77,8 @@ pub fn new_stream(t: CipherType, key: &[u8], iv: &[u8], mode: CryptoMode) -> Box | CipherType::Camellia256Cfb | CipherType::Camellia256Cfb1 | CipherType::Camellia256Cfb8 - | CipherType::Camellia256Cfb128 => StreamCipherVariant::new(openssl::OpenSSLCipher::new(t, key, iv, mode)), + | CipherType::Camellia256Cfb128 => Box::new(openssl::OpenSSLCipher::new(t, key, iv, mode)), _ => unreachable!("{} is not a stream cipher", t), - }; - BoxStreamCipher { - cipher: Box::new(cipher), } } diff --git a/src/crypto/table.rs b/src/crypto/table.rs index a0507fab..eae0ce68 100644 --- a/src/crypto/table.rs +++ b/src/crypto/table.rs @@ -56,7 +56,7 @@ impl TableCipher { } } - fn process(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + fn process(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { let mut buf = BytesMut::with_capacity(self.buffer_size(data)); unsafe { buf.set_len(self.buffer_size(data)); // Set length @@ -64,17 +64,17 @@ impl TableCipher { for (idx, d) in data.iter().enumerate() { buf[idx] = self.table[*d as usize]; } - out.put(buf); + out.put_slice(&buf); Ok(()) } } impl StreamCipher for TableCipher { - fn update(&mut self, data: &[u8], out: &mut B) -> CipherResult<()> { + fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> { self.process(data, out) } - fn finalize(&mut self, _: &mut B) -> CipherResult<()> { + fn finalize(&mut self, _: &mut BufMut) -> CipherResult<()> { Ok(()) } diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index 6dfda181..999ea405 100644 --- a/src/relay/tcprelay/stream.rs +++ b/src/relay/tcprelay/stream.rs @@ -5,7 +5,7 @@ use std::{ io::{self, BufRead, Read}, }; -use crate::crypto::{new_stream, BoxStreamCipher, CipherType, CryptoMode, StreamCipher}; +use crate::crypto::{new_stream, BoxStreamCipher, CipherType, CryptoMode}; use bytes::{BufMut, BytesMut}; use tokio_io::{AsyncRead, AsyncWrite}; diff --git a/src/relay/udprelay/crypto_io.rs b/src/relay/udprelay/crypto_io.rs index 6b1eaa43..fc58d711 100644 --- a/src/relay/udprelay/crypto_io.rs +++ b/src/relay/udprelay/crypto_io.rs @@ -22,7 +22,7 @@ use std::io; -use crate::crypto::{self, CipherCategory, CipherType, CryptoMode, StreamCipher}; +use crate::crypto::{self, CipherCategory, CipherType, CryptoMode}; /// Encrypt payload into ShadowSocks UDP encrypted packet pub fn encrypt_payload(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result> {