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

View file

@ -153,6 +153,24 @@ impl<'a> Read for CryptoReader<'a> {
CryptoReader::Aes { reader: r, .. } => r.read(buf), 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> { impl<'a> CryptoReader<'a> {
@ -219,6 +237,66 @@ impl<'a> Read for ZipFileReader<'a> {
ZipFileReader::Xz(r) => r.read(buf), 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> { impl<'a> ZipFileReader<'a> {
@ -1624,6 +1702,18 @@ impl<'a> Read for ZipFile<'a> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.get_reader()?.read(buf) 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> { impl<'a> Drop for ZipFile<'a> {