aes: Clarify counter overflow checking.

Create a more robust internal API for counter/nonce/IV management that makes the usage within AES-GCM more clearly correct. The new design is easier to test.
briansmith · Jun 23, 2024 · 3d10ef3 · 3d10ef3
1 parent 525047a
commit 3d10ef3
Show file tree

Hide file tree

Showing 8 changed files with 185 additions and 96 deletions.
diff --git a/src/aead/aes.rs b/src/aead/aes.rs
@@ -12,7 +12,7 @@
 // OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 // CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
-use super::{nonce::Nonce, quic::Sample, NONCE_LEN};
+use super::quic::Sample;
 use crate::{
     constant_time,
     cpu::{self, GetFeature as _},
@@ -21,12 +21,16 @@ use crate::{
 use cfg_if::cfg_if;
 use core::ops::RangeFrom;
 
-pub(super) use ffi::Counter;
+pub(super) use self::{
+    counter::{CounterOverflowError, Iv, IvBlock},
+    ffi::Counter,
+};
 
 #[macro_use]
 mod ffi;
 
 mod bs;
+mod counter;
 pub(super) mod fallback;
 pub(super) mod hw;
 pub(super) mod vp;
@@ -113,38 +117,11 @@ pub enum KeyBytes<'a> {
     AES_256(&'a [u8; AES_256_KEY_LEN]),
 }
 
-// `Counter` is `ffi::Counter` as its representation is dictated by its use in
-// the FFI.
-impl Counter {
-    pub fn one(nonce: Nonce) -> Self {
-        let mut value = [0u8; BLOCK_LEN];
-        value[..NONCE_LEN].copy_from_slice(nonce.as_ref());
-        value[BLOCK_LEN - 1] = 1;
-        Self(value)
-    }
-
-    pub fn increment(&mut self) -> Iv {
-        let iv = Iv(self.0);
-        self.increment_by_less_safe(1);
-        iv
-    }
-
-    fn increment_by_less_safe(&mut self, increment_by: u32) {
-        let [.., c0, c1, c2, c3] = &mut self.0;
-        let old_value: u32 = u32::from_be_bytes([*c0, *c1, *c2, *c3]);
-        let new_value = old_value + increment_by;
-        [*c0, *c1, *c2, *c3] = u32::to_be_bytes(new_value);
-    }
-}
-
-/// The IV for a single block encryption.
-///
-/// Intentionally not `Clone` to ensure each is used only once.
-pub struct Iv(Block);
-
-impl From<Counter> for Iv {
-    fn from(counter: Counter) -> Self {
-        Self(counter.0)
+pub(super) struct InOutLenInconsistentWithIvBlockLenError(());
+impl InOutLenInconsistentWithIvBlockLenError {
+    #[cold]
+    fn new() -> Self {
+        Self(())
     }
 }
 
@@ -158,28 +135,35 @@ pub(super) trait EncryptBlock {
 }
 
 pub(super) trait EncryptCtr32 {
-    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter);
+    fn ctr32_encrypt_within(
+        &self,
+        in_out: &mut [u8],
+        src: RangeFrom<usize>,
+        iv_block: IvBlock,
+    ) -> Result<(), InOutLenInconsistentWithIvBlockLenError>;
 }
 
 #[allow(dead_code)]
 fn encrypt_block_using_encrypt_iv_xor_block(key: &impl EncryptBlock, block: Block) -> Block {
-    key.encrypt_iv_xor_block(Iv(block), ZERO_BLOCK)
+    key.encrypt_iv_xor_block(Iv::new_less_safe(block), ZERO_BLOCK)
 }
 
 fn encrypt_iv_xor_block_using_encrypt_block(
     key: &impl EncryptBlock,
     iv: Iv,
     block: Block,
 ) -> Block {
-    let encrypted_iv = key.encrypt_block(iv.0);
+    let encrypted_iv = key.encrypt_block(iv.into_block_less_safe());
     constant_time::xor_16(encrypted_iv, block)
 }
 
 #[allow(dead_code)]
 fn encrypt_iv_xor_block_using_ctr32(key: &impl EncryptCtr32, iv: Iv, mut block: Block) -> Block {
-    let mut ctr = Counter(iv.0); // This is OK because we're only encrypting one block.
-    key.ctr32_encrypt_within(&mut block, 0.., &mut ctr);
-    block
+    let iv_block = IvBlock::from_iv(iv);
+    match key.ctr32_encrypt_within(&mut block, 0.., iv_block) {
+        Ok(()) => block,
+        Result::<_, InOutLenInconsistentWithIvBlockLenError>::Err(_) => unreachable!(),
+    }
 }
 
 #[cfg(test)]

diff --git a/src/aead/aes/fallback.rs b/src/aead/aes/fallback.rs
@@ -12,7 +12,10 @@
 // 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 super::{
+    Block, EncryptBlock, EncryptCtr32, InOutLenInconsistentWithIvBlockLenError, Iv, IvBlock,
+    KeyBytes, AES_KEY,
+};
 use crate::error;
 use core::ops::RangeFrom;
 
@@ -39,9 +42,20 @@ impl EncryptBlock for Key {
 }
 
 impl EncryptCtr32 for Key {
-    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+    fn ctr32_encrypt_within(
+        &self,
+        in_out: &mut [u8],
+        src: RangeFrom<usize>,
+        iv_block: IvBlock,
+    ) -> Result<(), InOutLenInconsistentWithIvBlockLenError> {
         unsafe {
-            ctr32_encrypt_blocks!(aes_nohw_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr)
+            ctr32_encrypt_blocks!(
+                aes_nohw_ctr32_encrypt_blocks,
+                in_out,
+                src,
+                &self.inner,
+                iv_block
+            )
         }
     }
 }
diff --git a/src/aead/aes/ffi.rs b/src/aead/aes/ffi.rs
@@ -12,9 +12,9 @@
 // OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 // CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
-use super::{Block, KeyBytes, BLOCK_LEN};
+use super::{Block, InOutLenInconsistentWithIvBlockLenError, IvBlock, KeyBytes, BLOCK_LEN};
 use crate::{bits::BitLength, c, error, polyfill::slice};
-use core::{num::NonZeroUsize, ops::RangeFrom};
+use core::ops::RangeFrom;
 
 /// nonce || big-endian counter.
 #[repr(transparent)]
@@ -127,9 +127,9 @@ 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 $(,)? ) => {{
+    ($name:ident, $in_out:expr, $src:expr, $key:expr, $iv_block:expr $(,)? ) => {{
         use crate::{
-            aead::aes::{ffi::AES_KEY, Counter, BLOCK_LEN},
+            aead::aes::{ffi::AES_KEY, BLOCK_LEN},
             c,
         };
         prefixed_extern! {
@@ -138,10 +138,10 @@ macro_rules! ctr32_encrypt_blocks {
                 output: *mut [u8; BLOCK_LEN],
                 blocks: c::NonZero_size_t,
                 key: &AES_KEY,
-                ivec: &Counter,
+                ivec: &[u8; BLOCK_LEN],
             );
         }
-        $key.ctr32_encrypt_blocks($name, $in_out, $src, $ctr)
+        $key.ctr32_encrypt_blocks($name, $in_out, $src, $iv_block)
     }};
 }
 
@@ -165,23 +165,22 @@ impl AES_KEY {
             output: *mut [u8; BLOCK_LEN],
             blocks: c::NonZero_size_t,
             key: &AES_KEY,
-            ivec: &Counter,
+            ivec: &[u8; BLOCK_LEN],
         ),
         in_out: &mut [u8],
         src: RangeFrom<usize>,
-        ctr: &mut Counter,
-    ) {
+        iv_block: IvBlock,
+    ) -> Result<(), InOutLenInconsistentWithIvBlockLenError> {
         let (input, leftover) = slice::as_chunks(&in_out[src]);
         debug_assert_eq!(leftover.len(), 0);
 
-        let blocks = match NonZeroUsize::new(input.len()) {
-            Some(blocks) => blocks,
-            None => {
-                return;
-            }
-        };
+        if input.len() != iv_block.len().get() {
+            return Err(InOutLenInconsistentWithIvBlockLenError::new());
+        }
+        debug_assert!(!input.is_empty());
 
-        let blocks_u32: u32 = blocks.get().try_into().unwrap();
+        let iv_block_len = iv_block.len();
+        let initial_iv = iv_block.into_initial_iv().into_block_less_safe();
 
         let input = input.as_ptr();
         let output: *mut [u8; BLOCK_LEN] = in_out.as_mut_ptr().cast();
@@ -196,9 +195,9 @@ impl AES_KEY {
         //  * The caller is responsible for ensuring `key` was initialized by the
         //    `set_encrypt_key!` invocation required by `f`.
         unsafe {
-            f(input, output, blocks, self, ctr);
+            f(input, output, iv_block_len, self, &initial_iv);
         }
 
-        ctr.increment_by_less_safe(blocks_u32);
+        Ok(())
     }
 }
diff --git a/src/aead/aes/hw.rs b/src/aead/aes/hw.rs
@@ -14,7 +14,10 @@
 
 #![cfg(any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64"))]
 
-use super::{Block, Counter, EncryptBlock, EncryptCtr32, Iv, KeyBytes, AES_KEY};
+use super::{
+    Block, EncryptBlock, EncryptCtr32, InOutLenInconsistentWithIvBlockLenError, Iv, IvBlock,
+    KeyBytes, AES_KEY,
+};
 use crate::{cpu, error};
 use core::ops::RangeFrom;
 
@@ -56,9 +59,22 @@ impl EncryptBlock for Key {
 }
 
 impl EncryptCtr32 for Key {
-    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+    fn ctr32_encrypt_within(
+        &self,
+        in_out: &mut [u8],
+        src: RangeFrom<usize>,
+        iv_block: IvBlock,
+    ) -> Result<(), InOutLenInconsistentWithIvBlockLenError> {
         #[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) }
+        unsafe {
+            ctr32_encrypt_blocks!(
+                aes_hw_ctr32_encrypt_blocks,
+                in_out,
+                src,
+                &self.inner,
+                iv_block
+            )
+        }
     }
 }
diff --git a/src/aead/aes/vp.rs b/src/aead/aes/vp.rs
@@ -19,7 +19,10 @@
     target_arch = "x86_64"
 ))]
 
-use super::{Block, Counter, EncryptBlock, EncryptCtr32, Iv, KeyBytes, AES_KEY};
+use super::{
+    Block, EncryptBlock, EncryptCtr32, InOutLenInconsistentWithIvBlockLenError, Iv, IvBlock,
+    KeyBytes, AES_KEY,
+};
 use crate::{cpu, error};
 use core::ops::RangeFrom;
 
@@ -57,14 +60,32 @@ impl EncryptBlock for Key {
 
 #[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) }
+    fn ctr32_encrypt_within(
+        &self,
+        in_out: &mut [u8],
+        src: RangeFrom<usize>,
+        iv_block: IvBlock,
+    ) -> Result<(), InOutLenInconsistentWithIvBlockLenError> {
+        unsafe {
+            ctr32_encrypt_blocks!(
+                vpaes_ctr32_encrypt_blocks,
+                in_out,
+                src,
+                &self.inner,
+                iv_block
+            )
+        }
     }
 }
 
 #[cfg(target_arch = "arm")]
 impl EncryptCtr32 for Key {
-    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+    fn ctr32_encrypt_within(
+        &self,
+        in_out: &mut [u8],
+        src: RangeFrom<usize>,
+        ctr: &mut Counter,
+    ) -> Result<(), CounterOverflowError> {
         use super::{bs, BLOCK_LEN};
 
         let in_out = {
@@ -122,9 +143,14 @@ impl EncryptBlock for Key {
 
 #[cfg(target_arch = "x86")]
 impl EncryptCtr32 for Key {
-    fn ctr32_encrypt_within(&self, in_out: &mut [u8], src: RangeFrom<usize>, ctr: &mut Counter) {
+    fn ctr32_encrypt_within(
+        &self,
+        in_out: &mut [u8],
+        src: RangeFrom<usize>,
+        ctr: &mut Counter,
+    ) -> Result<(), CounterOverflowError> {
         super::super::shift::shift_full_blocks(in_out, src, |input| {
-            self.encrypt_iv_xor_block(ctr.increment(), *input)
+            self.encrypt_iv_xor_block(ctr.increment()?, *input)
         });
     }
 }