//! Helper module to compute a CRC32 checksum use std::io; use std::io::prelude::*; use crc32fast::Hasher; /// Reader that validates the CRC32 when it reaches the EOF. pub struct Crc32Reader { inner: R, hasher: Hasher, check: u32, /// Signals if `inner` stores aes encrypted data. /// AE-2 encrypted data doesn't use crc and sets the value to 0. enabled: bool, } impl Crc32Reader { /// Get a new Crc32Reader which checks the inner reader against checksum. /// The check is disabled if `ae2_encrypted == true`. pub(crate) fn new(inner: R, checksum: u32, ae2_encrypted: bool) -> Crc32Reader { Crc32Reader { inner, hasher: Hasher::new(), check: checksum, enabled: !ae2_encrypted, } } fn check_matches(&self) -> bool { self.check == self.hasher.clone().finalize() } pub fn into_inner(self) -> R { self.inner } } #[cold] fn invalid_checksum() -> io::Error { io::Error::new(io::ErrorKind::InvalidData, "Invalid checksum") } impl Read for Crc32Reader { fn read(&mut self, buf: &mut [u8]) -> io::Result { let count = self.inner.read(buf)?; if self.enabled { if count == 0 && !buf.is_empty() && !self.check_matches() { return Err(invalid_checksum()); } self.hasher.update(&buf[..count]); } Ok(count) } fn read_to_end(&mut self, buf: &mut Vec) -> io::Result { let start = buf.len(); let n = self.inner.read_to_end(buf)?; if self.enabled { self.hasher.update(&buf[start..]); if !self.check_matches() { return Err(invalid_checksum()); } } Ok(n) } fn read_to_string(&mut self, buf: &mut String) -> io::Result { let start = buf.len(); let n = self.inner.read_to_string(buf)?; if self.enabled { self.hasher.update(&buf.as_bytes()[start..]); if !self.check_matches() { return Err(invalid_checksum()); } } Ok(n) } } #[cfg(test)] mod test { use super::*; #[test] fn test_empty_reader() { let data: &[u8] = b""; let mut buf = [0; 1]; let mut reader = Crc32Reader::new(data, 0, false); assert_eq!(reader.read(&mut buf).unwrap(), 0); let mut reader = Crc32Reader::new(data, 1, false); assert!(reader .read(&mut buf) .unwrap_err() .to_string() .contains("Invalid checksum")); } #[test] fn test_byte_by_byte() { let data: &[u8] = b"1234"; let mut buf = [0; 1]; let mut reader = Crc32Reader::new(data, 0x9be3e0a3, false); assert_eq!(reader.read(&mut buf).unwrap(), 1); assert_eq!(reader.read(&mut buf).unwrap(), 1); assert_eq!(reader.read(&mut buf).unwrap(), 1); assert_eq!(reader.read(&mut buf).unwrap(), 1); assert_eq!(reader.read(&mut buf).unwrap(), 0); // Can keep reading 0 bytes after the end assert_eq!(reader.read(&mut buf).unwrap(), 0); } #[test] fn test_zero_read() { let data: &[u8] = b"1234"; let mut buf = [0; 5]; let mut reader = Crc32Reader::new(data, 0x9be3e0a3, false); assert_eq!(reader.read(&mut buf[..0]).unwrap(), 0); assert_eq!(reader.read(&mut buf).unwrap(), 4); } }