From 681402342bfe69f72238e50b75deb5edeff741af Mon Sep 17 00:00:00 2001
From: Brian Smith <brian@briansmith.org>
Date: Wed, 3 Apr 2024 17:47:15 -0700
Subject: [PATCH] AES-GCM: Align the types in AArch64 AES-GCM assembly with
 BoringSSL.

Make it easier to compare *ring*'s declarations with BoringSSL's for
these functions by using u64 instead of usize for the `in_bits`
argument to aes_gcm_{dec,enc}_kernel. Use `NonZeroU64` to indicate
that the assembly functions do not work with zero-length inputs.
---
 src/aead/aes_gcm.rs | 20 ++++++++++++++------
 src/bits.rs         | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/aead/aes_gcm.rs b/src/aead/aes_gcm.rs
index 898cd94353..781ec3d4af 100644
--- a/src/aead/aes_gcm.rs
+++ b/src/aead/aes_gcm.rs
@@ -127,14 +127,18 @@ fn aes_gcm_seal(
         if !aes_key.is_aes_hw(cpu_features) || !auth.is_clmul() {
             in_out
         } else {
+            use crate::bits::BitLength;
+
             let whole_block_bits = auth.in_out_whole_block_bits();
-            if whole_block_bits.as_bits() > 0 {
-                use crate::{bits::BitLength, c};
+            let whole_block_bits_u64: BitLength<u64> = whole_block_bits.into();
+            if let Ok(whole_block_bits) = whole_block_bits_u64.try_into() {
+                use core::num::NonZeroU64;
+
                 let (htable, xi) = auth.inner();
                 prefixed_extern! {
                     fn aes_gcm_enc_kernel(
                         input: *const u8,
-                        in_bits: BitLength<c::size_t>,
+                        in_bits: BitLength<NonZeroU64>,
                         output: *mut u8,
                         Xi: &mut gcm::Xi,
                         ivec: &mut Counter,
@@ -243,14 +247,18 @@ fn aes_gcm_open(
         if !aes_key.is_aes_hw(cpu_features) || !auth.is_clmul() {
             in_out
         } else {
+            use crate::bits::BitLength;
+
             let whole_block_bits = auth.in_out_whole_block_bits();
-            if whole_block_bits.as_bits() > 0 {
-                use crate::{bits::BitLength, c};
+            let whole_block_bits_u64: BitLength<u64> = whole_block_bits.into();
+            if let Ok(whole_block_bits) = whole_block_bits_u64.try_into() {
+                use core::num::NonZeroU64;
+
                 let (htable, xi) = auth.inner();
                 prefixed_extern! {
                     fn aes_gcm_dec_kernel(
                         input: *const u8,
-                        in_bits: BitLength<c::size_t>,
+                        in_bits: BitLength<NonZeroU64>,
                         output: *mut u8,
                         Xi: &mut gcm::Xi,
                         ivec: &mut Counter,
diff --git a/src/bits.rs b/src/bits.rs
index 9dacad0b81..68e2722e37 100644
--- a/src/bits.rs
+++ b/src/bits.rs
@@ -107,3 +107,18 @@ impl BitLength<u64> {
         self.0.to_be_bytes()
     }
 }
+
+#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
+impl From<BitLength<usize>> for BitLength<u64> {
+    fn from(BitLength(value): BitLength<usize>) -> Self {
+        BitLength(polyfill::u64_from_usize(value))
+    }
+}
+
+impl TryFrom<BitLength<u64>> for BitLength<core::num::NonZeroU64> {
+    type Error = <core::num::NonZeroU64 as TryFrom<u64>>::Error;
+
+    fn try_from(BitLength(value): BitLength<u64>) -> Result<Self, Self::Error> {
+        value.try_into().map(BitLength)
+    }
+}