shadowsocks/
cipher.rs

1use aes::{Aes128, Aes192, Aes256};
2use cfb_mode::Cfb;
3use chacha20::ChaCha20;
4use cipher::{consts::U16, AsyncStreamCipher, BlockCipher, BlockEncrypt, NewCipher, StreamCipher};
5use ctr::Ctr128BE;
6use rand::distributions::Standard;
7use rand::{thread_rng, Rng};
8
9use crate::util::generate_key;
10
11pub trait SymmetricCipher {
12    /// Encrypt data in place.
13    fn encrypt(&mut self, data: &mut [u8]);
14
15    /// Decrypt data in place.
16    fn decrypt(&mut self, data: &mut [u8]);
17}
18
19impl<C: BlockCipher + BlockEncrypt> SymmetricCipher for Cfb<C> {
20    /// Encrypt data in place.
21    fn encrypt(&mut self, data: &mut [u8]) {
22        AsyncStreamCipher::encrypt(self, data)
23    }
24
25    /// Decrypt data in place.
26    fn decrypt(&mut self, data: &mut [u8]) {
27        AsyncStreamCipher::decrypt(self, data)
28    }
29}
30
31impl<B: BlockEncrypt + BlockCipher<BlockSize = U16>> SymmetricCipher for Ctr128BE<B> {
32    /// Encrypt data in place.
33    fn encrypt(&mut self, data: &mut [u8]) {
34        StreamCipher::apply_keystream(self, data)
35    }
36
37    /// Decrypt data in place.
38    fn decrypt(&mut self, data: &mut [u8]) {
39        StreamCipher::apply_keystream(self, data)
40    }
41}
42
43impl SymmetricCipher for ChaCha20 {
44    /// Encrypt data in place.
45    fn encrypt(&mut self, data: &mut [u8]) {
46        StreamCipher::apply_keystream(self, data)
47    }
48
49    /// Decrypt data in place.
50    fn decrypt(&mut self, data: &mut [u8]) {
51        StreamCipher::apply_keystream(self, data)
52    }
53}
54
55type Aes128Cfb = Cfb<Aes128>;
56type Aes192Cfb = Cfb<Aes192>;
57type Aes256Cfb = Cfb<Aes256>;
58type Aes128Ctr = Ctr128BE<Aes128>;
59type Aes192Ctr = Ctr128BE<Aes192>;
60type Aes256Ctr = Ctr128BE<Aes256>;
61
62pub struct Cipher {
63    pub key: Vec<u8>,
64    pub key_len: usize,
65    pub iv: Vec<u8>,
66    pub iv_len: usize,
67    pub enc: Option<Box<dyn SymmetricCipher + Send + 'static>>,
68    pub dec: Option<Box<dyn SymmetricCipher + Send + 'static>>,
69    cipher_method: CipherMethod,
70}
71
72#[derive(Clone, Copy, Debug)]
73enum CipherMethod {
74    Aes128Cfb,
75    Aes192Cfb,
76    Aes256Cfb,
77    Aes128Ctr,
78    Aes192Ctr,
79    Aes256Ctr,
80    ChaCha20,
81}
82
83impl Cipher {
84    pub fn new(method: &str, password: &str) -> Cipher {
85        let (key_len, cipher_method, iv_len) = match method {
86            "aes-128-cfb" => (16, CipherMethod::Aes128Cfb, 16),
87            "aes-192-cfb" => (24, CipherMethod::Aes192Cfb, 16),
88            "aes-256-cfb" => (32, CipherMethod::Aes256Cfb, 16),
89            "aes-128-ctr" => (16, CipherMethod::Aes128Ctr, 16),
90            "aes-192-ctr" => (24, CipherMethod::Aes192Ctr, 16),
91            "aes-256-ctr" => (32, CipherMethod::Aes256Ctr, 16),
92            "chacha20" => (32, CipherMethod::ChaCha20, 12),
93            _ => panic!("method not supported"),
94        };
95
96        let key = generate_key(password.as_bytes(), key_len);
97        Cipher {
98            key: Vec::from(&key[..]),
99            key_len,
100            iv_len,
101            iv: vec![0u8; iv_len],
102            enc: None,
103            dec: None,
104            cipher_method,
105        }
106    }
107
108    pub fn init_encrypt(&mut self) {
109        if self.iv.is_empty() {
110            let rng = thread_rng();
111            self.iv = rng.sample_iter(&Standard).take(self.iv_len).collect();
112        }
113        self.enc = Some(self.new_cipher(&self.iv));
114    }
115
116    fn new_cipher(&self, iv: &[u8]) -> Box<dyn SymmetricCipher + Send + 'static> {
117        let key: &[u8] = &self.key;
118        match self.cipher_method {
119            CipherMethod::Aes128Cfb => {
120                Box::new(Aes128Cfb::new_from_slices(key, iv).expect("init cipher error"))
121            }
122            CipherMethod::Aes192Cfb => {
123                Box::new(Aes192Cfb::new_from_slices(key, iv).expect("init cipher error"))
124            }
125            CipherMethod::Aes256Cfb => {
126                Box::new(Aes256Cfb::new_from_slices(key, iv).expect("init cipher error"))
127            }
128            CipherMethod::Aes128Ctr => Box::new(Aes128Ctr::new(key.into(), iv.into())),
129            CipherMethod::Aes192Ctr => Box::new(Aes192Ctr::new(key.into(), iv.into())),
130            CipherMethod::Aes256Ctr => Box::new(Aes256Ctr::new(key.into(), iv.into())),
131            CipherMethod::ChaCha20 => Box::new(ChaCha20::new(key.into(), iv.into())),
132        }
133    }
134
135    pub fn init_decrypt(&mut self) {
136        self.dec = Some(self.new_cipher(&self.iv));
137    }
138
139    pub fn encrypt(&mut self, input: &mut [u8]) {
140        if let Some(enc) = &mut self.enc {
141            enc.encrypt(input);
142        }
143    }
144
145    pub fn decrypt(&mut self, input: &mut [u8]) {
146        if let Some(dec) = &mut self.dec {
147            dec.decrypt(input)
148        }
149    }
150
151    pub fn reset(&self) -> Cipher {
152        Cipher {
153            key: self.key.clone(),
154            iv: vec![0u8; self.iv_len],
155            iv_len: self.iv_len,
156            key_len: self.key_len,
157            enc: None,
158            dec: None,
159            cipher_method: self.cipher_method,
160        }
161    }
162}