Small speedups for AES-GCM:

* Eliminate a copy in the GHASH(length) computation at the end.
* In the tag computation, instead of copying the 128-bit GHASH
  result into a u8 array and doing a bytewise XOR loop, load
  it into a u128 and do a vector XOR.

Author: brian-pane

graviola/src/low/aarch64/cpu.rs

@@ -1,6 +1,9 @@
 // Written for Graviola by Joe Birr-Pixton, 2024.
 // SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0
 
+use core::arch::aarch64::{uint8x16_t, vgetq_lane_u64, vreinterpretq_u64_u8};
+use core::mem;
+
 pub(crate) fn enter_cpu_state() -> u32 {
     dit::maybe_enable()
 }
@@ -155,6 +158,20 @@ pub(in crate::low) unsafe fn ct_compare_bytes(a: *const u8, b: *const u8, len: u
     }
 }
 
+/// Compare `a` and `b` in constant time, returning 0 if they are equal
+/// and some nonzero value if they are not equal.
+#[target_feature(enable = "neon")]
+pub(in crate::low) fn ct_compare_u128(a: u128, b: u128) -> u64 {
+    // SAFETY: u128 and uint8x16_t have the compatible memory layouts.
+    let a: uint8x16_t = unsafe { mem::transmute(a) };
+    let a = vreinterpretq_u64_u8(a);
+    // SAFETY: u128 and uint8x16_t have the compatible memory layouts.
+    let b: uint8x16_t = unsafe { mem::transmute(b) };
+    let b = vreinterpretq_u64_u8(b);
+    (vgetq_lane_u64::<0>(a) ^ vgetq_lane_u64::<0>(b)) |
+        (vgetq_lane_u64::<1>(a) ^ vgetq_lane_u64::<1>(b))
+}
+
 /// This macro interdicts is_aarch64_feature_detected to allow testability.
 macro_rules! have_cpu_feature {
     ("neon") => {

graviola/src/low/aarch64/ghash.rs

@@ -102,11 +102,18 @@ impl<'a> Ghash<'a> {
         }
     }
 
-    pub(crate) fn into_bytes(self) -> [u8; 16] {
-        to_u128(self.current).to_be_bytes()
+    #[cfg(test)]
+    fn into_bytes(self) -> [u8; 16] {
+        self.into_u128().to_be_bytes()
     }
 
-    fn one_block(&mut self, block: u128) {
+    pub(crate) fn into_u128(self) -> u128 {
+        // SAFETY: u128 and uint64x2_t have the same size and meaning of bits
+        unsafe { mem::transmute(self.current) }
+    }
+
+    // Input the 16 bytes of `block` to the computation.
+    pub(crate) fn one_block(&mut self, block: u128) {
         // SAFETY: this crate requires the `neon` cpu feature
         self.current = unsafe { veorq_u64(self.current, from_u128(block)) };
         self.current = mul(self.current, self.table.powers[0]);
@@ -326,12 +333,6 @@ fn from_u128(u: u128) -> uint64x2_t {
     unsafe { mem::transmute(u) }
 }
 
-#[inline]
-fn to_u128(u: uint64x2_t) -> u128 {
-    // SAFETY: u128 and uint64x2_t have the same size and meaning of bits
-    unsafe { mem::transmute(u) }
-}
-
 #[inline]
 #[target_feature(enable = "neon")]
 // Make a copy of `u` with the bytes reversed, and cast to `uint64x2_t`.

graviola/src/low/generic/ct_equal.rs

@@ -1,7 +1,7 @@
 // Written for Graviola by Joe Birr-Pixton, 2024.
 // SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0
 
-use crate::low::ct_compare_bytes;
+use crate::low::{ct_compare_bytes, ct_compare_u128};
 
 pub(crate) fn ct_equal(a: &[u8], b: &[u8]) -> bool {
     if a.len() != b.len() {
@@ -12,3 +12,7 @@ pub(crate) fn ct_equal(a: &[u8], b: &[u8]) -> bool {
     let diff = unsafe { ct_compare_bytes(a.as_ptr(), b.as_ptr(), a.len()) };
     diff == 0
 }
+
+pub(crate) fn ct_equal_u128(a: u128, b: u128) -> bool {
+    (unsafe { ct_compare_u128(a, b) }) == 0
+}

graviola/src/low/mod.rs

@@ -32,7 +32,7 @@ mod posint;
 pub(crate) use entry::Entry;
 pub(crate) use generic::blockwise::Blockwise;
 pub(crate) use generic::ct_copy::{ct_copy, ct_select_i16};
-pub(crate) use generic::ct_equal::ct_equal;
+pub(crate) use generic::ct_equal::{ct_equal, ct_equal_u128};
 pub(crate) use generic::poly1305;
 pub(crate) use generic::zeroise::{zeroise, zeroise_value};
 pub(crate) use posint::{PosInt, SecretPosInt};
@@ -45,7 +45,7 @@ cfg_if::cfg_if! {
     if #[cfg(target_arch = "x86_64")] {
         mod x86_64;
 
-        pub(in crate::low) use x86_64::cpu::{enter_cpu_state, zero_bytes, ct_compare_bytes, leave_cpu_state, verify_cpu_features};
+        pub(in crate::low) use x86_64::cpu::{enter_cpu_state, zero_bytes, ct_compare_bytes, ct_compare_u128, leave_cpu_state, verify_cpu_features};
         pub(crate) use x86_64::chacha20;
         pub(crate) use x86_64::aes::AesKey;
         pub(crate) use x86_64::aes_gcm;
@@ -127,7 +127,7 @@ cfg_if::cfg_if! {
     } else if #[cfg(target_arch = "aarch64")] {
         mod aarch64;
 
-        pub(in crate::low) use aarch64::cpu::{enter_cpu_state, zero_bytes, ct_compare_bytes, leave_cpu_state, verify_cpu_features};
+        pub(in crate::low) use aarch64::cpu::{enter_cpu_state, zero_bytes, ct_compare_bytes, ct_compare_u128, leave_cpu_state, verify_cpu_features};
         pub(crate) use aarch64::aes::AesKey;
         pub(crate) use aarch64::aes_gcm;
         pub(crate) use aarch64::bignum_add::bignum_add;

graviola/src/low/x86_64/cpu.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0
 
 use core::arch::x86_64::*;
+use core::mem;
 
 pub(crate) fn enter_cpu_state() -> u32 {
     // DOIT: "Data Operand Independent Timing" -- turning this on
@@ -115,6 +116,17 @@ pub(in crate::low) unsafe fn ct_compare_bytes(a: *const u8, b: *const u8, len: u
     acc
 }
 
+/// Compare `a` and `b` in constant time, returning 0 if they are equal
+/// and some nonzero value if they are not equal.
+#[target_feature(enable = "sse2")]
+pub(in crate::low) fn ct_compare_u128(a: u128, b: u128) -> u64 {
+    let a: __m128i = unsafe { mem::transmute(a) };
+    let b: __m128i = unsafe { mem::transmute(b) };
+    let cmp = _mm_cmpeq_epi8(a, b);
+    let mask = _mm_movemask_epi8(cmp) as u16;
+    !mask as _
+}
+
 /// This macro interdicts is_x86_feature_detected to
 /// allow testability.
 macro_rules! have_cpu_feature {

graviola/src/low/x86_64/ghash.rs

@@ -206,7 +206,7 @@ impl<'a> Ghash<'a> {
 
         for chunk in whole_blocks.by_ref() {
             let u = u128::from_be_bytes(chunk.try_into().unwrap());
-            self.one_block(u128_to_m128i(u));
+            self.one_block(u);
         }
 
         let bytes = whole_blocks.remainder();
@@ -215,15 +215,22 @@ impl<'a> Ghash<'a> {
             block[..bytes.len()].copy_from_slice(bytes);
 
             let u = u128::from_be_bytes(block);
-            self.one_block(u128_to_m128i(u));
+            self.one_block(u);
         }
     }
 
+    #[cfg(test)]
     pub(crate) fn into_bytes(self) -> [u8; 16] {
         // SAFETY: this crate requires the `sse2` and `ssse3` cpu features
         unsafe { self._into_bytes() }
     }
 
+    pub(crate) fn into_u128(self) -> u128 {
+        // SAFETY: sizeof(u128) == sizeof(__m128i), all bits have same meaning
+        unsafe { mem::transmute(self.current) }
+    }
+
+    #[cfg(test)]
     #[target_feature(enable = "sse2,ssse3")]
     fn _into_bytes(self) -> [u8; 16] {
         let mut out: i128 = 0;
@@ -233,7 +240,9 @@ impl<'a> Ghash<'a> {
         out.to_le_bytes()
     }
 
-    fn one_block(&mut self, block: __m128i) {
+    // Input the 16 bytes of `block` to the computation.
+    pub(crate) fn one_block(&mut self, block: u128) {
+        let block = u128_to_m128i(block);
         // SAFETY: this crate requires the `avx` and `pclmulqdq` cpu features
         unsafe {
             self.current = _mm_xor_si128(self.current, block);

graviola/src/mid/aes_gcm.rs

@@ -3,7 +3,7 @@
 
 use crate::Error;
 use crate::low::ghash::{Ghash, GhashTable};
-use crate::low::{AesKey, Entry, aes_gcm, ct_equal};
+use crate::low::{AesKey, Entry, aes_gcm, ct_equal_u128};
 
 /// An AES-GCM key.
 ///
@@ -66,16 +66,13 @@ impl AesGcm {
         // computations. see low::generic::aes_gcm for model version.
         aes_gcm::encrypt(&self.key, &mut ghash, &counter, aad, cipher_inout);
 
-        let mut lengths = [0u8; 16];
-        lengths[..8].copy_from_slice(&((aad.len() * 8) as u64).to_be_bytes());
-        lengths[8..].copy_from_slice(&((cipher_inout.len() * 8) as u64).to_be_bytes());
-        ghash.add(&lengths);
+        let lengths: u128 = Self::pack_bit_lengths(aad, cipher_inout);
+        ghash.one_block(lengths);
 
-        let final_xi = ghash.into_bytes();
-
-        for ((out, x), e) in tag_out.iter_mut().zip(final_xi.iter()).zip(e_y0.iter()) {
-            *out = *x ^ *e;
-        }
+        let final_xi = ghash.into_u128();
+        let e = u128::from_be_bytes(e_y0);
+        let tag = final_xi ^ e;
+        tag_out.copy_from_slice(&tag.to_be_bytes());
     }
 
     /// Decrypts and verifies the given message.
@@ -107,23 +104,33 @@ impl AesGcm {
 
         aes_gcm::decrypt(&self.key, &mut ghash, &counter, aad, cipher_inout);
 
-        let mut lengths = [0u8; 16];
-        lengths[..8].copy_from_slice(&((aad.len() * 8) as u64).to_be_bytes());
-        lengths[8..].copy_from_slice(&((cipher_inout.len() * 8) as u64).to_be_bytes());
-        ghash.add(&lengths);
-
-        let mut actual_tag = ghash.into_bytes();
-        for (out, e) in actual_tag.iter_mut().zip(e_y0.iter()) {
-            *out ^= *e;
+        let lengths: u128 = Self::pack_bit_lengths(aad, cipher_inout);
+        ghash.one_block(lengths);
+
+        let actual_tag = ghash.into_u128() ^ u128::from_be_bytes(e_y0);
+        let tag: Result<[u8; 16], _> = tag.try_into();
+        match tag {
+            Ok(tag) => {
+                // TODO add a constant-time u128 comparison function
+                if ct_equal_u128(u128::from_be_bytes(tag), actual_tag) {
+                    Ok(())
+                } else {
+                    // avoid unauthenticated plaintext leak
+                    cipher_inout.fill(0x00);
+                    Err(Error::DecryptFailed)
+                }
+            },
+            Err(_) => {
+                cipher_inout.fill(0x00);
+                Err(Error::DecryptFailed)
+            }
         }
+    }
 
-        if ct_equal(&actual_tag, tag) {
-            Ok(())
-        } else {
-            // avoid unauthenticated plaintext leak
-            cipher_inout.fill(0x00);
-            Err(Error::DecryptFailed)
-        }
+    fn pack_bit_lengths(high: &[u8], low: &[u8]) -> u128 {
+        let lengths = (high.len() * 8) as u128;
+        let lengths = lengths << 64;
+        lengths | ((low.len() * 8) as u128)
     }
 
     fn nonce_to_y0(&self, nonce: &[u8; 12]) -> [u8; 16] {