cipher: fix seeking implementation in the stream cipher wrapper

cipher/src/stream.rs

@@ -249,8 +249,8 @@ macro_rules! impl_seek_num {
                 }
 
                 fn into_block_byte<T: StreamCipherCounter>(self, block_size: u8) -> Result<(T, u8), OverflowError> {
-                    let bs: Self = block_size.into();
-                    let byte = (self % bs) as u8;
+                    let bs = Self::from(block_size);
+                    let byte = u8::try_from(self % bs).expect("bs fits into u8");
                     let block = T::try_from(self / bs).map_err(|_| OverflowError)?;
                     Ok((block, byte))
                 }

cipher/src/stream/core_api.rs

@@ -37,11 +37,10 @@ pub trait StreamCipherClosure: BlockSizeUser {
 
 /// Block-level synchronous stream ciphers.
 pub trait StreamCipherCore: BlockSizeUser + Sized {
-    /// Return number of remaining blocks before cipher wraps around.
+    /// Return number of remaining blocks before the cipher wraps around.
     ///
     /// Returns `None` if number of remaining blocks can not be computed
-    /// (e.g. in ciphers based on the sponge construction) or it's too big
-    /// to fit into `usize`.
+    /// (e.g. in the case of sponge-based stream ciphers) or it’s too big to fit into `usize`.
     fn remaining_blocks(&self) -> Option<usize>;
 
     /// Process data using backend provided to the rank-2 closure.
@@ -91,23 +90,23 @@ pub trait StreamCipherCore: BlockSizeUser + Sized {
 
     /// Try to apply keystream to data not divided into blocks.
     ///
-    /// Consumes cipher since it may consume final keystream block only
-    /// partially.
+    /// Consumes cipher since it may consume the final keystream block only partially.
     ///
-    /// Returns an error if number of remaining blocks is not sufficient
-    /// for processing the input data.
+    /// Returns an error if the number of remaining blocks is not sufficient
+    /// for processing of the input data.
     #[inline]
     fn try_apply_keystream_partial(
         mut self,
         mut buf: InOutBuf<'_, '_, u8>,
     ) -> Result<(), StreamCipherError> {
-        if let Some(rem) = self.remaining_blocks() {
-            let blocks = if buf.len() % Self::BlockSize::USIZE == 0 {
-                buf.len() % Self::BlockSize::USIZE
-            } else {
-                buf.len() % Self::BlockSize::USIZE + 1
-            };
-            if blocks > rem {
+        if let Some(rem_blocks) = self.remaining_blocks() {
+            // Note that if `rem_blocks` is equal to zero, it means that
+            // the next generated block will be the last in the keystream and
+            // the cipher core will wrap to its initial state.
+            // Since we consume `self`, it's fine to generate the last keystream block,
+            // so we can use division instead of `div_ceil` to compute `req_blocks`.
+            let req_blocks = buf.len() / Self::BlockSize::USIZE;
+            if req_blocks > rem_blocks {
                 return Err(StreamCipherError);
             }
         }
@@ -164,7 +163,10 @@ pub trait StreamCipherCounter:
     + TryInto<u64>
     + TryInto<u128>
     + TryInto<usize>
+    + Copy
 {
+    /// Returns `true` if `self` is equal to the max counter value.
+    fn is_max(&self) -> bool;
 }
 
 /// Block-level seeking trait for stream ciphers.
@@ -181,7 +183,13 @@ pub trait StreamCipherSeekCore: StreamCipherCore {
 
 macro_rules! impl_counter {
     {$($t:ty )*} => {
-        $( impl StreamCipherCounter for $t { } )*
+        $(
+            impl StreamCipherCounter for $t {
+                fn is_max(&self) -> bool {
+                    *self == <$t>::MAX
+                }
+            }
+        )*
     };
 }
 

cipher/src/stream/wrapper.rs

@@ -1,3 +1,5 @@
+use crate::StreamCipherCounter;
+
 use super::{
     OverflowError, SeekNum, StreamCipher, StreamCipherCore, StreamCipherSeek, StreamCipherSeekCore,
     errors::StreamCipherError,
@@ -51,20 +53,14 @@ impl<T: StreamCipherCore> StreamCipherCoreWrapper<T> {
     }
 
     fn check_remaining(&self, data_len: usize) -> Result<(), StreamCipherError> {
-        let rem_blocks = match self.core.remaining_blocks() {
-            Some(v) => v,
-            None => return Ok(()),
+        let Some(rem_blocks) = self.core.remaining_blocks() else {
+            return Ok(());
         };
-
-        let buf_rem = self.buffer.remaining();
-        let data_len = match data_len.checked_sub(buf_rem) {
-            Some(0) | None => return Ok(()),
-            Some(res) => res,
+        let Some(data_len) = data_len.checked_sub(self.buffer.remaining()) else {
+            return Ok(());
         };
-
-        let bs = T::BlockSize::USIZE;
-        let blocks = data_len.div_ceil(bs);
-        if blocks > rem_blocks {
+        let req_blocks = data_len.div_ceil(T::BlockSize::USIZE);
+        if req_blocks > rem_blocks {
             Err(StreamCipherError)
         } else {
             Ok(())
@@ -129,8 +125,11 @@ impl<T: StreamCipherSeekCore> StreamCipherSeek for StreamCipherCoreWrapper<T> {
     }
 
     fn try_seek<SN: SeekNum>(&mut self, new_pos: SN) -> Result<(), StreamCipherError> {
-        let (block_pos, byte_pos) = new_pos.into_block_byte(T::BlockSize::U8)?;
-        // For correct implementations of `SeekNum` compiler should be able to
+        let (block_pos, byte_pos) = new_pos.into_block_byte::<T::Counter>(T::BlockSize::U8)?;
+        if byte_pos != 0 && block_pos.is_max() {
+            return Err(StreamCipherError);
+        }
+        // For correct implementations of `SeekNum` the compiler should be able to
         // eliminate this assert
         assert!(byte_pos < T::BlockSize::U8);
 

cipher/tests/stream.rs

@@ -101,6 +101,8 @@ fn dummy_stream_cipher_core() {
 
 #[cfg(feature = "stream-wrapper")]
 mod wrapper {
+    use core::panic;
+
     use super::*;
     use cipher::{StreamCipher, StreamCipherCoreWrapper, StreamCipherSeek};
 
@@ -127,26 +129,37 @@ mod wrapper {
     #[test]
     fn dummy_stream_cipher_seek_limit() {
         let mut cipher = DummyStreamCipher::new(&KEY.into(), &IV.into());
+        let mut buf = [0u8; 64];
+
+        let block_size = DummyStreamCipherCore::block_size();
+        let block_size_u128 = u128::try_from(block_size).unwrap();
+        let keystream_end = 1u128 << 68;
+        let last_block_pos = keystream_end - block_size_u128;
+
+        // Seeking to the last block or past it should return error
+        for offset in 0..block_size_u128 {
+            let res = cipher.try_seek(keystream_end - offset);
+            assert!(res.is_err());
+            let res = cipher.try_seek(keystream_end + offset);
+            assert!(res.is_err());
+        }
 
-        let pos = ((u64::MAX as u128) << 4) - 20;
-        cipher.try_seek(pos).unwrap();
-
-        let mut buf = [0u8; 30];
-        let res = cipher.try_apply_keystream(&mut buf);
-        assert!(res.is_err());
-        let cur_pos: u128 = cipher.current_pos();
-        assert_eq!(cur_pos, pos);
-
-        let res = cipher.try_apply_keystream(&mut buf[..19]);
-        assert!(res.is_ok());
-        let cur_pos: u128 = cipher.current_pos();
-        assert_eq!(cur_pos, pos + 19);
-
-        cipher.try_seek(pos).unwrap();
-
-        // TODO: fix as part of https://github.com/RustCrypto/traits/issues/1808
-        // let res = cipher.try_apply_keystream(&mut buf[..20]);
-        // assert!(res.is_err());
+        // Trying to apply the last keystream block should return error
+        for offset in block_size..buf.len() {
+            for len in 0..buf.len() {
+                let pos = keystream_end - u128::try_from(offset).unwrap();
+                let res = cipher.try_seek(pos);
+                assert!(res.is_ok());
+                let res = cipher.try_apply_keystream(&mut buf[..len]);
+                let expected_pos = pos + u128::try_from(len).unwrap();
+                if expected_pos > last_block_pos {
+                    assert!(res.is_err());
+                } else {
+                    assert!(res.is_ok());
+                    assert_eq!(cipher.current_pos::<u128>(), expected_pos);
+                }
+            }
+        }
     }
 
     #[cfg(feature = "dev")]