use flate2::{Decompress, FlushDecompress, Status}; use std::{io, io::BufRead}; /// Provide streaming decompression using the `std::io::Read` trait. /// If `std::io::BufReader` is used, an allocation for the input buffer will be performed. pub struct InflateReader { pub(crate) inner: R, pub(crate) decompressor: Decompress, } impl InflateReader where R: io::BufRead, { pub fn from_read(read: R) -> InflateReader { // TODO: Performance opportunity - a buf reader that doesn't allocate InflateReader { decompressor: Decompress::new(true), inner: read, } } pub fn reset(&mut self, read: R) { self.inner = read; self.decompressor.reset(true); } } impl io::Read for InflateReader where R: BufRead, { fn read(&mut self, into: &mut [u8]) -> io::Result { read(&mut self.inner, &mut self.decompressor, into) } } /// The boxed variant is faster for what we do (moving the decompressor in and out a lot) pub struct InflateReaderBoxed { pub(crate) inner: R, pub(crate) decompressor: Box, } impl io::Read for InflateReaderBoxed where R: BufRead, { fn read(&mut self, into: &mut [u8]) -> io::Result { read(&mut self.inner, &mut self.decompressor, into) } } fn read(obj: &mut R, data: &mut Decompress, dst: &mut [u8]) -> io::Result where R: BufRead, { loop { let (read, consumed, ret, eof); { let input = obj.fill_buf()?; eof = input.is_empty(); let before_out = data.total_out(); let before_in = data.total_in(); let flush = if eof { FlushDecompress::Finish } else { FlushDecompress::None }; ret = data.decompress(input, dst, flush); read = (data.total_out() - before_out) as usize; consumed = (data.total_in() - before_in) as usize; } obj.consume(consumed); match ret { // If we haven't ready any data and we haven't hit EOF yet, // then we need to keep asking for more data because if we // return that 0 bytes of data have been read then it will // be interpreted as EOF. Ok(Status::Ok) | Ok(Status::BufError) if read == 0 && !eof && !dst.is_empty() => continue, Ok(Status::Ok) | Ok(Status::BufError) | Ok(Status::StreamEnd) => return Ok(read), Err(..) => return Err(io::Error::new(io::ErrorKind::InvalidInput, "corrupt deflate stream")), } } } #[cfg(test)] mod tests { use super::*; use git_object::bstr::ByteSlice; use std::{io::Read, path::PathBuf}; fn fixture_path(path: &str) -> PathBuf { PathBuf::from("tests/fixtures").join(path) } #[test] fn small_file_decompress() -> Result<(), Box> { let r = InflateReader::from_read(io::BufReader::new(std::fs::File::open(fixture_path( "objects/37/d4e6c5c48ba0d245164c4e10d5f41140cab980", ))?)); let mut bytes = r.bytes(); let content = bytes.by_ref().take(16).collect::, _>>()?; assert_eq!(content.as_slice().as_bstr(), b"blob 9\0hi there\n".as_bstr()); assert!(bytes.next().is_none()); Ok(()) } }