diff --git a/src/aead/aes.rs b/src/aead/aes.rs index a1bf4094b..2db1739f0 100644 --- a/src/aead/aes.rs +++ b/src/aead/aes.rs @@ -17,8 +17,10 @@ use crate::{ constant_time, cpu::{self, GetFeature as _}, error, + polyfill::unwrap_const, }; use cfg_if::cfg_if; +use core::num::NonZeroU32; pub(super) use ffi::Counter; @@ -123,15 +125,17 @@ impl Counter { } pub fn increment(&mut self) -> Iv { + const ONE: NonZeroU32 = unwrap_const(NonZeroU32::new(1)); + let iv = Iv(self.0); - self.increment_by_less_safe(1); + self.increment_by_less_safe(ONE); iv } - fn increment_by_less_safe(&mut self, increment_by: u32) { + fn increment_by_less_safe(&mut self, increment_by: NonZeroU32) { 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; + let new_value = old_value + increment_by.get(); [*c0, *c1, *c2, *c3] = u32::to_be_bytes(new_value); } } diff --git a/src/aead/aes/ffi.rs b/src/aead/aes/ffi.rs index 97c931f5c..1826124eb 100644 --- a/src/aead/aes/ffi.rs +++ b/src/aead/aes/ffi.rs @@ -14,7 +14,7 @@ use super::{Block, InOut, KeyBytes, BLOCK_LEN}; use crate::{bits::BitLength, c, error}; -use core::num::NonZeroUsize; +use core::num::{NonZeroU32, NonZeroUsize}; /// nonce || big-endian counter. #[repr(transparent)] @@ -182,15 +182,14 @@ impl AES_KEY { let input: *const [u8; BLOCK_LEN] = input.cast(); let output: *mut [u8; BLOCK_LEN] = output.cast(); - let blocks_u32: u32 = blocks.get().try_into().unwrap(); + let blocks_u32: NonZeroU32 = blocks.try_into().unwrap(); // SAFETY: // * `input` points to `blocks` blocks. // * `output` points to space for `blocks` blocks to be written. // * input == output.add(n), where n == src.start, and the caller is // responsible for ensuing this sufficient for `f` to work correctly. - // * The caller is responsible for ensuring `f` can handle any value of - // `blocks` including zero. + // * `blocks` is non-zero so `f` doesn't have to work for empty slices. // * The caller is responsible for ensuring `key` was initialized by the // `set_encrypt_key!` invocation required by `f`. unsafe {