aes-gcm: Clarify CPU feature detection.

Although every key has been represented with the same types
`aes::AES_KEY` and `gcm::HTable` regardless of which implementation is
used, in reality those types are polymorphic in ways that aren't
captured by the type system currently. Thus, the
`set_encrypt_key!` function must be matched with the corresponding
`encrypt_block!` and/or `ctr32_encrypt_blocks!` function. Previously,
we did CPU feature detection for each function call and assumed that
CPU feature detection is idempotent. Now, we do CPU feature detection
during key construction and make the lesser assumption that at least
those same CPU features are available as long as the key exists.

This is a step towards making further improvements in CPU-feature-based
dispatching.

One side-effect of this change is that GCM keys (and thus AES-GCM keys)
are now much smaller on targets that don't support any assembly
implementation, as they now just store a single `U128` instead of a
whole `HTable`.

```
git difftool HEAD^1:src/aead/aes.rs src/aead/aes/bs.rs
git difftool HEAD^1:src/aead/aes.rs src/aead/aes/vp.rs
```

src/aead/aes/bs.rs

@@ -0,0 +1,62 @@
+// Copyright 2018-2024 Brian Smith.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES
+// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
+// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+#![cfg(target_arch = "arm")]
+
+use super::{Counter, AES_KEY};
+use core::ops::RangeFrom;
+
+/// SAFETY:
+///   * The caller must ensure that if blocks > 0 then either `input` and
+///     `output` do not overlap at all, or input == output.add(n) for some
+///     (nonnegative) n.
+///   * if blocks > 0, The caller must ensure `input` points to `blocks` blocks
+///     and that `output` points to writable space for `blocks` blocks.
+///   * The caller must ensure that `vpaes_key` was initialized with
+///     `vpaes_set_encrypt_key`.
+///   * Upon returning, `blocks` blocks will have been read from `input` and
+///     written to `output`.
+pub(super) unsafe fn ctr32_encrypt_blocks_with_vpaes_key(
+    in_out: &mut [u8],
+    src: RangeFrom<usize>,
+    vpaes_key: &AES_KEY,
+    ctr: &mut Counter,
+) {
+    prefixed_extern! {
+        // bsaes_ctr32_encrypt_blocks requires transformation of an existing
+        // VPAES key; there is no `bsaes_set_encrypt_key`.
+        fn vpaes_encrypt_key_to_bsaes(bsaes_key: *mut AES_KEY, vpaes_key: &AES_KEY);
+    }
+
+    // SAFETY:
+    //   * The caller ensures `vpaes_key` was initialized by
+    //     `vpaes_set_encrypt_key`.
+    //   * `bsaes_key was zeroed above, and `vpaes_encrypt_key_to_bsaes`
+    //     is assumed to initialize `bsaes_key`.
+    let bsaes_key = unsafe { AES_KEY::derive(vpaes_encrypt_key_to_bsaes, vpaes_key) };
+
+    // The code for `vpaes_encrypt_key_to_bsaes` notes "vpaes stores one
+    // fewer round count than bsaes, but the number of keys is the same,"
+    // so use this as a sanity check.
+    debug_assert_eq!(bsaes_key.rounds(), vpaes_key.rounds() + 1);
+
+    // SAFETY:
+    //  * `bsaes_key` is in bsaes format after calling
+    //    `vpaes_encrypt_key_to_bsaes`.
+    //  * `bsaes_ctr32_encrypt_blocks` satisfies the contract for
+    //    `ctr32_encrypt_blocks`.
+    unsafe {
+        ctr32_encrypt_blocks!(bsaes_ctr32_encrypt_blocks, in_out, src, &bsaes_key, ctr,);
+    }
+}

src/aead/aes/fallback.rs

@@ -0,0 +1,47 @@
+// Copyright 2018-2024 Brian Smith.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES
+// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
+// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+use super::{Block, Counter, EncryptBlock, EncryptCtr32, Iv, KeyBytes, AES_KEY};
+use crate::error;
+use core::ops::RangeFrom;
+
+#[derive(Clone)]
+pub struct Key {
+    inner: AES_KEY,
+}
+
+impl Key {
+    pub(in super::super) fn new(bytes: KeyBytes<'_>) -> Result<Self, error::Unspecified> {
+        let inner = unsafe { set_encrypt_key!(aes_nohw_set_encrypt_key, bytes) }?;
+        Ok(Self { inner })
+    }
+}
+
+impl EncryptBlock for Key {
+    fn encrypt_block(&self, block: Block) -> Block {
+        unsafe { encrypt_block!(aes_nohw_encrypt, block, &self.inner) }
+    }
+
+    fn encrypt_iv_xor_block(&self, iv: Iv, block: Block) -> Block {
+        super::encrypt_iv_xor_block_using_encrypt_block(self, iv, block)
+    }
+}
+
+impl EncryptCtr32 for Key {
+    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+        unsafe {
+            ctr32_encrypt_blocks!(aes_nohw_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr)
+        }
+    }
+}

src/aead/aes/ffi.rs

@@ -13,7 +13,7 @@
 // CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 use super::{Block, KeyBytes, BLOCK_LEN};
-use crate::{bits::BitLength, c, cpu, error, polyfill::slice};
+use crate::{bits::BitLength, c, error, polyfill::slice};
 use core::{num::NonZeroUsize, ops::RangeFrom};
 
 /// nonce || big-endian counter.
@@ -36,7 +36,6 @@ impl AES_KEY {
     pub(super) unsafe fn new(
         f: unsafe extern "C" fn(*const u8, BitLength<c::int>, *mut AES_KEY) -> c::int,
         bytes: KeyBytes<'_>,
-        _cpu_features: cpu::Features,
     ) -> Result<Self, error::Unspecified> {
         let mut key = Self {
             rd_key: [0; 4 * (MAX_ROUNDS + 1)],
@@ -63,7 +62,6 @@ impl AES_KEY {
     pub(super) unsafe fn derive(
         f: for<'a> unsafe extern "C" fn(*mut AES_KEY, &'a AES_KEY),
         src: &Self,
-        _cpu_features: cpu::Features,
     ) -> Self {
         let mut r = AES_KEY {
             rd_key: [0u32; 4 * (MAX_ROUNDS + 1)],
@@ -89,12 +87,12 @@ impl AES_KEY {
 // In BoringSSL, the C prototypes for these are in
 // crypto/fipsmodule/aes/internal.h.
 macro_rules! set_encrypt_key {
-    ( $name:ident, $key_bytes:expr, $cpu_features:expr $(,)? ) => {{
+    ( $name:ident, $key_bytes:expr $(,)? ) => {{
         use crate::{bits::BitLength, c};
         prefixed_extern! {
             fn $name(user_key: *const u8, bits: BitLength<c::int>, key: *mut AES_KEY) -> c::int;
         }
-        $crate::aead::aes::ffi::AES_KEY::new($name, $key_bytes, $cpu_features)
+        $crate::aead::aes::ffi::AES_KEY::new($name, $key_bytes)
     }};
 }
 
@@ -129,7 +127,7 @@ impl AES_KEY {
 ///   * The caller must ensure that fhe function `$name` satisfies the conditions
 ///     for the `f` parameter to `ctr32_encrypt_blocks`.
 macro_rules! ctr32_encrypt_blocks {
-    ($name:ident, $in_out:expr, $src:expr, $key:expr, $ctr:expr, $cpu_features:expr ) => {{
+    ($name:ident, $in_out:expr, $src:expr, $key:expr, $ctr:expr $(,)? ) => {{
         use crate::{
             aead::aes::{ffi::AES_KEY, Counter, BLOCK_LEN},
             c,
@@ -143,7 +141,7 @@ macro_rules! ctr32_encrypt_blocks {
                 ivec: &Counter,
             );
         }
-        $key.ctr32_encrypt_blocks($name, $in_out, $src, $ctr, $cpu_features)
+        $key.ctr32_encrypt_blocks($name, $in_out, $src, $ctr)
     }};
 }
 
@@ -172,7 +170,6 @@ impl AES_KEY {
         in_out: &mut [u8],
         src: RangeFrom<usize>,
         ctr: &mut Counter,
-        cpu_features: cpu::Features,
     ) {
         let (input, leftover) = slice::as_chunks(&in_out[src]);
         debug_assert_eq!(leftover.len(), 0);
@@ -189,8 +186,6 @@ impl AES_KEY {
         let input = input.as_ptr();
         let output: *mut [u8; BLOCK_LEN] = in_out.as_mut_ptr().cast();
 
-        let _: cpu::Features = cpu_features;
-
         // SAFETY:
         //  * `input` points to `blocks` blocks.
         //  * `output` points to space for `blocks` blocks to be written.
@@ -200,8 +195,6 @@ impl AES_KEY {
         //    `blocks` including zero.
         //  * The caller is responsible for ensuring `key` was initialized by the
         //    `set_encrypt_key!` invocation required by `f`.
-        //  * CPU feature detection has been done so `f` can inspect
-        //    CPU features.
         unsafe {
             f(input, output, blocks, self, ctr);
         }

src/aead/aes/hw.rs

@@ -0,0 +1,66 @@
+// Copyright 2018-2024 Brian Smith.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES
+// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
+// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+#![cfg(any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64"))]
+
+use super::{Block, Counter, EncryptBlock, EncryptCtr32, Iv, KeyBytes, AES_KEY};
+use crate::{cpu, error};
+use core::ops::RangeFrom;
+
+#[cfg(target_arch = "aarch64")]
+pub(in super::super) type RequiredCpuFeatures = cpu::arm::Aes;
+
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+pub(in super::super) type RequiredCpuFeatures = cpu::intel::Aes;
+
+#[derive(Clone)]
+pub struct Key {
+    inner: AES_KEY,
+}
+
+impl Key {
+    pub(in super::super) fn new(
+        bytes: KeyBytes<'_>,
+        _cpu: RequiredCpuFeatures,
+    ) -> Result<Self, error::Unspecified> {
+        let inner = unsafe { set_encrypt_key!(aes_hw_set_encrypt_key, bytes) }?;
+        Ok(Self { inner })
+    }
+
+    #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
+    #[must_use]
+    pub(in super::super) fn inner_less_safe(&self) -> &AES_KEY {
+        &self.inner
+    }
+}
+
+impl EncryptBlock for Key {
+    fn encrypt_block(&self, block: Block) -> Block {
+        super::encrypt_block_using_encrypt_iv_xor_block(self, block)
+    }
+
+    fn encrypt_iv_xor_block(&self, iv: Iv, block: Block) -> Block {
+        super::encrypt_iv_xor_block_using_ctr32(self, iv, block)
+    }
+}
+
+impl EncryptCtr32 for Key {
+    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+        #[cfg(target_arch = "x86_64")]
+        let _: cpu::Features = cpu::features();
+        unsafe {
+            ctr32_encrypt_blocks!(aes_hw_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr,)
+        }
+    }
+}

src/aead/aes/vp.rs

@@ -0,0 +1,130 @@
+// Copyright 2018-2024 Brian Smith.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES
+// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
+// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+#![cfg(any(
+    target_arch = "aarch64",
+    target_arch = "arm",
+    target_arch = "x86",
+    target_arch = "x86_64"
+))]
+
+use super::{Block, Counter, EncryptBlock, EncryptCtr32, Iv, KeyBytes, AES_KEY};
+use crate::{cpu, error};
+use core::ops::RangeFrom;
+
+#[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
+type RequiredCpuFeatures = cpu::arm::Neon;
+
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+type RequiredCpuFeatures = cpu::intel::Ssse3;
+
+#[derive(Clone)]
+pub(in super::super) struct Key {
+    inner: AES_KEY,
+}
+
+impl Key {
+    pub(in super::super) fn new(
+        bytes: KeyBytes<'_>,
+        _cpu: RequiredCpuFeatures,
+    ) -> Result<Self, error::Unspecified> {
+        let inner = unsafe { set_encrypt_key!(vpaes_set_encrypt_key, bytes) }?;
+        Ok(Self { inner })
+    }
+}
+
+#[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "x86_64"))]
+impl EncryptBlock for Key {
+    fn encrypt_block(&self, block: Block) -> Block {
+        super::encrypt_block_using_encrypt_iv_xor_block(self, block)
+    }
+
+    fn encrypt_iv_xor_block(&self, iv: Iv, block: Block) -> Block {
+        super::encrypt_iv_xor_block_using_ctr32(self, iv, block)
+    }
+}
+
+#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
+impl EncryptCtr32 for Key {
+    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+        unsafe { ctr32_encrypt_blocks!(vpaes_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr,) }
+    }
+}
+
+#[cfg(target_arch = "arm")]
+impl EncryptCtr32 for Key {
+    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+        use super::{bs, BLOCK_LEN};
+
+        let in_out = {
+            let blocks = in_out[src.clone()].len() / BLOCK_LEN;
+
+            // bsaes operates in batches of 8 blocks.
+            let bsaes_blocks = if blocks >= 8 && (blocks % 8) < 6 {
+                // It's faster to use bsaes for all the full batches and then
+                // switch to vpaes for the last partial batch (if any).
+                blocks - (blocks % 8)
+            } else if blocks >= 8 {
+                // It's faster to let bsaes handle everything including
+                // the last partial batch.
+                blocks
+            } else {
+                // It's faster to let vpaes handle everything.
+                0
+            };
+            let bsaes_in_out_len = bsaes_blocks * BLOCK_LEN;
+
+            // SAFETY:
+            //  * self.inner was initialized with `vpaes_set_encrypt_key` above,
+            //    as required by `bsaes_ctr32_encrypt_blocks_with_vpaes_key`.
+            unsafe {
+                bs::ctr32_encrypt_blocks_with_vpaes_key(
+                    &mut in_out[..(src.start + bsaes_in_out_len)],
+                    src.clone(),
+                    &self.inner,
+                    ctr,
+                );
+            }
+
+            &mut in_out[bsaes_in_out_len..]
+        };
+
+        // SAFETY:
+        //  * self.inner was initialized with `vpaes_set_encrypt_key` above,
+        //    as required by `vpaes_ctr32_encrypt_blocks`.
+        //  * `vpaes_ctr32_encrypt_blocks` satisfies the contract for
+        //    `ctr32_encrypt_blocks`.
+        unsafe { ctr32_encrypt_blocks!(vpaes_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr,) }
+    }
+}
+
+#[cfg(target_arch = "x86")]
+impl EncryptBlock for Key {
+    fn encrypt_block(&self, block: Block) -> Block {
+        unsafe { encrypt_block!(vpaes_encrypt, block, &self.inner) }
+    }
+
+    fn encrypt_iv_xor_block(&self, iv: Iv, block: Block) -> Block {
+        super::encrypt_iv_xor_block_using_encrypt_block(self, iv, block)
+    }
+}
+
+#[cfg(target_arch = "x86")]
+impl EncryptCtr32 for Key {
+    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+        super::super::shift::shift_full_blocks(in_out, src, |input| {
+            self.encrypt_iv_xor_block(ctr.increment(), *input)
+        });
+    }
+}