Move EVP_AEAD_CTX to a heap allocated structure (#210)

aws · Aug 9, 2023 · c74566e · c74566e
1 parent 15daf8e
commit c74566e
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 36 deletions.
diff --git a/aws-lc-rs/src/aead.rs b/aws-lc-rs/src/aead.rs
@@ -764,7 +764,7 @@ where
         let add_str = aad.0;
 
         if 1 != EVP_AEAD_CTX_seal(
-            aead_ctx,
+            *aead_ctx.as_const(),
             mut_in_out.as_mut_ptr(),
             out_len.as_mut_ptr(),
             plaintext_len + TAG_LEN,
@@ -803,7 +803,7 @@ pub(crate) fn aead_open_combined(
         let aad_str = aad.0;
         let mut out_len = MaybeUninit::<usize>::uninit();
         if 1 != EVP_AEAD_CTX_open(
-            aead_ctx,
+            *aead_ctx.as_const(),
             in_out.as_mut_ptr(),
             out_len.as_mut_ptr(),
             plaintext_len,

diff --git a/aws-lc-rs/src/aead/aead_ctx.rs b/aws-lc-rs/src/aead/aead_ctx.rs
@@ -6,22 +6,21 @@ use crate::cipher::chacha;
 
 use crate::cipher::aes::{AES_128_KEY_LEN, AES_256_KEY_LEN};
 use crate::error::Unspecified;
+use crate::ptr::LcPtr;
 use aws_lc::{
-    EVP_AEAD_CTX_cleanup, EVP_AEAD_CTX_init, EVP_AEAD_CTX_zero, EVP_aead_aes_128_gcm,
-    EVP_aead_aes_256_gcm, EVP_aead_chacha20_poly1305, EVP_AEAD_CTX,
+    EVP_AEAD_CTX_new, EVP_aead_aes_128_gcm, EVP_aead_aes_256_gcm, EVP_aead_chacha20_poly1305,
+    EVP_AEAD_CTX,
 };
-use std::mem::MaybeUninit;
-use std::ptr::null_mut;
 
 #[allow(
     clippy::large_enum_variant,
     variant_size_differences,
     non_camel_case_types
 )]
 pub(crate) enum AeadCtx {
-    AES_128_GCM(EVP_AEAD_CTX),
-    AES_256_GCM(EVP_AEAD_CTX),
-    CHACHA20_POLY1305(EVP_AEAD_CTX),
+    AES_128_GCM(LcPtr<*mut EVP_AEAD_CTX>),
+    AES_256_GCM(LcPtr<*mut EVP_AEAD_CTX>),
+    CHACHA20_POLY1305(LcPtr<*mut EVP_AEAD_CTX>),
 }
 
 unsafe impl Send for AeadCtx {}
@@ -61,36 +60,18 @@ impl AeadCtx {
     fn build_context(
         aead_fn: unsafe extern "C" fn() -> *const aws_lc::evp_aead_st,
         key_bytes: &[u8],
-    ) -> Result<EVP_AEAD_CTX, Unspecified> {
-        let mut aead_ctx = MaybeUninit::<EVP_AEAD_CTX>::uninit();
-        unsafe {
-            let aead = aead_fn();
+    ) -> Result<LcPtr<*mut EVP_AEAD_CTX>, Unspecified> {
+        let aead = unsafe { aead_fn() };
 
-            if 1 != EVP_AEAD_CTX_init(
-                aead_ctx.as_mut_ptr(),
+        let aead_ctx = unsafe {
+            LcPtr::new(EVP_AEAD_CTX_new(
                 aead,
                 key_bytes.as_ptr().cast(),
                 key_bytes.len(),
                 TAG_LEN,
-                null_mut(),
-            ) {
-                return Err(Unspecified);
-            }
-            Ok(aead_ctx.assume_init())
-        }
-    }
-}
+            ))?
+        };
 
-impl Drop for AeadCtx {
-    fn drop(&mut self) {
-        unsafe {
-            let ctx = match self {
-                AeadCtx::AES_128_GCM(ctx)
-                | AeadCtx::AES_256_GCM(ctx)
-                | AeadCtx::CHACHA20_POLY1305(ctx) => ctx,
-            };
-            EVP_AEAD_CTX_cleanup(ctx);
-            EVP_AEAD_CTX_zero(ctx);
-        }
+        Ok(aead_ctx)
     }
 }
diff --git a/aws-lc-rs/src/aead/aes_gcm.rs b/aws-lc-rs/src/aead/aes_gcm.rs
@@ -31,7 +31,7 @@ pub(crate) fn aead_seal_separate(
         let mut out_tag_len = MaybeUninit::<usize>::uninit();
 
         if 1 != EVP_AEAD_CTX_seal_scatter(
-            aead_ctx,
+            *aead_ctx.as_const(),
             in_out.as_mut_ptr(),
             tag.as_mut_ptr().cast(),
             out_tag_len.as_mut_ptr(),

diff --git a/aws-lc-rs/src/ptr.rs b/aws-lc-rs/src/ptr.rs
@@ -3,7 +3,7 @@
 
 use std::ops::Deref;
 
-use aws_lc::OPENSSL_free;
+use aws_lc::{EVP_AEAD_CTX_free, OPENSSL_free, EVP_AEAD_CTX};
 
 use mirai_annotations::verify_unreachable;
 
@@ -197,6 +197,7 @@ create_pointer!(ECDSA_SIG, ECDSA_SIG_free);
 create_pointer!(BIGNUM, BN_free);
 create_pointer!(EVP_PKEY, EVP_PKEY_free);
 create_pointer!(RSA, RSA_free);
+create_pointer!(EVP_AEAD_CTX, EVP_AEAD_CTX_free);
 
 #[cfg(test)]
 mod tests {