AES tests (#340)

Co-authored-by: Daniel García <[email protected]>
bitwarden · Nov 21, 2023 · bebf731 · bebf731
1 parent 27045bf
commit bebf731
Show file tree

Hide file tree

Showing 4 changed files with 110 additions and 17 deletions.
diff --git a/crates/bitwarden/src/crypto/aes_ops.rs b/crates/bitwarden/src/crypto/aes_ops.rs
@@ -13,10 +13,9 @@ use aes::cipher::{
     BlockEncryptMut, KeyIvInit,
 };
 use hmac::Mac;
-use rand::RngCore;
 
 use crate::{
-    crypto::{EncString, PbkdfSha256Hmac, PBKDF_SHA256_HMAC_OUT_SIZE},
+    crypto::{PbkdfSha256Hmac, PBKDF_SHA256_HMAC_OUT_SIZE},
     error::{CryptoError, Result},
 };
 
@@ -63,10 +62,11 @@ pub fn decrypt_aes256_hmac(
 ///
 /// A AesCbc256_B64 EncString
 #[allow(unused)]
-pub fn encrypt_aes256(data_dec: &[u8], key: GenericArray<u8, U32>) -> Result<EncString> {
-    let (iv, data) = encrypt_aes256_internal(data_dec, key);
+pub fn encrypt_aes256(data_dec: &[u8], key: GenericArray<u8, U32>) -> ([u8; 16], Vec<u8>) {
+    let rng = rand::thread_rng();
+    let (iv, data) = encrypt_aes256_internal(rng, data_dec, key);
 
-    Ok(EncString::AesCbc256_B64 { iv, data })
+    (iv, data)
 }
 
 /// Encrypt using AES-256 in CBC mode with MAC.
@@ -80,21 +80,26 @@ pub fn encrypt_aes256_hmac(
     data_dec: &[u8],
     mac_key: GenericArray<u8, U32>,
     key: GenericArray<u8, U32>,
-) -> Result<EncString> {
-    let (iv, data) = encrypt_aes256_internal(data_dec, key);
+) -> Result<([u8; 16], [u8; 32], Vec<u8>)> {
+    let rng = rand::thread_rng();
+    let (iv, data) = encrypt_aes256_internal(rng, data_dec, key);
     let mac = validate_mac(&mac_key, &iv, &data)?;
 
-    Ok(EncString::AesCbc256_HmacSha256_B64 { iv, mac, data })
+    Ok((iv, mac, data))
 }
 
 /// Encrypt using AES-256 in CBC mode.
 ///
 /// Used internally by:
 /// - [encrypt_aes256]
 /// - [encrypt_aes256_hmac]
-fn encrypt_aes256_internal(data_dec: &[u8], key: GenericArray<u8, U32>) -> ([u8; 16], Vec<u8>) {
+fn encrypt_aes256_internal(
+    mut rng: impl rand::RngCore,
+    data_dec: &[u8],
+    key: GenericArray<u8, U32>,
+) -> ([u8; 16], Vec<u8>) {
     let mut iv = [0u8; 16];
-    rand::thread_rng().fill_bytes(&mut iv);
+    rng.fill_bytes(&mut iv);
     let data = cbc::Encryptor::<aes::Aes256>::new(&key, &iv.into())
         .encrypt_padded_vec_mut::<Pkcs7>(data_dec);
 
@@ -112,3 +117,78 @@ fn validate_mac(mac_key: &[u8], iv: &[u8], data: &[u8]) -> Result<[u8; 32]> {
 
     Ok(mac)
 }
+
+#[cfg(test)]
+mod tests {
+    use aes::cipher::generic_array::sequence::GenericSequence;
+    use base64::Engine;
+    use rand::SeedableRng;
+
+    use crate::util::BASE64_ENGINE;
+
+    use super::*;
+
+    /// Helper function for generating a `GenericArray` of size 32 with each element being
+    /// a multiple of a given increment, starting from a given offset.
+    fn generate_generic_array(offset: u8, increment: u8) -> GenericArray<u8, U32> {
+        GenericArray::generate(|i| offset + i as u8 * increment)
+    }
+
+    /// Helper function for generating a vector of a given size with each element being
+    /// a multiple of a given increment, starting from a given offset.
+    fn generate_vec(length: usize, offset: u8, increment: u8) -> Vec<u8> {
+        (0..length).map(|i| offset + i as u8 * increment).collect()
+    }
+
+    #[test]
+    fn test_encrypt_aes256_internal() {
+        let key = generate_generic_array(0, 1);
+
+        let rng = rand_chacha::ChaCha8Rng::from_seed([0u8; 32]);
+        let result = encrypt_aes256_internal(rng, "EncryptMe!".as_bytes(), key);
+        assert_eq!(
+            result,
+            (
+                [62, 0, 239, 47, 137, 95, 64, 214, 127, 91, 184, 232, 31, 9, 165, 161],
+                vec![214, 76, 187, 97, 58, 146, 212, 140, 95, 164, 177, 204, 179, 133, 172, 148]
+            )
+        );
+    }
+
+    #[test]
+    fn test_validate_mac() {
+        let mac_key = generate_vec(16, 0, 16);
+
+        let iv = generate_vec(16, 0, 16);
+        let data = generate_vec(16, 0, 16);
+
+        let result = validate_mac(&mac_key, &iv, &data);
+
+        assert!(result.is_ok());
+        let mac = result.unwrap();
+        assert_eq!(mac.len(), 32);
+    }
+
+    #[test]
+    fn test_decrypt_aes256() {
+        let iv = generate_vec(16, 0, 1);
+        let iv: &[u8; 16] = iv.as_slice().try_into().unwrap();
+        let key = generate_generic_array(0, 1);
+        let data = BASE64_ENGINE.decode("ByUF8vhyX4ddU9gcooznwA==").unwrap();
+
+        let decrypted = decrypt_aes256(iv, data, key).unwrap();
+
+        assert_eq!(String::from_utf8(decrypted).unwrap(), "EncryptMe!");
+    }
+
+    #[test]
+    fn test_encrypt_decrypt_aes256() {
+        let key = generate_generic_array(0, 1);
+        let data = "EncryptMe!";
+
+        let (iv, encrypted) = encrypt_aes256(data.as_bytes(), key);
+        let decrypted = decrypt_aes256(&iv, encrypted, key).unwrap();
+
+        assert_eq!(String::from_utf8(decrypted).unwrap(), "EncryptMe!");
+    }
+}
diff --git a/crates/bitwarden/src/crypto/enc_string.rs b/crates/bitwarden/src/crypto/enc_string.rs
@@ -1,5 +1,6 @@
 use std::{fmt::Display, str::FromStr};
 
+use aes::cipher::{generic_array::GenericArray, typenum::U32};
 use base64::Engine;
 use serde::{de::Visitor, Deserialize};
 
@@ -331,6 +332,15 @@ impl serde::Serialize for EncString {
 }
 
 impl EncString {
+    pub(crate) fn encrypt_aes256_hmac(
+        data_dec: &[u8],
+        mac_key: GenericArray<u8, U32>,
+        key: GenericArray<u8, U32>,
+    ) -> Result<EncString> {
+        let (iv, mac, data) = super::encrypt_aes256_hmac(data_dec, mac_key, key)?;
+        Ok(EncString::AesCbc256_HmacSha256_B64 { iv, mac, data })
+    }
+
     /// The numerical representation of the encryption type of the [EncString].
     const fn enc_type(&self) -> u8 {
         match self {
@@ -357,7 +367,7 @@ fn invalid_len_error(expected: usize) -> impl Fn(Vec<u8>) -> EncStringParseError
 impl LocateKey for EncString {}
 impl KeyEncryptable<EncString> for &[u8] {
     fn encrypt_with_key(self, key: &SymmetricCryptoKey) -> Result<EncString> {
-        super::encrypt_aes256_hmac(self, key.mac_key.ok_or(CryptoError::InvalidMac)?, key.key)
+        EncString::encrypt_aes256_hmac(self, key.mac_key.ok_or(CryptoError::InvalidMac)?, key.key)
     }
 }
 

diff --git a/crates/bitwarden/src/crypto/master_key.rs b/crates/bitwarden/src/crypto/master_key.rs
@@ -4,8 +4,8 @@ use rand::Rng;
 use sha2::Digest;
 
 use super::{
-    encrypt_aes256_hmac, hkdf_expand, EncString, KeyDecryptable, PbkdfSha256Hmac,
-    SymmetricCryptoKey, UserKey, PBKDF_SHA256_HMAC_OUT_SIZE,
+    hkdf_expand, EncString, KeyDecryptable, PbkdfSha256Hmac, SymmetricCryptoKey, UserKey,
+    PBKDF_SHA256_HMAC_OUT_SIZE,
 };
 use crate::{client::kdf::Kdf, error::Result, util::BASE64_ENGINE};
 
@@ -61,8 +61,11 @@ fn make_user_key(
     rng.fill(&mut user_key);
 
     let stretched_key = stretch_master_key(master_key)?;
-    let protected =
-        encrypt_aes256_hmac(&user_key, stretched_key.mac_key.unwrap(), stretched_key.key)?;
+    let protected = EncString::encrypt_aes256_hmac(
+        &user_key,
+        stretched_key.mac_key.unwrap(),
+        stretched_key.key,
+    )?;
 
     let u: &[u8] = &user_key;
     Ok((UserKey::new(SymmetricCryptoKey::try_from(u)?), protected))

diff --git a/crates/bitwarden/src/crypto/rsa.rs b/crates/bitwarden/src/crypto/rsa.rs
@@ -5,7 +5,7 @@ use rsa::{
 };
 
 use crate::{
-    crypto::{encrypt_aes256_hmac, EncString, SymmetricCryptoKey},
+    crypto::{EncString, SymmetricCryptoKey},
     error::{Error, Result},
     util::BASE64_ENGINE,
 };
@@ -33,7 +33,7 @@ pub(super) fn make_key_pair(key: &SymmetricCryptoKey) -> Result<RsaKeyPair> {
         .to_pkcs8_der()
         .map_err(|_| Error::Internal("unable to create private key"))?;
 
-    let protected = encrypt_aes256_hmac(pkcs.as_bytes(), key.mac_key.unwrap(), key.key)?;
+    let protected = EncString::encrypt_aes256_hmac(pkcs.as_bytes(), key.mac_key.unwrap(), key.key)?;
 
     Ok(RsaKeyPair {
         public: b64,