Improve several Read methods on ZipFile

This commit is contained in:
Benoît du Garreau 2024-07-11 14:31:31 +02:00
parent b61fd1198e
commit 83b1273fab
2 changed files with 131 additions and 10 deletions

View file

@ -12,7 +12,7 @@ pub struct Crc32Reader<R> {
check: u32,
/// Signals if `inner` stores aes encrypted data.
/// AE-2 encrypted data doesn't use crc and sets the value to 0.
ae2_encrypted: bool,
enabled: bool,
}
impl<R> Crc32Reader<R> {
@ -23,7 +23,7 @@ impl<R> Crc32Reader<R> {
inner,
hasher: Hasher::new(),
check: checksum,
ae2_encrypted,
enabled: !ae2_encrypted,
}
}
@ -36,20 +36,51 @@ impl<R> Crc32Reader<R> {
}
}
#[cold]
fn invalid_checksum() -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, "Invalid checksum")
}
impl<R: Read> Read for Crc32Reader<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let invalid_check = !buf.is_empty() && !self.check_matches() && !self.ae2_encrypted;
let count = self.inner.read(buf)?;
let count = match self.inner.read(buf) {
Ok(0) if invalid_check => {
return Err(io::Error::new(io::ErrorKind::Other, "Invalid checksum"))
if self.enabled {
if count == 0 && !buf.is_empty() && !self.check_matches() {
return Err(invalid_checksum());
}
Ok(n) => n,
Err(e) => return Err(e),
};
self.hasher.update(&buf[0..count]);
self.hasher.update(&buf[..count]);
}
Ok(count)
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
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<usize> {
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)]

View file

@ -153,6 +153,24 @@ impl<'a> Read for CryptoReader<'a> {
CryptoReader::Aes { reader: r, .. } => r.read(buf),
}
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
match self {
CryptoReader::Plaintext(r) => r.read_to_end(buf),
CryptoReader::ZipCrypto(r) => r.read_to_end(buf),
#[cfg(feature = "aes-crypto")]
CryptoReader::Aes { reader: r, .. } => r.read_to_end(buf),
}
}
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
match self {
CryptoReader::Plaintext(r) => r.read_to_string(buf),
CryptoReader::ZipCrypto(r) => r.read_to_string(buf),
#[cfg(feature = "aes-crypto")]
CryptoReader::Aes { reader: r, .. } => r.read_to_string(buf),
}
}
}
impl<'a> CryptoReader<'a> {
@ -219,6 +237,66 @@ impl<'a> Read for ZipFileReader<'a> {
ZipFileReader::Xz(r) => r.read(buf),
}
}
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
match self {
ZipFileReader::NoReader => panic!("ZipFileReader was in an invalid state"),
ZipFileReader::Raw(r) => r.read_exact(buf),
ZipFileReader::Stored(r) => r.read_exact(buf),
#[cfg(feature = "_deflate-any")]
ZipFileReader::Deflated(r) => r.read_exact(buf),
#[cfg(feature = "deflate64")]
ZipFileReader::Deflate64(r) => r.read_exact(buf),
#[cfg(feature = "bzip2")]
ZipFileReader::Bzip2(r) => r.read_exact(buf),
#[cfg(feature = "zstd")]
ZipFileReader::Zstd(r) => r.read_exact(buf),
#[cfg(feature = "lzma")]
ZipFileReader::Lzma(r) => r.read_exact(buf),
#[cfg(feature = "xz")]
ZipFileReader::Xz(r) => r.read_exact(buf),
}
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
match self {
ZipFileReader::NoReader => panic!("ZipFileReader was in an invalid state"),
ZipFileReader::Raw(r) => r.read_to_end(buf),
ZipFileReader::Stored(r) => r.read_to_end(buf),
#[cfg(feature = "_deflate-any")]
ZipFileReader::Deflated(r) => r.read_to_end(buf),
#[cfg(feature = "deflate64")]
ZipFileReader::Deflate64(r) => r.read_to_end(buf),
#[cfg(feature = "bzip2")]
ZipFileReader::Bzip2(r) => r.read_to_end(buf),
#[cfg(feature = "zstd")]
ZipFileReader::Zstd(r) => r.read_to_end(buf),
#[cfg(feature = "lzma")]
ZipFileReader::Lzma(r) => r.read_to_end(buf),
#[cfg(feature = "xz")]
ZipFileReader::Xz(r) => r.read_to_end(buf),
}
}
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
match self {
ZipFileReader::NoReader => panic!("ZipFileReader was in an invalid state"),
ZipFileReader::Raw(r) => r.read_to_string(buf),
ZipFileReader::Stored(r) => r.read_to_string(buf),
#[cfg(feature = "_deflate-any")]
ZipFileReader::Deflated(r) => r.read_to_string(buf),
#[cfg(feature = "deflate64")]
ZipFileReader::Deflate64(r) => r.read_to_string(buf),
#[cfg(feature = "bzip2")]
ZipFileReader::Bzip2(r) => r.read_to_string(buf),
#[cfg(feature = "zstd")]
ZipFileReader::Zstd(r) => r.read_to_string(buf),
#[cfg(feature = "lzma")]
ZipFileReader::Lzma(r) => r.read_to_string(buf),
#[cfg(feature = "xz")]
ZipFileReader::Xz(r) => r.read_to_string(buf),
}
}
}
impl<'a> ZipFileReader<'a> {
@ -1624,6 +1702,18 @@ impl<'a> Read for ZipFile<'a> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.get_reader()?.read(buf)
}
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
self.get_reader()?.read_exact(buf)
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
self.get_reader()?.read_to_end(buf)
}
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
self.get_reader()?.read_to_string(buf)
}
}
impl<'a> Drop for ZipFile<'a> {