ChaCha20: Fix stream-length overflow check

chacha20/src/chacha.rs

@@ -162,8 +162,18 @@ impl<R: Rounds> StreamCipherSeek for ChaCha<R> {
 
     fn try_seek<T: SeekNum>(&mut self, pos: T) -> Result<(), LoopError> {
         let res = pos.to_block_byte(BLOCK_SIZE as u8)?;
+        let old_counter = self.counter;
+        let old_buffer_pos = self.buffer_pos;
+
         self.counter = res.0;
         self.buffer_pos = res.1;
+
+        if let Err(e) = self.check_data_len(&[0]) {
+            self.counter = old_counter;
+            self.buffer_pos = old_buffer_pos;
+            return Err(e);
+        }
+
         if self.buffer_pos != 0 {
             self.generate_block(self.counter);
         }
@@ -174,16 +184,18 @@ impl<R: Rounds> StreamCipherSeek for ChaCha<R> {
 impl<R: Rounds> ChaCha<R> {
     /// Check data length
     fn check_data_len(&self, data: &[u8]) -> Result<(), LoopError> {
-        let leftover_bytes = BUFFER_SIZE - self.buffer_pos as usize;
-        if data.len() < leftover_bytes {
-            return Ok(());
-        }
-        let blocks = 1 + (data.len() - leftover_bytes) / BLOCK_SIZE;
-        let res = self.counter.checked_add(blocks as u64).ok_or(LoopError)?;
-        if res <= MAX_BLOCKS as u64 {
-            Ok(())
-        } else {
+        let byte_after_last = self
+            .counter
+            .checked_mul(BLOCK_SIZE as u64)
+            .ok_or(LoopError)?
+            .checked_add(self.buffer_pos as u64)
+            .ok_or(LoopError)?
+            .checked_add(data.len() as u64)
+            .ok_or(LoopError)?;
+        if byte_after_last > ((MAX_BLOCKS as u64) + 1) * (BLOCK_SIZE as u64) {
             Err(LoopError)
+        } else {
+            Ok(())
         }
     }
 

chacha20/tests/lib.rs

@@ -6,6 +6,111 @@ use chacha20::ChaCha20;
 cipher::stream_cipher_test!(chacha20_core, ChaCha20, "chacha20");
 cipher::stream_cipher_seek_test!(chacha20_seek, ChaCha20);
 
+mod overflow {
+    use cipher::{NewCipher, StreamCipher, StreamCipherSeek};
+
+    const OFFSET_256GB: u64 = 256u64 << 30;
+
+    #[test]
+    fn bad_overflow_check1() {
+        let mut cipher = chacha20::ChaCha20::new(&Default::default(), &Default::default());
+        cipher
+            .try_seek(OFFSET_256GB - 1)
+            .expect("Couldn't seek to nearly 256GB");
+        let mut data = [0u8; 1];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect("Couldn't encrypt the last byte of 256GB");
+        assert_eq!(cipher.try_current_pos::<u64>().unwrap(), OFFSET_256GB);
+        let mut data = [0u8; 1];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect_err("Could encrypt past the last byte of 256GB");
+    }
+
+    #[test]
+    fn bad_overflow_check2() {
+        let mut cipher = chacha20::ChaCha20::new(&Default::default(), &Default::default());
+        cipher
+            .try_seek(OFFSET_256GB - 1)
+            .expect("Couldn't seek to nearly 256GB");
+        let mut data = [0u8; 2];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect_err("Could encrypt over the 256GB boundary");
+    }
+
+    #[test]
+    fn bad_overflow_check3() {
+        let mut cipher = chacha20::ChaCha20::new(&Default::default(), &Default::default());
+        cipher
+            .try_seek(OFFSET_256GB - 1)
+            .expect("Couldn't seek to nearly 256GB");
+        let mut data = [0u8; 1];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect("Couldn't encrypt the last byte of 256GB");
+        assert_eq!(cipher.try_current_pos::<u64>().unwrap(), OFFSET_256GB);
+        let mut data = [0u8; 63];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect_err("Could encrypt past the last byte of 256GB");
+    }
+
+    #[test]
+    fn bad_overflow_check4() {
+        let mut cipher = chacha20::ChaCha20::new(&Default::default(), &Default::default());
+        cipher
+            .try_seek(OFFSET_256GB - 1)
+            .expect("Couldn't seek to nearly 256GB");
+        let mut data = [0u8; 1];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect("Couldn't encrypt the last byte of 256GB");
+        assert_eq!(cipher.try_current_pos::<u64>().unwrap(), OFFSET_256GB);
+        let mut data = [0u8; 64];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect_err("Could encrypt past the last byte of 256GB");
+    }
+
+    #[test]
+    fn bad_overflow_check5() {
+        let mut cipher = chacha20::ChaCha20::new(&Default::default(), &Default::default());
+        cipher
+            .try_seek(OFFSET_256GB - 1)
+            .expect("Couldn't seek to nearly 256GB");
+        let mut data = [0u8; 1];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect("Couldn't encrypt the last byte of 256GB");
+        assert_eq!(cipher.try_current_pos::<u64>().unwrap(), OFFSET_256GB);
+        let mut data = [0u8; 65];
+        cipher
+            .try_apply_keystream(&mut data)
+            .expect_err("Could encrypt past the last byte of 256GB");
+    }
+
+    #[test]
+    fn bad_overflow_check6() {
+        let mut cipher = chacha20::ChaCha20::new(&Default::default(), &Default::default());
+        cipher
+            .try_seek(OFFSET_256GB)
+            .expect_err("Could seek to 256GB");
+    }
+
+    #[test]
+    fn bad_overflow_check7() {
+        let mut cipher = chacha20::ChaCha20::new(&Default::default(), &Default::default());
+        if let Ok(()) = cipher.try_seek(OFFSET_256GB + 63) {
+            let mut data = [0u8; 1];
+            cipher
+                .try_apply_keystream(&mut data)
+                .expect_err("Could encrypt the 64th byte past the 256GB boundary");
+        }
+    }
+}
+
 #[cfg(feature = "xchacha")]
 #[rustfmt::skip]
 mod xchacha20 {