AEAD: Don't store cpu::Features in the key.

Have the inner AEAD API take `cpu::features()` for all operations.
Then we will be able to write CPU-capability-based tests using (a
variation of) the inner API, which will (when implemented) eliminate
the need to use SDE and the other various hacks we use for testing all
the implementations.

src/aead.rs

@@ -134,13 +134,20 @@ impl hkdf::KeyType for &'static Algorithm {
 pub struct Algorithm {
     init: fn(key: &[u8], cpu_features: cpu::Features) -> Result<KeyInner, error::Unspecified>,
 
-    seal: fn(key: &KeyInner, nonce: Nonce, aad: Aad<&[u8]>, in_out: &mut [u8]) -> Tag,
+    seal: fn(
+        key: &KeyInner,
+        nonce: Nonce,
+        aad: Aad<&[u8]>,
+        in_out: &mut [u8],
+        cpu_features: cpu::Features,
+    ) -> Tag,
     open: fn(
         key: &KeyInner,
         nonce: Nonce,
         aad: Aad<&[u8]>,
         in_out: &mut [u8],
         src: RangeFrom<usize>,
+        cpu_features: cpu::Features,
     ) -> Tag,
 
     key_len: usize,

src/aead/aes.rs

@@ -29,7 +29,6 @@ use core::ops::RangeFrom;
 #[derive(Clone)]
 pub(super) struct Key {
     inner: AES_KEY,
-    cpu_features: cpu::Features,
 }
 
 macro_rules! set_encrypt_key {
@@ -172,15 +171,12 @@ impl Key {
             }
         };
 
-        Ok(Self {
-            inner: key,
-            cpu_features,
-        })
+        Ok(Self { inner: key })
     }
 
     #[inline]
-    pub fn encrypt_block(&self, a: Block) -> Block {
-        match detect_implementation(self.cpu_features) {
+    pub fn encrypt_block(&self, a: Block, cpu_features: cpu::Features) -> Block {
+        match detect_implementation(cpu_features) {
             #[cfg(any(
                 target_arch = "aarch64",
                 target_arch = "arm",
@@ -203,8 +199,8 @@ impl Key {
     }
 
     #[inline]
-    pub fn encrypt_iv_xor_block(&self, iv: Iv, input: Block) -> Block {
-        let encrypted_iv = self.encrypt_block(iv.into_block_less_safe());
+    pub fn encrypt_iv_xor_block(&self, iv: Iv, input: Block, cpu_features: cpu::Features) -> Block {
+        let encrypted_iv = self.encrypt_block(iv.into_block_less_safe(), cpu_features);
         encrypted_iv ^ input
     }
 
@@ -214,12 +210,13 @@ impl Key {
         in_out: &mut [u8],
         src: RangeFrom<usize>,
         ctr: &mut Counter,
+        cpu_features: cpu::Features,
     ) {
         let in_out_len = in_out[src.clone()].len();
 
         assert_eq!(in_out_len % BLOCK_LEN, 0);
 
-        match detect_implementation(self.cpu_features) {
+        match detect_implementation(cpu_features) {
             #[cfg(any(
                 target_arch = "aarch64",
                 target_arch = "arm",
@@ -271,7 +268,7 @@ impl Key {
             #[cfg(target_arch = "x86")]
             Implementation::VPAES_BSAES => {
                 super::shift::shift_full_blocks(in_out, src, |input| {
-                    self.encrypt_iv_xor_block(ctr.increment(), Block::from(input))
+                    self.encrypt_iv_xor_block(ctr.increment(), Block::from(input), cpu_features)
                 });
             }
 
@@ -283,7 +280,7 @@ impl Key {
     }
 
     pub fn new_mask(&self, sample: Sample) -> [u8; 5] {
-        let block = self.encrypt_block(Block::from(&sample));
+        let block = self.encrypt_block(Block::from(&sample), cpu::features());
 
         let mut out: [u8; 5] = [0; 5];
         out.copy_from_slice(&block.as_ref()[..5]);
@@ -293,11 +290,8 @@ impl Key {
 
     #[cfg(target_arch = "x86_64")]
     #[must_use]
-    pub fn is_aes_hw(&self) -> bool {
-        matches!(
-            detect_implementation(self.cpu_features),
-            Implementation::HWAES
-        )
+    pub fn is_aes_hw(&self, cpu_features: cpu::Features) -> bool {
+        matches!(detect_implementation(cpu_features), Implementation::HWAES)
     }
 
     #[cfg(target_arch = "x86_64")]
@@ -445,6 +439,7 @@ mod tests {
 
     #[test]
     pub fn test_aes() {
+        let cpu_features = cpu::features();
         test::run(test_file!("aes_tests.txt"), |section, test_case| {
             assert_eq!(section, "");
             let key = consume_key(test_case, "Key");
@@ -453,7 +448,7 @@ mod tests {
             let expected_output = test_case.consume_bytes("Output");
 
             let block = Block::from(input);
-            let output = key.encrypt_block(block);
+            let output = key.encrypt_block(block, cpu_features);
             assert_eq!(output.as_ref(), &expected_output[..]);
 
             Ok(())

src/aead/aes_gcm.rs

@@ -63,13 +63,22 @@ fn init(
     cpu_features: cpu::Features,
 ) -> Result<aead::KeyInner, error::Unspecified> {
     let aes_key = aes::Key::new(key, variant, cpu_features)?;
-    let gcm_key = gcm::Key::new(aes_key.encrypt_block(Block::zero()), cpu_features);
+    let gcm_key = gcm::Key::new(
+        aes_key.encrypt_block(Block::zero(), cpu_features),
+        cpu_features,
+    );
     Ok(aead::KeyInner::AesGcm(Key { gcm_key, aes_key }))
 }
 
 const CHUNK_BLOCKS: usize = 3 * 1024 / 16;
 
-fn aes_gcm_seal(key: &aead::KeyInner, nonce: Nonce, aad: Aad<&[u8]>, in_out: &mut [u8]) -> Tag {
+fn aes_gcm_seal(
+    key: &aead::KeyInner,
+    nonce: Nonce,
+    aad: Aad<&[u8]>,
+    in_out: &mut [u8],
+    cpu_features: cpu::Features,
+) -> Tag {
     let Key { gcm_key, aes_key } = match key {
         aead::KeyInner::AesGcm(key) => key,
         _ => unreachable!(),
@@ -80,11 +89,11 @@ fn aes_gcm_seal(key: &aead::KeyInner, nonce: Nonce, aad: Aad<&[u8]>, in_out: &mu
 
     let total_in_out_len = in_out.len();
     let aad_len = aad.0.len();
-    let mut auth = gcm::Context::new(gcm_key, aad);
+    let mut auth = gcm::Context::new(gcm_key, aad, cpu_features);
 
     #[cfg(target_arch = "x86_64")]
     let in_out = {
-        if !aes_key.is_aes_hw() || !auth.is_avx() {
+        if !aes_key.is_aes_hw(cpu_features) || !auth.is_avx() {
             in_out
         } else {
             use crate::c;
@@ -124,20 +133,27 @@ fn aes_gcm_seal(key: &aead::KeyInner, nonce: Nonce, aad: Aad<&[u8]>, in_out: &mu
     };
 
     for chunk in whole.chunks_mut(CHUNK_BLOCKS * BLOCK_LEN) {
-        aes_key.ctr32_encrypt_within(chunk, 0.., &mut ctr);
+        aes_key.ctr32_encrypt_within(chunk, 0.., &mut ctr, cpu_features);
         auth.update_blocks(chunk);
     }
 
     if !remainder.is_empty() {
         let mut input = Block::zero();
         input.overwrite_part_at(0, remainder);
-        let mut output = aes_key.encrypt_iv_xor_block(ctr.into(), input);
+        let mut output = aes_key.encrypt_iv_xor_block(ctr.into(), input, cpu_features);
         output.zero_from(remainder.len());
         auth.update_block(output);
         remainder.copy_from_slice(&output.as_ref()[..remainder.len()]);
     }
 
-    finish(aes_key, auth, tag_iv, aad_len, total_in_out_len)
+    finish(
+        aes_key,
+        auth,
+        tag_iv,
+        aad_len,
+        total_in_out_len,
+        cpu_features,
+    )
 }
 
 fn aes_gcm_open(
@@ -146,6 +162,7 @@ fn aes_gcm_open(
     aad: Aad<&[u8]>,
     in_out: &mut [u8],
     src: RangeFrom<usize>,
+    cpu_features: cpu::Features,
 ) -> Tag {
     let Key { gcm_key, aes_key } = match key {
         aead::KeyInner::AesGcm(key) => key,
@@ -156,15 +173,15 @@ fn aes_gcm_open(
     let tag_iv = ctr.increment();
 
     let aad_len = aad.0.len();
-    let mut auth = gcm::Context::new(gcm_key, aad);
+    let mut auth = gcm::Context::new(gcm_key, aad, cpu_features);
 
     let in_prefix_len = src.start;
 
     let total_in_out_len = in_out.len() - in_prefix_len;
 
     #[cfg(target_arch = "x86_64")]
     let in_out = {
-        if !aes_key.is_aes_hw() || !auth.is_avx() {
+        if !aes_key.is_aes_hw(cpu_features) || !auth.is_avx() {
             in_out
         } else {
             use crate::c;
@@ -218,6 +235,7 @@ fn aes_gcm_open(
                 &mut in_out[output..][..(chunk_len + in_prefix_len)],
                 in_prefix_len..,
                 &mut ctr,
+                cpu_features,
             );
             output += chunk_len;
             input += chunk_len;
@@ -229,10 +247,17 @@ fn aes_gcm_open(
         let mut input = Block::zero();
         input.overwrite_part_at(0, remainder);
         auth.update_block(input);
-        aes_key.encrypt_iv_xor_block(ctr.into(), input)
+        aes_key.encrypt_iv_xor_block(ctr.into(), input, cpu_features)
     });
 
-    finish(aes_key, auth, tag_iv, aad_len, total_in_out_len)
+    finish(
+        aes_key,
+        auth,
+        tag_iv,
+        aad_len,
+        total_in_out_len,
+        cpu_features,
+    )
 }
 
 fn finish(
@@ -241,6 +266,7 @@ fn finish(
     tag_iv: aes::Iv,
     aad_len: usize,
     in_out_len: usize,
+    cpu_features: cpu::Features,
 ) -> Tag {
     // Authenticate the final block containing the input lengths.
     let aad_bits = polyfill::u64_from_usize(aad_len) << 3;
@@ -251,7 +277,7 @@ fn finish(
 
     // Finalize the tag and return it.
     gcm_ctx.pre_finish(|pre_tag| {
-        let encrypted_iv = aes_key.encrypt_block(tag_iv.into_block_less_safe());
+        let encrypted_iv = aes_key.encrypt_block(tag_iv.into_block_less_safe(), cpu_features);
         let tag = pre_tag ^ encrypted_iv;
         Tag(*tag.as_ref())
     })

src/aead/chacha.rs

@@ -14,7 +14,6 @@
 // CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 use super::{quic::Sample, Nonce};
-use crate::cpu;
 
 #[cfg(any(
     test,
@@ -33,20 +32,14 @@ use core::ops::RangeFrom;
 #[derive(Clone)]
 pub struct Key {
     words: [u32; KEY_LEN / 4],
-    cpu_features: cpu::Features,
 }
 
 impl Key {
-    pub(super) fn new(value: [u8; KEY_LEN], cpu_features: cpu::Features) -> Self {
+    pub(super) fn new(value: [u8; KEY_LEN]) -> Self {
         Self {
             words: value.array_split_map(u32::from_le_bytes),
-            cpu_features,
         }
     }
-
-    pub(super) fn cpu_features(&self) -> cpu::Features {
-        self.cpu_features
-    }
 }
 
 impl Key {
@@ -261,7 +254,7 @@ mod tests {
 
             let key = test_case.consume_bytes("Key");
             let key: &[u8; KEY_LEN] = key.as_slice().try_into()?;
-            let key = Key::new(*key, cpu::features());
+            let key = Key::new(*key);
 
             let ctr = test_case.consume_usize("Ctr");
             let nonce = test_case.consume_bytes("Nonce");

src/aead/chacha20_poly1305.rs

@@ -39,20 +39,18 @@ pub static CHACHA20_POLY1305: aead::Algorithm = aead::Algorithm {
 /// Copies |key| into |ctx_buf|.
 fn chacha20_poly1305_init(
     key: &[u8],
-    cpu_features: cpu::Features,
+    _cpu_features: cpu::Features,
 ) -> Result<aead::KeyInner, error::Unspecified> {
     let key: [u8; chacha::KEY_LEN] = key.try_into()?;
-    Ok(aead::KeyInner::ChaCha20Poly1305(chacha::Key::new(
-        key,
-        cpu_features,
-    )))
+    Ok(aead::KeyInner::ChaCha20Poly1305(chacha::Key::new(key)))
 }
 
 fn chacha20_poly1305_seal(
     key: &aead::KeyInner,
     nonce: Nonce,
     aad: Aad<&[u8]>,
     in_out: &mut [u8],
+    cpu_features: cpu::Features,
 ) -> Tag {
     let chacha20_key = match key {
         aead::KeyInner::ChaCha20Poly1305(key) => key,
@@ -61,9 +59,7 @@ fn chacha20_poly1305_seal(
 
     #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
     {
-        if cpu::intel::SSE41.available(chacha20_key.cpu_features())
-            || cpu::arm::NEON.available(chacha20_key.cpu_features())
-        {
+        if cpu::intel::SSE41.available(cpu_features) || cpu::arm::NEON.available(cpu_features) {
             // XXX: BoringSSL uses `alignas(16)` on `key` instead of on the
             // structure, but Rust can't do that yet; see
             // https://github.com/rust-lang/rust/issues/73557.
@@ -121,7 +117,7 @@ fn chacha20_poly1305_seal(
     let mut counter = Counter::zero(nonce);
     let mut auth = {
         let key = derive_poly1305_key(chacha20_key, counter.increment());
-        poly1305::Context::from_key(key)
+        poly1305::Context::from_key(key, cpu_features)
     };
 
     poly1305_update_padded_16(&mut auth, aad.as_ref());
@@ -136,6 +132,7 @@ fn chacha20_poly1305_open(
     aad: Aad<&[u8]>,
     in_out: &mut [u8],
     src: RangeFrom<usize>,
+    cpu_features: cpu::Features,
 ) -> Tag {
     let chacha20_key = match key {
         aead::KeyInner::ChaCha20Poly1305(key) => key,
@@ -144,9 +141,7 @@ fn chacha20_poly1305_open(
 
     #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
     {
-        if cpu::intel::SSE41.available(chacha20_key.cpu_features())
-            || cpu::arm::NEON.available(chacha20_key.cpu_features())
-        {
+        if cpu::intel::SSE41.available(cpu_features) || cpu::arm::NEON.available(cpu_features) {
             // XXX: BoringSSL uses `alignas(16)` on `key` instead of on the
             // structure, but Rust can't do that yet; see
             // https://github.com/rust-lang/rust/issues/73557.
@@ -200,7 +195,7 @@ fn chacha20_poly1305_open(
     let mut counter = Counter::zero(nonce);
     let mut auth = {
         let key = derive_poly1305_key(chacha20_key, counter.increment());
-        poly1305::Context::from_key(key)
+        poly1305::Context::from_key(key, cpu_features)
     };
 
     poly1305_update_padded_16(&mut auth, aad.as_ref());
@@ -258,7 +253,7 @@ fn poly1305_update_padded_16(ctx: &mut poly1305::Context, input: &[u8]) {
 pub(super) fn derive_poly1305_key(chacha_key: &chacha::Key, iv: Iv) -> poly1305::Key {
     let mut key_bytes = [0u8; poly1305::KEY_LEN];
     chacha_key.encrypt_iv_xor_in_place(iv, &mut key_bytes);
-    poly1305::Key::new(key_bytes, chacha_key.cpu_features())
+    poly1305::Key::new(key_bytes)
 }
 
 #[cfg(test)]