chacha: Pass cpu::Features into encrypt.

src/aead/chacha.rs

@@ -14,6 +14,7 @@
 // CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 use super::{overlapping, quic::Sample, Nonce};
+use crate::cpu;
 
 #[cfg(any(
     test,
@@ -47,31 +48,33 @@ impl Key {
     // Encrypts `in_out` with the counter 0 and returns counter 1,
     // where the counter is derived from the nonce `nonce`.
     #[inline]
-    pub fn encrypt_single_block_with_ctr_0<const N: usize>(
+    pub(super) fn encrypt_single_block_with_ctr_0<const N: usize>(
         &self,
         nonce: Nonce,
         in_out: &mut [u8; N],
+        cpu: cpu::Features,
     ) -> Counter {
         assert!(N <= BLOCK_LEN);
         let (zero, one) = Counter::zero_one_less_safe(nonce);
-        self.encrypt(zero, in_out.as_mut().into());
+        self.encrypt(zero, in_out.as_mut().into(), cpu);
         one
     }
 
     #[inline]
     pub fn new_mask(&self, sample: Sample) -> [u8; 5] {
+        let cpu = cpu::features(); // TODO: Remove this.
         let (ctr, nonce) = sample.split_at(4);
         let ctr = u32::from_le_bytes(ctr.try_into().unwrap());
         let nonce = Nonce::assume_unique_for_key(nonce.try_into().unwrap());
         let ctr = Counter::from_nonce_and_ctr(nonce, ctr);
 
         let mut out: [u8; 5] = [0; 5];
-        self.encrypt(ctr, out.as_mut().into());
+        self.encrypt(ctr, out.as_mut().into(), cpu);
         out
     }
 
     #[inline(always)]
-    pub fn encrypt(&self, counter: Counter, in_out: Overlapping<'_>) {
+    pub(super) fn encrypt(&self, counter: Counter, in_out: Overlapping<'_>, _cpu: cpu::Features) {
         #[cfg(any(
             all(target_arch = "aarch64", target_endian = "little"),
             all(target_arch = "arm", target_endian = "little"),
@@ -198,7 +201,9 @@ mod tests {
     // Smoketest the fallback implementation.
     #[test]
     fn chacha20_test_fallback() {
-        chacha20_test(MAX_ALIGNMENT_AND_OFFSET_SUBSET, fallback::ChaCha20_ctr32);
+        chacha20_test(MAX_ALIGNMENT_AND_OFFSET_SUBSET, |key, ctr, in_out, _cpu| {
+            fallback::ChaCha20_ctr32(key, ctr, in_out)
+        });
     }
 
     // Verifies the encryption is successful when done on overlapping buffers.
@@ -210,8 +215,10 @@ mod tests {
     // works around that.
     fn chacha20_test(
         max_alignment_and_offset: (usize, usize),
-        f: impl for<'k, 'o> Fn(&'k Key, Counter, Overlapping<'o>),
+        f: impl for<'k, 'o> Fn(&'k Key, Counter, Overlapping<'o>, cpu::Features),
     ) {
+        let cpu = cpu::features();
+
         // Reuse a buffer to avoid slowing down the tests with allocations.
         let mut buf = vec![0u8; 1300];
 
@@ -240,6 +247,7 @@ mod tests {
                     &output[..len],
                     &mut buf,
                     max_alignment_and_offset,
+                    cpu,
                     &f,
                 );
             }
@@ -256,7 +264,8 @@ mod tests {
         expected: &[u8],
         buf: &mut [u8],
         (max_alignment, max_offset): (usize, usize),
-        f: &impl for<'k, 'o> Fn(&'k Key, Counter, Overlapping<'o>),
+        cpu: cpu::Features,
+        f: &impl for<'k, 'o> Fn(&'k Key, Counter, Overlapping<'o>, cpu::Features),
     ) {
         const ARBITRARY: u8 = 123;
 
@@ -276,7 +285,7 @@ mod tests {
                 let in_out = Overlapping::new(buf, src)
                     .map_err(error::erase::<IndexError>)
                     .unwrap();
-                f(key, ctr, in_out);
+                f(key, ctr, in_out, cpu);
                 assert_eq!(&buf[..input.len()], expected)
             }
         }

src/aead/chacha20_poly1305/mod.rs

@@ -71,13 +71,13 @@ pub(super) fn seal_fallback(
     nonce: Nonce,
     aad: Aad<&[u8]>,
     in_out: &mut [u8],
-    cpu_features: cpu::Features,
+    cpu: cpu::Features,
 ) -> Result<Tag, InputTooLongError> {
-    let (counter, poly1305_key) = begin(chacha20_key, nonce, aad, in_out)?;
-    let mut auth = poly1305::Context::from_key(poly1305_key, cpu_features);
+    let (counter, poly1305_key) = begin(chacha20_key, nonce, aad, in_out, cpu)?;
+    let mut auth = poly1305::Context::from_key(poly1305_key, cpu);
 
     poly1305_update_padded_16(&mut auth, aad.as_ref());
-    chacha20_key.encrypt(counter, in_out.into());
+    chacha20_key.encrypt(counter, in_out.into(), cpu);
     poly1305_update_padded_16(&mut auth, in_out);
     Ok(finish(auth, aad.as_ref().len(), in_out.len()))
 }
@@ -105,15 +105,15 @@ pub(super) fn open_fallback(
     nonce: Nonce,
     aad: Aad<&[u8]>,
     in_out: Overlapping<'_>,
-    cpu_features: cpu::Features,
+    cpu: cpu::Features,
 ) -> Result<Tag, InputTooLongError> {
-    let (counter, poly1305_key) = begin(chacha20_key, nonce, aad, in_out.input())?;
-    let mut auth = poly1305::Context::from_key(poly1305_key, cpu_features);
+    let (counter, poly1305_key) = begin(chacha20_key, nonce, aad, in_out.input(), cpu)?;
+    let mut auth = poly1305::Context::from_key(poly1305_key, cpu);
 
     poly1305_update_padded_16(&mut auth, aad.as_ref());
     poly1305_update_padded_16(&mut auth, in_out.input());
     let in_out_len = in_out.len();
-    chacha20_key.encrypt(counter, in_out);
+    chacha20_key.encrypt(counter, in_out, cpu);
     Ok(finish(auth, aad.as_ref().len(), in_out_len))
 }
 
@@ -137,11 +137,12 @@ pub(super) fn begin(
     nonce: Nonce,
     aad: Aad<&[u8]>,
     input: &[u8],
+    cpu: cpu::Features,
 ) -> Result<(Counter, poly1305::Key), InputTooLongError> {
     check_input_lengths(aad, input)?;
 
     let mut key_bytes = [0u8; poly1305::KEY_LEN];
-    let counter = key.encrypt_single_block_with_ctr_0(nonce, &mut key_bytes);
+    let counter = key.encrypt_single_block_with_ctr_0(nonce, &mut key_bytes, cpu);
     let poly1305_key = poly1305::Key::new(key_bytes);
     Ok((counter, poly1305_key))
 }

src/aead/chacha20_poly1305_openssh.rs

@@ -77,26 +77,28 @@ impl SealingKey {
         let (len_in_out, data_and_padding_in_out): (&mut [u8; PACKET_LENGTH_LEN], _) =
             slice::split_first_chunk_mut(plaintext_in_ciphertext_out).unwrap();
 
-        let cpu_features = cpu::features();
+        let cpu = cpu::features();
         // XXX/TODO(SemVer): Refactor API to return an error.
         let (counter, poly_key) = chacha20_poly1305::begin(
             &self.key.k_2,
             make_nonce(sequence_number),
             Aad::from(len_in_out),
             data_and_padding_in_out,
+            cpu,
         )
         .map_err(error::erase::<InputTooLongError>)
         .unwrap();
 
-        let _: Counter = self
-            .key
-            .k_1
-            .encrypt_single_block_with_ctr_0(make_nonce(sequence_number), len_in_out);
+        let _: Counter = self.key.k_1.encrypt_single_block_with_ctr_0(
+            make_nonce(sequence_number),
+            len_in_out,
+            cpu,
+        );
         self.key
             .k_2
-            .encrypt(counter, data_and_padding_in_out.into());
+            .encrypt(counter, data_and_padding_in_out.into(), cpu);
 
-        let Tag(tag) = poly1305::sign(poly_key, plaintext_in_ciphertext_out, cpu_features);
+        let Tag(tag) = poly1305::sign(poly_key, plaintext_in_ciphertext_out, cpu);
         *tag_out = tag;
     }
 }
@@ -123,11 +125,13 @@ impl OpeningKey {
         sequence_number: u32,
         encrypted_packet_length: [u8; PACKET_LENGTH_LEN],
     ) -> [u8; PACKET_LENGTH_LEN] {
+        let cpu = cpu::features();
         let mut packet_length = encrypted_packet_length;
-        let _: Counter = self
-            .key
-            .k_1
-            .encrypt_single_block_with_ctr_0(make_nonce(sequence_number), &mut packet_length);
+        let _: Counter = self.key.k_1.encrypt_single_block_with_ctr_0(
+            make_nonce(sequence_number),
+            &mut packet_length,
+            cpu,
+        );
         packet_length
     }
 
@@ -155,6 +159,7 @@ impl OpeningKey {
             make_nonce(sequence_number),
             Aad::from(packet_length),
             after_packet_length,
+            cpu,
         )
         .map_err(error::erase::<InputTooLongError>)?;
 
@@ -167,7 +172,9 @@ impl OpeningKey {
         // Won't panic because the length was checked above.
         let after_packet_length = &mut ciphertext_in_plaintext_out[PACKET_LENGTH_LEN..];
 
-        self.key.k_2.encrypt(counter, after_packet_length.into());
+        self.key
+            .k_2
+            .encrypt(counter, after_packet_length.into(), cpu);
 
         Ok(after_packet_length)
     }