From 22339bcef0d73e29d32024a3090c49ddf9c08561 Mon Sep 17 00:00:00 2001
From: Jack Lloyd <jack@signal.org>
Date: Thu, 15 Oct 2020 13:51:11 -0400
Subject: [PATCH] aesni: Add target_feature annotations to allow intrinsic
 inlining (#165)

---
 aes/aesni/src/aes128.rs | 33 +++++++++++++++++++++------------
 aes/aesni/src/aes192.rs | 33 +++++++++++++++++++++------------
 aes/aesni/src/aes256.rs | 33 +++++++++++++++++++++------------
 3 files changed, 63 insertions(+), 36 deletions(-)

diff --git a/aes/aesni/src/aes128.rs b/aes/aesni/src/aes128.rs
index b686801c..c227f58c 100644
--- a/aes/aesni/src/aes128.rs
+++ b/aes/aesni/src/aes128.rs
@@ -19,8 +19,9 @@ pub struct Aes128 {
 impl Aes128 {
     #[inline(always)]
     pub(crate) fn encrypt8(&self, mut blocks: [__m128i; 8]) -> [__m128i; 8] {
-        let keys = self.encrypt_keys;
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aesni128_encrypt8(keys: &[__m128i; 11], blocks: &mut [__m128i; 8]) {
             xor8!(blocks, keys[0]);
             aesenc8!(blocks, keys[1]);
             aesenc8!(blocks, keys[2]);
@@ -33,13 +34,15 @@ impl Aes128 {
             aesenc8!(blocks, keys[9]);
             aesenclast8!(blocks, keys[10]);
         }
+        unsafe { aesni128_encrypt8(&self.encrypt_keys, &mut blocks) };
         blocks
     }
 
     #[inline(always)]
-    pub(crate) fn encrypt(&self, mut block: __m128i) -> __m128i {
-        let keys = self.encrypt_keys;
-        unsafe {
+    pub(crate) fn encrypt(&self, block: __m128i) -> __m128i {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aesni128_encrypt1(keys: &[__m128i; 11], mut block: __m128i) -> __m128i {
             block = _mm_xor_si128(block, keys[0]);
             block = _mm_aesenc_si128(block, keys[1]);
             block = _mm_aesenc_si128(block, keys[2]);
@@ -52,6 +55,7 @@ impl Aes128 {
             block = _mm_aesenc_si128(block, keys[9]);
             _mm_aesenclast_si128(block, keys[10])
         }
+        unsafe { aesni128_encrypt1(&self.encrypt_keys, block) }
     }
 }
 
@@ -88,11 +92,11 @@ impl BlockCipher for Aes128 {
 
     #[inline]
     fn decrypt_block(&self, block: &mut Block128) {
-        let keys = self.decrypt_keys;
-
-        // Safety: `loadu` and `storeu` support unaligned access
-        #[allow(clippy::cast_ptr_alignment)]
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aes128_decrypt1(block: &mut Block128, keys: &[__m128i; 11]) {
+            // Safety: `loadu` and `storeu` support unaligned access
+            #[allow(clippy::cast_ptr_alignment)]
             let mut b = _mm_loadu_si128(block.as_ptr() as *const __m128i);
             b = _mm_xor_si128(b, keys[10]);
             b = _mm_aesdec_si128(b, keys[9]);
@@ -107,6 +111,8 @@ impl BlockCipher for Aes128 {
             b = _mm_aesdeclast_si128(b, keys[0]);
             _mm_storeu_si128(block.as_mut_ptr() as *mut __m128i, b);
         }
+
+        unsafe { aes128_decrypt1(block, &self.decrypt_keys) }
     }
 
     #[inline]
@@ -119,8 +125,9 @@ impl BlockCipher for Aes128 {
 
     #[inline]
     fn decrypt_blocks(&self, blocks: &mut Block128x8) {
-        let keys = self.decrypt_keys;
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aes128_decrypt8(blocks: &mut Block128x8, keys: &[__m128i; 11]) {
             let mut b = load8!(blocks);
             xor8!(b, keys[10]);
             aesdec8!(b, keys[9]);
@@ -135,6 +142,8 @@ impl BlockCipher for Aes128 {
             aesdeclast8!(b, keys[0]);
             store8!(blocks, b);
         }
+
+        unsafe { aes128_decrypt8(blocks, &self.decrypt_keys) }
     }
 }
 
diff --git a/aes/aesni/src/aes192.rs b/aes/aesni/src/aes192.rs
index cea04ae2..eb4cb0f8 100644
--- a/aes/aesni/src/aes192.rs
+++ b/aes/aesni/src/aes192.rs
@@ -19,8 +19,9 @@ pub struct Aes192 {
 impl Aes192 {
     #[inline(always)]
     pub(crate) fn encrypt8(&self, mut blocks: [__m128i; 8]) -> [__m128i; 8] {
-        let keys = self.encrypt_keys;
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aesni192_encrypt8(keys: &[__m128i; 13], blocks: &mut [__m128i; 8]) {
             xor8!(blocks, keys[0]);
             aesenc8!(blocks, keys[1]);
             aesenc8!(blocks, keys[2]);
@@ -35,13 +36,15 @@ impl Aes192 {
             aesenc8!(blocks, keys[11]);
             aesenclast8!(blocks, keys[12]);
         }
+        unsafe { aesni192_encrypt8(&self.encrypt_keys, &mut blocks) };
         blocks
     }
 
     #[inline(always)]
-    pub(crate) fn encrypt(&self, mut block: __m128i) -> __m128i {
-        let keys = self.encrypt_keys;
-        unsafe {
+    pub(crate) fn encrypt(&self, block: __m128i) -> __m128i {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aesni192_encrypt1(keys: &[__m128i; 13], mut block: __m128i) -> __m128i {
             block = _mm_xor_si128(block, keys[0]);
             block = _mm_aesenc_si128(block, keys[1]);
             block = _mm_aesenc_si128(block, keys[2]);
@@ -56,6 +59,7 @@ impl Aes192 {
             block = _mm_aesenc_si128(block, keys[11]);
             _mm_aesenclast_si128(block, keys[12])
         }
+        unsafe { aesni192_encrypt1(&self.encrypt_keys, block) }
     }
 }
 
@@ -90,11 +94,11 @@ impl BlockCipher for Aes192 {
 
     #[inline]
     fn decrypt_block(&self, block: &mut Block128) {
-        let keys = self.decrypt_keys;
-
-        // Safety: `loadu` and `storeu` support unaligned access
-        #[allow(clippy::cast_ptr_alignment)]
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aes192_decrypt1(block: &mut Block128, keys: &[__m128i; 13]) {
+            // Safety: `loadu` and `storeu` support unaligned access
+            #[allow(clippy::cast_ptr_alignment)]
             let mut b = _mm_loadu_si128(block.as_ptr() as *const __m128i);
             b = _mm_xor_si128(b, keys[12]);
             b = _mm_aesdec_si128(b, keys[11]);
@@ -111,6 +115,8 @@ impl BlockCipher for Aes192 {
             b = _mm_aesdeclast_si128(b, keys[0]);
             _mm_storeu_si128(block.as_mut_ptr() as *mut __m128i, b);
         }
+
+        unsafe { aes192_decrypt1(block, &self.decrypt_keys) }
     }
 
     #[inline]
@@ -123,8 +129,9 @@ impl BlockCipher for Aes192 {
 
     #[inline]
     fn decrypt_blocks(&self, blocks: &mut Block128x8) {
-        let keys = self.decrypt_keys;
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aes192_decrypt8(blocks: &mut Block128x8, keys: &[__m128i; 13]) {
             let mut b = load8!(blocks);
             xor8!(b, keys[12]);
             aesdec8!(b, keys[11]);
@@ -141,6 +148,8 @@ impl BlockCipher for Aes192 {
             aesdeclast8!(b, keys[0]);
             store8!(blocks, b);
         }
+
+        unsafe { aes192_decrypt8(blocks, &self.decrypt_keys) }
     }
 }
 
diff --git a/aes/aesni/src/aes256.rs b/aes/aesni/src/aes256.rs
index da32db38..193b7ca7 100644
--- a/aes/aesni/src/aes256.rs
+++ b/aes/aesni/src/aes256.rs
@@ -19,8 +19,9 @@ pub struct Aes256 {
 impl Aes256 {
     #[inline(always)]
     pub(crate) fn encrypt8(&self, mut blocks: [__m128i; 8]) -> [__m128i; 8] {
-        let keys = self.encrypt_keys;
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aesni256_encrypt8(keys: &[__m128i; 15], blocks: &mut [__m128i; 8]) {
             xor8!(blocks, keys[0]);
             aesenc8!(blocks, keys[1]);
             aesenc8!(blocks, keys[2]);
@@ -37,13 +38,15 @@ impl Aes256 {
             aesenc8!(blocks, keys[13]);
             aesenclast8!(blocks, keys[14]);
         }
+        unsafe { aesni256_encrypt8(&self.encrypt_keys, &mut blocks) };
         blocks
     }
 
     #[inline(always)]
-    pub(crate) fn encrypt(&self, mut block: __m128i) -> __m128i {
-        let keys = self.encrypt_keys;
-        unsafe {
+    pub(crate) fn encrypt(&self, block: __m128i) -> __m128i {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aesni256_encrypt1(keys: &[__m128i; 15], mut block: __m128i) -> __m128i {
             block = _mm_xor_si128(block, keys[0]);
             block = _mm_aesenc_si128(block, keys[1]);
             block = _mm_aesenc_si128(block, keys[2]);
@@ -60,6 +63,7 @@ impl Aes256 {
             block = _mm_aesenc_si128(block, keys[13]);
             _mm_aesenclast_si128(block, keys[14])
         }
+        unsafe { aesni256_encrypt1(&self.encrypt_keys, block) }
     }
 }
 
@@ -94,11 +98,11 @@ impl BlockCipher for Aes256 {
 
     #[inline]
     fn decrypt_block(&self, block: &mut Block128) {
-        let keys = self.decrypt_keys;
-
-        // Safety: `loadu` and `storeu` support unaligned access
-        #[allow(clippy::cast_ptr_alignment)]
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aes256_decrypt1(block: &mut Block128, keys: &[__m128i; 15]) {
+            // Safety: `loadu` and `storeu` support unaligned access
+            #[allow(clippy::cast_ptr_alignment)]
             let mut b = _mm_loadu_si128(block.as_ptr() as *const __m128i);
             b = _mm_xor_si128(b, keys[14]);
             b = _mm_aesdec_si128(b, keys[13]);
@@ -117,6 +121,8 @@ impl BlockCipher for Aes256 {
             b = _mm_aesdeclast_si128(b, keys[0]);
             _mm_storeu_si128(block.as_mut_ptr() as *mut __m128i, b);
         }
+
+        unsafe { aes256_decrypt1(block, &self.decrypt_keys) }
     }
 
     #[inline]
@@ -129,8 +135,9 @@ impl BlockCipher for Aes256 {
 
     #[inline]
     fn decrypt_blocks(&self, blocks: &mut Block128x8) {
-        let keys = self.decrypt_keys;
-        unsafe {
+        #[inline]
+        #[target_feature(enable = "aes")]
+        unsafe fn aes256_decrypt8(blocks: &mut Block128x8, keys: &[__m128i; 15]) {
             let mut b = load8!(blocks);
             xor8!(b, keys[14]);
             aesdec8!(b, keys[13]);
@@ -149,6 +156,8 @@ impl BlockCipher for Aes256 {
             aesdeclast8!(b, keys[0]);
             store8!(blocks, b);
         }
+
+        unsafe { aes256_decrypt8(blocks, &self.decrypt_keys) }
     }
 }