aes: Clarify counter overflow checking.

Move `Conuter` and `Iv` to aes/counter.rs.

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.

`git difftool HEAD^1:src/aead/aes.rs src/aead/aes/counter.rs`

src/aead/aes.rs

@@ -21,7 +21,10 @@ use crate::{
 use cfg_if::cfg_if;
 use core::ops::RangeFrom;
 
-pub(super) use self::{counter::Iv, ffi::Counter};
+pub(super) use self::{
+    counter::{CounterOverflowError, Iv, IvBlock},
+    ffi::Counter,
+};
 
 #[macro_use]
 mod ffi;
@@ -114,6 +117,14 @@ pub enum KeyBytes<'a> {
     AES_256(&'a [u8; AES_256_KEY_LEN]),
 }
 
+pub(super) struct InOutLenInconsistentWithIvBlockLenError(());
+impl InOutLenInconsistentWithIvBlockLenError {
+    #[cold]
+    fn new() -> Self {
+        Self(())
+    }
+}
+
 pub(super) type Block = [u8; BLOCK_LEN];
 pub(super) const BLOCK_LEN: usize = 16;
 pub(super) const ZERO_BLOCK: Block = [0u8; BLOCK_LEN];
@@ -124,7 +135,12 @@ 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)]
@@ -144,11 +160,11 @@ fn encrypt_iv_xor_block_using_encrypt_block(
 
 #[allow(dead_code)]
 fn encrypt_iv_xor_block_using_ctr32(key: &impl EncryptCtr32, iv: Iv, mut block: Block) -> Block {
-    // This is OK because we're only encrypting one block, and `iv` is already
-    // reserved for us to use.
-    let mut ctr = Counter(iv.into_block_less_safe());
-    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)]

src/aead/aes/bs.rs

@@ -14,7 +14,7 @@
 
 #![cfg(target_arch = "arm")]
 
-use super::{Counter, AES_KEY};
+use super::{IvBlock, AES_KEY};
 use core::ops::RangeFrom;
 
 /// SAFETY:
@@ -31,8 +31,8 @@ pub(super) unsafe fn ctr32_encrypt_blocks_with_vpaes_key(
     in_out: &mut [u8],
     src: RangeFrom<usize>,
     vpaes_key: &AES_KEY,
-    ctr: &mut Counter,
-) {
+    iv_block: IvBlock,
+) -> Result<(), super::InOutLenInconsistentWithIvBlockLenError> {
     prefixed_extern! {
         // bsaes_ctr32_encrypt_blocks requires transformation of an existing
         // VPAES key; there is no `bsaes_set_encrypt_key`.
@@ -57,6 +57,12 @@ pub(super) unsafe fn ctr32_encrypt_blocks_with_vpaes_key(
     //  * `bsaes_ctr32_encrypt_blocks` satisfies the contract for
     //    `ctr32_encrypt_blocks`.
     unsafe {
-        ctr32_encrypt_blocks!(bsaes_ctr32_encrypt_blocks, in_out, src, &bsaes_key, ctr);
+        ctr32_encrypt_blocks!(
+            bsaes_ctr32_encrypt_blocks,
+            in_out,
+            src,
+            &bsaes_key,
+            iv_block
+        )
     }
 }

src/aead/aes/counter.rs

@@ -17,48 +17,179 @@ use super::{
     ffi::Counter,
     Block, BLOCK_LEN,
 };
+use crate::polyfill::{nonzerousize_from_nonzerou32, unwrap_const};
+use core::num::{NonZeroU32, NonZeroUsize};
 
 // `Counter` is `ffi::Counter` as its representation is dictated by its use in
 // the FFI.
 impl Counter {
-    pub fn one(nonce: Nonce) -> Self {
+    pub fn one_two(nonce: Nonce) -> (Iv, Self) {
         let mut value = [0u8; BLOCK_LEN];
         value[..NONCE_LEN].copy_from_slice(nonce.as_ref());
         value[BLOCK_LEN - 1] = 1;
-        Self(value)
+        let iv = Iv::new_less_safe(value);
+        value[BLOCK_LEN - 1] = 2;
+        (iv, Self(value))
     }
 
-    pub fn increment(&mut self) -> Iv {
+    pub fn try_into_iv(self) -> Result<Iv, CounterOverflowError> {
         let iv = Iv(self.0);
-        self.increment_by_less_safe(1);
-        iv
+        let [.., c0, c1, c2, c3] = &self.0;
+        let old_value: u32 = u32::from_be_bytes([*c0, *c1, *c2, *c3]);
+        if old_value == 0 {
+            return Err(CounterOverflowError::new());
+        }
+        Ok(iv)
+    }
+
+    pub fn increment_by(
+        &mut self,
+        increment_by: NonZeroUsize,
+    ) -> Result<IvBlock, CounterOverflowError> {
+        #[cold]
+        #[inline(never)]
+        fn overflowed(sum: u32) -> Result<u32, CounterOverflowError> {
+            match sum {
+                0 => Ok(0),
+                _ => Err(CounterOverflowError::new()),
+            }
+        }
+
+        let iv = Iv(self.0);
+
+        let increment_by = match NonZeroU32::try_from(increment_by) {
+            Ok(value) => value,
+            _ => return Err(CounterOverflowError::new()),
+        };
+
+        let [.., c0, c1, c2, c3] = &mut self.0;
+        let old_value: u32 = u32::from_be_bytes([*c0, *c1, *c2, *c3]);
+        if old_value == 0 {
+            return Err(CounterOverflowError::new());
+        }
+        let new_value = match old_value.overflowing_add(increment_by.get()) {
+            (sum, false) => sum,
+            (sum, true) => overflowed(sum)?,
+        };
+        [*c0, *c1, *c2, *c3] = u32::to_be_bytes(new_value);
+
+        Ok(IvBlock {
+            initial_iv: iv,
+            len: increment_by,
+        })
     }
 
-    pub(super) fn increment_by_less_safe(&mut self, increment_by: u32) {
+    #[cfg(target_arch = "x86")]
+    pub(super) fn increment_unchecked_less_safe(&mut self) -> Iv {
+        let iv = Iv(self.0);
+
         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;
+        debug_assert_ne!(old_value, 0);
+        // TODO: unchecked_add?
+        let new_value = old_value.wrapping_add(1);
+        // Note that it *is* valid for new_value to be zero!
         [*c0, *c1, *c2, *c3] = u32::to_be_bytes(new_value);
+
+        iv
+    }
+}
+
+pub(in super::super) struct CounterOverflowError(());
+
+impl CounterOverflowError {
+    #[cold]
+    fn new() -> Self {
+        Self(())
+    }
+}
+
+pub(in super::super) struct IvBlock {
+    initial_iv: Iv,
+    // invariant: 0 < len && len <= u32::MAX
+    len: NonZeroU32,
+}
+
+impl IvBlock {
+    pub(super) fn from_iv(iv: Iv) -> Self {
+        const _1: NonZeroU32 = unwrap_const(NonZeroU32::new(1));
+        Self {
+            initial_iv: iv,
+            len: _1,
+        }
+    }
+
+    // This conversion cannot fail.
+    pub fn len(&self) -> NonZeroUsize {
+        nonzerousize_from_nonzerou32(self.len)
+    }
+
+    // "Less safe" because this subverts the IV reuse prevention machinery. The
+    // caller must ensure the IV is used only once.
+    pub(super) fn into_initial_iv(self) -> Iv {
+        self.initial_iv
+    }
+
+    #[cfg(target_arch = "arm")]
+    pub(super) fn split_at(self, num_blocks: usize) -> (Option<IvBlock>, Option<IvBlock>) {
+        let num_before = match u32::try_from(num_blocks).ok().and_then(NonZeroU32::new) {
+            Some(num_blocks) => num_blocks,
+            None => return (None, Some(self)),
+        };
+        let num_after = match self
+            .len
+            .get()
+            .checked_sub(num_before.get())
+            .and_then(NonZeroU32::new)
+        {
+            Some(num_after) => num_after,
+            None => return (Some(self), None),
+        };
+        let mut ctr = Counter(self.initial_iv.0);
+        match ctr.increment_by(nonzerousize_from_nonzerou32(num_before)) {
+            Ok(before) => {
+                let after = Self {
+                    initial_iv: Iv::new_less_safe(ctr.0),
+                    len: num_after,
+                };
+                (Some(before), Some(after))
+            }
+            Result::<_, CounterOverflowError>::Err(_) => {
+                unreachable!()
+            }
+        }
+    }
+
+    #[cfg(target_arch = "x86")]
+    pub(super) fn into_counter_less_safe(
+        self,
+        input_blocks: usize,
+    ) -> Result<Counter, super::InOutLenInconsistentWithIvBlockLenError> {
+        if input_blocks != self.len().get() {
+            return Err(super::InOutLenInconsistentWithIvBlockLenError::new());
+        }
+        Ok(Counter(self.initial_iv.0))
     }
 }
 
 /// The IV for a single block encryption.
 ///
 /// Intentionally not `Clone` to ensure each is used only once.
-pub struct Iv(Block);
+pub(in super::super) struct Iv(Block);
 
 impl Iv {
+    // This is "less safe" because it subverts the counter reuse protection.
+    // The caller needs to ensure that the IV isn't reused.
     pub(super) fn new_less_safe(value: Block) -> Self {
         Self(value)
     }
 
+    /// "Less safe" because it defeats attempts to use the type system to prevent reuse of the IV.
+    #[inline]
     pub(super) fn into_block_less_safe(self) -> Block {
         self.0
     }
 }
 
-impl From<Counter> for Iv {
-    fn from(counter: Counter) -> Self {
-        Self(counter.0)
-    }
-}
+#[cfg(test)]
+mod tests {}

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
+            )
         }
     }
 }