1use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
9use generic_array::GenericArray;
10use hmac::Mac;
11use subtle::ConstantTimeEq;
12use typenum::U32;
13
14use crate::{
15 error::{CryptoError, Result},
16 util::{PbkdfSha256Hmac, PBKDF_SHA256_HMAC_OUT_SIZE},
17};
18
19pub(crate) fn decrypt_aes256(
23 iv: &[u8; 16],
24 data: Vec<u8>,
25 key: &GenericArray<u8, U32>,
26) -> Result<Vec<u8>> {
27 let iv = GenericArray::from_slice(iv);
29 let mut data = data;
30 let decrypted_key_slice = cbc::Decryptor::<aes::Aes256>::new(key, iv)
31 .decrypt_padded_mut::<Pkcs7>(&mut data)
32 .map_err(|_| CryptoError::KeyDecrypt)?;
33
34 let decrypted_len = decrypted_key_slice.len();
37 data.truncate(decrypted_len);
38
39 Ok(data)
40}
41
42pub(crate) fn decrypt_aes256_hmac(
46 iv: &[u8; 16],
47 mac: &[u8; 32],
48 data: Vec<u8>,
49 mac_key: &GenericArray<u8, U32>,
50 key: &GenericArray<u8, U32>,
51) -> Result<Vec<u8>> {
52 let res = generate_mac(mac_key, iv, &data)?;
53 if res.ct_ne(mac).into() {
54 return Err(CryptoError::InvalidMac);
55 }
56 decrypt_aes256(iv, data, key)
57}
58
59pub(crate) fn encrypt_aes256_hmac(
65 data_dec: &[u8],
66 mac_key: &GenericArray<u8, U32>,
67 key: &GenericArray<u8, U32>,
68) -> Result<([u8; 16], [u8; 32], Vec<u8>)> {
69 let rng = rand::thread_rng();
70 let (iv, data) = encrypt_aes256_internal(rng, data_dec, key);
71 let mac = generate_mac(mac_key, &iv, &data)?;
72
73 Ok((iv, mac, data))
74}
75
76fn encrypt_aes256_internal(
81 mut rng: impl rand::RngCore,
82 data_dec: &[u8],
83 key: &GenericArray<u8, U32>,
84) -> ([u8; 16], Vec<u8>) {
85 let mut iv = [0u8; 16];
86 rng.fill_bytes(&mut iv);
87 let data = cbc::Encryptor::<aes::Aes256>::new(key, &iv.into())
88 .encrypt_padded_vec_mut::<Pkcs7>(data_dec);
89
90 (iv, data)
91}
92
93fn generate_mac(mac_key: &[u8], iv: &[u8], data: &[u8]) -> Result<[u8; 32]> {
95 let mut hmac =
96 PbkdfSha256Hmac::new_from_slice(mac_key).expect("hmac new_from_slice should not fail");
97 hmac.update(iv);
98 hmac.update(data);
99 let mac: [u8; PBKDF_SHA256_HMAC_OUT_SIZE] = (*hmac.finalize().into_bytes())
100 .try_into()
101 .map_err(|_| CryptoError::InvalidMac)?;
102
103 Ok(mac)
104}
105
106#[cfg(test)]
107mod tests {
108 use base64::{engine::general_purpose::STANDARD, Engine};
109 use generic_array::{sequence::GenericSequence, ArrayLength};
110 use rand::SeedableRng;
111
112 use super::*;
113
114 fn generate_generic_array<N: ArrayLength<u8>>(
117 offset: u8,
118 increment: u8,
119 ) -> GenericArray<u8, N> {
120 GenericArray::generate(|i| offset + i as u8 * increment)
121 }
122
123 fn generate_vec(length: usize, offset: u8, increment: u8) -> Vec<u8> {
126 (0..length).map(|i| offset + i as u8 * increment).collect()
127 }
128
129 #[test]
130 fn test_encrypt_aes256_internal() {
131 let key = generate_generic_array(0, 1);
132
133 let rng = rand_chacha::ChaCha8Rng::from_seed([0u8; 32]);
134 let result = encrypt_aes256_internal(rng, "EncryptMe!".as_bytes(), &key);
135 assert_eq!(
136 result,
137 (
138 [62, 0, 239, 47, 137, 95, 64, 214, 127, 91, 184, 232, 31, 9, 165, 161],
139 vec![214, 76, 187, 97, 58, 146, 212, 140, 95, 164, 177, 204, 179, 133, 172, 148]
140 )
141 );
142 }
143
144 #[test]
145 fn test_generate_mac() {
146 let mac_key = generate_vec(16, 0, 16);
147
148 let iv = generate_vec(16, 0, 16);
149 let data = generate_vec(16, 0, 16);
150
151 let result = generate_mac(&mac_key, &iv, &data);
152
153 assert!(result.is_ok());
154 let mac = result.unwrap();
155 assert_eq!(mac.len(), 32);
156 }
157
158 #[test]
159 fn test_decrypt_aes256() {
160 let iv = generate_vec(16, 0, 1);
161 let iv: &[u8; 16] = iv.as_slice().try_into().unwrap();
162 let key = generate_generic_array(0, 1);
163 let data = STANDARD.decode("ByUF8vhyX4ddU9gcooznwA==").unwrap();
164
165 let decrypted = decrypt_aes256(iv, data, &key).unwrap();
166
167 assert_eq!(String::from_utf8(decrypted).unwrap(), "EncryptMe!");
168 }
169}