refactor(crypto): implement of AesCbc*Variant and AesCtrVariant (#750)

awslabs · Dec 22, 2024 · f3f4ecc · f3f4ecc
1 parent 38ab94f
commit f3f4ecc
Show file tree

Hide file tree

Showing 3 changed files with 144 additions and 117 deletions.
diff --git a/modules/llrt_crypto/src/subtle/decrypt.rs b/modules/llrt_crypto/src/subtle/decrypt.rs
@@ -1,25 +1,16 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-use aes::cipher::{block_padding::Pkcs7, typenum::U12, BlockDecryptMut, KeyIvInit};
+use aes::cipher::typenum::U12;
 use aes_gcm::Nonce;
-use ctr::cipher::StreamCipher;
 use llrt_utils::{bytes::ObjectBytes, result::ResultExt};
 use rquickjs::{ArrayBuffer, Class, Ctx, Result};
 
-use crate::subtle::{
-    Aes128Ctr128, Aes128Ctr32, Aes128Ctr64, Aes192Ctr128, Aes192Ctr32, Aes192Ctr64, Aes256Ctr128,
-    Aes256Ctr32, Aes256Ctr64, CryptoKey,
-};
-
 use super::{
     algorithm_missmatch_error, encryption_algorithm::EncryptionAlgorithm,
-    key_algorithm::KeyAlgorithm, rsa_private_key, AesGcmVariant,
+    key_algorithm::KeyAlgorithm, rsa_private_key, AesCbcDecVariant, AesCtrVariant, AesGcmVariant,
+    CryptoKey,
 };
 
-type Aes128CbcDec = cbc::Decryptor<aes::Aes128>;
-type Aes192CbcDec = cbc::Decryptor<aes::Aes192>;
-type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;
-
 pub async fn subtle_decrypt<'js>(
     ctx: Ctx<'js>,
     algorithm: EncryptionAlgorithm,
@@ -42,31 +33,18 @@ fn decrypt(
     match algorithm {
         EncryptionAlgorithm::AesCbc { iv } => {
             if let KeyAlgorithm::Aes { length } = key.algorithm {
-                match length {
-                    128 => decrypt_aes_cbc_gen::<Aes128CbcDec>(ctx, handle, iv, data),
-                    192 => decrypt_aes_cbc_gen::<Aes192CbcDec>(ctx, handle, iv, data),
-                    256 => decrypt_aes_cbc_gen::<Aes256CbcDec>(ctx, handle, iv, data),
-                    _ => unreachable!(), // 'length' has already been sanitized.
-                }
+                let variant = AesCbcDecVariant::new(length, handle, iv).or_throw(ctx)?;
+                variant.decrypt(data).or_throw(ctx)
             } else {
                 algorithm_missmatch_error(ctx)
             }
         },
         EncryptionAlgorithm::AesCtr { counter, length } => {
             let encryption_length = length;
             if let KeyAlgorithm::Aes { length } = key.algorithm {
-                match (length, encryption_length) {
-                    (128, 32) => decrypt_aes_ctr_gen::<Aes128Ctr32>(ctx, handle, counter, data),
-                    (128, 64) => decrypt_aes_ctr_gen::<Aes128Ctr64>(ctx, handle, counter, data),
-                    (128, 128) => decrypt_aes_ctr_gen::<Aes128Ctr128>(ctx, handle, counter, data),
-                    (192, 32) => decrypt_aes_ctr_gen::<Aes192Ctr32>(ctx, handle, counter, data),
-                    (192, 64) => decrypt_aes_ctr_gen::<Aes192Ctr64>(ctx, handle, counter, data),
-                    (192, 128) => decrypt_aes_ctr_gen::<Aes192Ctr128>(ctx, handle, counter, data),
-                    (256, 32) => decrypt_aes_ctr_gen::<Aes256Ctr32>(ctx, handle, counter, data),
-                    (256, 64) => decrypt_aes_ctr_gen::<Aes256Ctr64>(ctx, handle, counter, data),
-                    (256, 128) => decrypt_aes_ctr_gen::<Aes256Ctr128>(ctx, handle, counter, data),
-                    _ => unreachable!(), // 'length' has already been sanitized.
-                }
+                let mut variant = AesCtrVariant::new(length, *encryption_length, handle, counter)
+                    .or_throw(ctx)?;
+                variant.decrypt(data).or_throw(ctx)
             } else {
                 algorithm_missmatch_error(ctx)
             }
@@ -98,24 +76,3 @@ fn decrypt(
         },
     }
 }
-
-fn decrypt_aes_cbc_gen<T>(ctx: &Ctx<'_>, key: &[u8], iv: &[u8], data: &[u8]) -> Result<Vec<u8>>
-where
-    T: KeyIvInit + BlockDecryptMut,
-{
-    T::new(key.into(), iv.into())
-        .decrypt_padded_vec_mut::<Pkcs7>(data)
-        .or_throw(ctx)
-}
-
-fn decrypt_aes_ctr_gen<T>(ctx: &Ctx<'_>, key: &[u8], counter: &[u8], data: &[u8]) -> Result<Vec<u8>>
-where
-    T: KeyIvInit + StreamCipher,
-{
-    let mut cipher = T::new(key.into(), counter.into());
-
-    let mut plaintext = data.to_vec();
-    cipher.try_apply_keystream(&mut plaintext).or_throw(ctx)?;
-
-    Ok(plaintext)
-}
diff --git a/modules/llrt_crypto/src/subtle/encrypt.rs b/modules/llrt_crypto/src/subtle/encrypt.rs
@@ -1,26 +1,17 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-use aes::cipher::{block_padding::Pkcs7, typenum::U12, KeyIvInit};
+use aes::cipher::typenum::U12;
 use aes_gcm::Nonce;
-use ctr::cipher::{BlockEncryptMut, StreamCipher};
 use llrt_utils::{bytes::ObjectBytes, result::ResultExt};
 use rquickjs::{ArrayBuffer, Class, Ctx, Result};
 use rsa::rand_core::OsRng;
 
-use crate::subtle::{
-    Aes128Ctr128, Aes128Ctr32, Aes128Ctr64, Aes192Ctr128, Aes192Ctr32, Aes192Ctr64, Aes256Ctr128,
-    Aes256Ctr32, Aes256Ctr64, CryptoKey,
-};
-
 use super::{
     algorithm_missmatch_error, encryption_algorithm::EncryptionAlgorithm,
-    key_algorithm::KeyAlgorithm, rsa_private_key, AesGcmVariant,
+    key_algorithm::KeyAlgorithm, rsa_private_key, AesCbcEncVariant, AesCtrVariant, AesGcmVariant,
+    CryptoKey,
 };
 
-type Aes128CbcEnc = cbc::Encryptor<aes::Aes128>;
-type Aes192CbcEnc = cbc::Encryptor<aes::Aes192>;
-type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
-
 pub async fn subtle_encrypt<'js>(
     ctx: Ctx<'js>,
     algorithm: EncryptionAlgorithm,
@@ -44,31 +35,18 @@ fn encrypt(
     match algorithm {
         EncryptionAlgorithm::AesCbc { iv } => {
             if let KeyAlgorithm::Aes { length } = key.algorithm {
-                match length {
-                    128 => encrypt_aes_cbc_gen::<Aes128CbcEnc>(ctx, handle, iv, data),
-                    192 => encrypt_aes_cbc_gen::<Aes192CbcEnc>(ctx, handle, iv, data),
-                    256 => encrypt_aes_cbc_gen::<Aes256CbcEnc>(ctx, handle, iv, data),
-                    _ => unreachable!(), // 'length' has already been sanitized.
-                }
+                let variant = AesCbcEncVariant::new(length, handle, iv).or_throw(ctx)?;
+                Ok(variant.encrypt(data))
             } else {
                 algorithm_missmatch_error(ctx)
             }
         },
         EncryptionAlgorithm::AesCtr { counter, length } => {
             let encryption_length = length;
             if let KeyAlgorithm::Aes { length } = key.algorithm {
-                match (length, encryption_length) {
-                    (128, 32) => encrypt_aes_ctr_gen::<Aes128Ctr32>(ctx, handle, counter, data),
-                    (128, 64) => encrypt_aes_ctr_gen::<Aes128Ctr64>(ctx, handle, counter, data),
-                    (128, 128) => encrypt_aes_ctr_gen::<Aes128Ctr128>(ctx, handle, counter, data),
-                    (192, 32) => encrypt_aes_ctr_gen::<Aes192Ctr32>(ctx, handle, counter, data),
-                    (192, 64) => encrypt_aes_ctr_gen::<Aes192Ctr64>(ctx, handle, counter, data),
-                    (192, 128) => encrypt_aes_ctr_gen::<Aes192Ctr128>(ctx, handle, counter, data),
-                    (256, 32) => encrypt_aes_ctr_gen::<Aes256Ctr32>(ctx, handle, counter, data),
-                    (256, 64) => encrypt_aes_ctr_gen::<Aes256Ctr64>(ctx, handle, counter, data),
-                    (256, 128) => encrypt_aes_ctr_gen::<Aes256Ctr128>(ctx, handle, counter, data),
-                    _ => unreachable!(), // 'length' has already been sanitized.
-                }
+                let mut variant = AesCtrVariant::new(length, *encryption_length, handle, counter)
+                    .or_throw(ctx)?;
+                variant.encrypt(data).or_throw(ctx)
             } else {
                 algorithm_missmatch_error(ctx)
             }
@@ -102,22 +80,3 @@ fn encrypt(
         },
     }
 }
-
-fn encrypt_aes_cbc_gen<T>(_ctx: &Ctx<'_>, key: &[u8], iv: &[u8], data: &[u8]) -> Result<Vec<u8>>
-where
-    T: KeyIvInit + BlockEncryptMut,
-{
-    Ok(T::new(key.into(), iv.into()).encrypt_padded_vec_mut::<Pkcs7>(data))
-}
-
-fn encrypt_aes_ctr_gen<T>(ctx: &Ctx<'_>, key: &[u8], counter: &[u8], data: &[u8]) -> Result<Vec<u8>>
-where
-    T: KeyIvInit + StreamCipher,
-{
-    let mut cipher = T::new(key.into(), counter.into());
-
-    let mut ciphertext = data.to_vec();
-    cipher.try_apply_keystream(&mut ciphertext).or_throw(ctx)?;
-
-    Ok(ciphertext)
-}
diff --git a/modules/llrt_crypto/src/subtle/mod.rs b/modules/llrt_crypto/src/subtle/mod.rs
@@ -23,19 +23,16 @@ pub use digest::subtle_digest;
 pub use encrypt::subtle_encrypt;
 pub use export_key::subtle_export_key;
 pub use generate_key::subtle_generate_key;
-use llrt_utils::object::ObjectExt;
-use ring::signature;
-use rquickjs::Object;
-use rquickjs::Value;
-use rsa::{pkcs1::DecodeRsaPrivateKey, Oaep, RsaPrivateKey};
 pub use sign::subtle_sign;
 pub use verify::subtle_verify;
 
 use aes::{
     cipher::{
+        block_padding::{Pkcs7, UnpadError},
         consts::{U13, U14, U15, U16},
         typenum::U12,
-        InvalidLength,
+        BlockDecryptMut, BlockEncryptMut, InvalidLength, KeyIvInit, StreamCipher,
+        StreamCipherError,
     },
     Aes128, Aes192, Aes256,
 };
@@ -44,20 +41,134 @@ use aes_gcm::{
     AesGcm, KeyInit,
 };
 use ctr::{Ctr128BE, Ctr32BE, Ctr64BE};
-use llrt_utils::{result::ResultExt, str_enum};
-use rquickjs::{Ctx, Exception, Result};
+use llrt_utils::{object::ObjectExt, result::ResultExt, str_enum};
+use ring::signature;
+use rquickjs::{Ctx, Exception, Object, Result, Value};
+use rsa::{pkcs1::DecodeRsaPrivateKey, Oaep, RsaPrivateKey};
 
 use crate::sha_hash::ShaAlgorithm;
 
-type Aes128Ctr32 = Ctr32BE<aes::Aes128>;
-type Aes128Ctr64 = Ctr64BE<aes::Aes128>;
-type Aes128Ctr128 = Ctr128BE<aes::Aes128>;
-type Aes192Ctr32 = Ctr32BE<aes::Aes192>;
-type Aes192Ctr64 = Ctr64BE<aes::Aes192>;
-type Aes192Ctr128 = Ctr128BE<aes::Aes192>;
-type Aes256Ctr32 = Ctr32BE<aes::Aes256>;
-type Aes256Ctr64 = Ctr64BE<aes::Aes256>;
-type Aes256Ctr128 = Ctr128BE<aes::Aes256>;
+pub enum AesCbcEncVariant {
+    Aes128(cbc::Encryptor<aes::Aes128>),
+    Aes192(cbc::Encryptor<aes::Aes192>),
+    Aes256(cbc::Encryptor<aes::Aes256>),
+}
+
+impl AesCbcEncVariant {
+    pub fn new(key_len: u16, key: &[u8], iv: &[u8]) -> std::result::Result<Self, InvalidLength> {
+        let variant: AesCbcEncVariant = match key_len {
+            128 => Self::Aes128(cbc::Encryptor::new_from_slices(key, iv)?),
+            192 => Self::Aes192(cbc::Encryptor::new_from_slices(key, iv)?),
+            256 => Self::Aes256(cbc::Encryptor::new_from_slices(key, iv)?),
+            _ => return Err(InvalidLength),
+        };
+
+        Ok(variant)
+    }
+
+    pub fn encrypt(&self, data: &[u8]) -> Vec<u8> {
+        match self {
+            Self::Aes128(v) => v.clone().encrypt_padded_vec_mut::<Pkcs7>(data),
+            Self::Aes192(v) => v.clone().encrypt_padded_vec_mut::<Pkcs7>(data),
+            Self::Aes256(v) => v.clone().encrypt_padded_vec_mut::<Pkcs7>(data),
+        }
+    }
+}
+
+pub enum AesCbcDecVariant {
+    Aes128(cbc::Decryptor<aes::Aes128>),
+    Aes192(cbc::Decryptor<aes::Aes192>),
+    Aes256(cbc::Decryptor<aes::Aes256>),
+}
+
+impl AesCbcDecVariant {
+    pub fn new(key_len: u16, key: &[u8], iv: &[u8]) -> std::result::Result<Self, InvalidLength> {
+        let variant: AesCbcDecVariant = match key_len {
+            128 => Self::Aes128(cbc::Decryptor::new_from_slices(key, iv)?),
+            192 => Self::Aes192(cbc::Decryptor::new_from_slices(key, iv)?),
+            256 => Self::Aes256(cbc::Decryptor::new_from_slices(key, iv)?),
+            _ => return Err(InvalidLength),
+        };
+
+        Ok(variant)
+    }
+
+    pub fn decrypt(&self, data: &[u8]) -> std::result::Result<Vec<u8>, UnpadError> {
+        Ok(match self {
+            Self::Aes128(v) => v.clone().decrypt_padded_vec_mut::<Pkcs7>(data)?,
+            Self::Aes192(v) => v.clone().decrypt_padded_vec_mut::<Pkcs7>(data)?,
+            Self::Aes256(v) => v.clone().decrypt_padded_vec_mut::<Pkcs7>(data)?,
+        })
+    }
+}
+
+pub enum AesCtrVariant {
+    Aes128Ctr32(Ctr32BE<aes::Aes128>),
+    Aes128Ctr64(Ctr64BE<aes::Aes128>),
+    Aes128Ctr128(Ctr128BE<aes::Aes128>),
+    Aes192Ctr32(Ctr32BE<aes::Aes192>),
+    Aes192Ctr64(Ctr64BE<aes::Aes192>),
+    Aes192Ctr128(Ctr128BE<aes::Aes192>),
+    Aes256Ctr32(Ctr32BE<aes::Aes256>),
+    Aes256Ctr64(Ctr64BE<aes::Aes256>),
+    Aes256Ctr128(Ctr128BE<aes::Aes256>),
+}
+
+impl AesCtrVariant {
+    pub fn new(
+        key_len: u16,
+        encryption_length: u32,
+        key: &[u8],
+        counter: &[u8],
+    ) -> std::result::Result<Self, InvalidLength> {
+        let variant: AesCtrVariant = match (key_len, encryption_length) {
+            (128, 32) => Self::Aes128Ctr32(Ctr32BE::new_from_slices(key, counter)?),
+            (128, 64) => Self::Aes128Ctr64(Ctr64BE::new_from_slices(key, counter)?),
+            (128, 128) => Self::Aes128Ctr128(Ctr128BE::new_from_slices(key, counter)?),
+            (192, 32) => Self::Aes192Ctr32(Ctr32BE::new_from_slices(key, counter)?),
+            (192, 64) => Self::Aes192Ctr64(Ctr64BE::new_from_slices(key, counter)?),
+            (192, 128) => Self::Aes192Ctr128(Ctr128BE::new_from_slices(key, counter)?),
+            (256, 32) => Self::Aes256Ctr32(Ctr32BE::new_from_slices(key, counter)?),
+            (256, 64) => Self::Aes256Ctr64(Ctr64BE::new_from_slices(key, counter)?),
+            (256, 128) => Self::Aes256Ctr128(Ctr128BE::new_from_slices(key, counter)?),
+            _ => return Err(InvalidLength),
+        };
+
+        Ok(variant)
+    }
+
+    pub fn encrypt(&mut self, data: &[u8]) -> std::result::Result<Vec<u8>, StreamCipherError> {
+        let mut ciphertext = data.to_vec();
+        match self {
+            Self::Aes128Ctr32(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes128Ctr64(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes128Ctr128(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes192Ctr32(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes192Ctr64(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes192Ctr128(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes256Ctr32(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes256Ctr64(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes256Ctr128(v) => v.try_apply_keystream(&mut ciphertext)?,
+        }
+        Ok(ciphertext)
+    }
+
+    pub fn decrypt(&mut self, data: &[u8]) -> std::result::Result<Vec<u8>, StreamCipherError> {
+        let mut ciphertext = data.to_vec();
+        match self {
+            Self::Aes128Ctr32(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes128Ctr64(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes128Ctr128(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes192Ctr32(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes192Ctr64(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes192Ctr128(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes256Ctr32(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes256Ctr64(v) => v.try_apply_keystream(&mut ciphertext)?,
+            Self::Aes256Ctr128(v) => v.try_apply_keystream(&mut ciphertext)?,
+        }
+        Ok(ciphertext)
+    }
+}
 
 pub enum AesGcmVariant {
     Aes128Gcm96(AesGcm<Aes128, U12, U12>),