diff --git a/chacha20/src/chacha.rs b/chacha20/src/chacha.rs index 40ea3f6f..6ac501da 100644 --- a/chacha20/src/chacha.rs +++ b/chacha20/src/chacha.rs @@ -162,8 +162,18 @@ impl StreamCipherSeek for ChaCha { fn try_seek(&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 StreamCipherSeek for ChaCha { impl ChaCha { /// 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(()) } } diff --git a/chacha20/tests/lib.rs b/chacha20/tests/lib.rs index 6e5cf3d3..0dd2328c 100644 --- a/chacha20/tests/lib.rs +++ b/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::().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::().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::().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::().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 {