blob: 35ce3b46f0ccf504d89c735879e283e353500a85 [file] [log] [blame]
//! I/O streams for wrapping `BufRead` types as encoders/decoders
use std::io::prelude::*;
use std::io;
use super::CompressParams;
use raw::{self, Decompress, DeStatus, Compress, CompressOp, CoStatus};
#[derive(Clone, Copy, Eq, PartialEq)]
enum DoneStatus {
Processing,
Finishing,
Done,
}
/// A brotli encoder, or compressor.
///
/// This structure implements a `BufRead` interface and will read uncompressed
/// data from an underlying stream and emit a stream of compressed data.
pub struct BrotliEncoder<R: BufRead> {
obj: R,
data: Compress,
done: DoneStatus,
err: Option<raw::Error>,
}
/// A brotli decoder, or decompressor.
///
/// This structure implements a `BufRead` interface and takes a stream of
/// compressed data as input, providing the decompressed data when read from.
pub struct BrotliDecoder<R: BufRead> {
obj: R,
data: Decompress,
err: Option<raw::Error>,
}
impl<R: BufRead> BrotliEncoder<R> {
/// Creates a new encoder which will read uncompressed data from the given
/// stream and emit the compressed stream.
///
/// The `level` argument here is typically 0-11.
pub fn new(r: R, level: u32) -> BrotliEncoder<R> {
let mut data = Compress::new();
data.set_params(CompressParams::new().quality(level));
BrotliEncoder {
obj: r,
data: data,
done: DoneStatus::Processing,
err: None,
}
}
/// Creates a new encoder with a custom `CompressParams`.
pub fn from_params(r: R, params: &CompressParams) -> BrotliEncoder<R> {
let mut data = Compress::new();
data.set_params(params);
BrotliEncoder {
obj: r,
data: data,
done: DoneStatus::Processing,
err: None,
}
}
/// Acquires a reference to the underlying stream
pub fn get_ref(&self) -> &R {
&self.obj
}
/// Acquires a mutable reference to the underlying stream
///
/// Note that mutation of the stream may result in surprising results if
/// this encoder is continued to be used.
pub fn get_mut(&mut self) -> &mut R {
&mut self.obj
}
/// Consumes this encoder, returning the underlying reader.
pub fn into_inner(self) -> R {
self.obj
}
}
impl<R: BufRead> Read for BrotliEncoder<R> {
fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
if buf.is_empty() { return Ok(0) }
// If the compressor has failed at some point, this is set.
// Unfortunately we have no idea what status is in the compressor
// was in when it failed so we can't do anything except bail again.
if let Some(ref err) = self.err {
return Err(err.clone().into())
}
if let Some(data) = self.data.take_output(Some(buf.len())) {
buf[..data.len()].copy_from_slice(data);
return Ok(data.len())
}
match self.done {
DoneStatus::Done => return Ok(0),
DoneStatus::Finishing => return tryfinish(self, buf),
DoneStatus::Processing => (),
}
loop {
let amt_in;
let amt_out;
{
let input = &mut try!(self.obj.fill_buf());
let avail_in = input.len();
if avail_in == 0 {
break
}
let output = &mut buf;
let avail_out = output.len();
if let Err(err) = self.data.compress(CompressOp::Process, input, output) {
self.err = Some(err.clone().into());
return Err(err.into())
}
amt_in = avail_in - input.len();
amt_out = avail_out - output.len();
}
self.obj.consume(amt_in);
if amt_out == 0 {
assert!(amt_in != 0);
continue
}
return Ok(amt_out)
}
self.done = DoneStatus::Finishing;
return tryfinish(self, buf);
fn tryfinish<R: BufRead>(enc: &mut BrotliEncoder<R>, mut buf: &mut [u8])
-> io::Result<usize> {
let output = &mut buf;
let avail_out = output.len();
let iscomplete = match enc.data.compress(CompressOp::Finish, &mut &[][..], output) {
Ok(c) => c,
Err(err) => {
enc.err = err.clone().into();
return Err(err.into())
},
};
let written = avail_out - output.len();
assert!(written != 0 || iscomplete == CoStatus::Finished);
if iscomplete == CoStatus::Finished {
enc.done = DoneStatus::Done
}
Ok(written)
}
}
}
impl<R: BufRead> BrotliDecoder<R> {
/// Creates a new decoder which will decompress data read from the given
/// stream.
pub fn new(r: R) -> BrotliDecoder<R> {
BrotliDecoder {
data: Decompress::new(),
obj: r,
err: None,
}
}
/// Acquires a reference to the underlying stream
pub fn get_ref(&self) -> &R {
&self.obj
}
/// Acquires a mutable reference to the underlying stream
///
/// Note that mutation of the stream may result in surprising results if
/// this encoder is continued to be used.
pub fn get_mut(&mut self) -> &mut R {
&mut self.obj
}
/// Consumes this decoder, returning the underlying reader.
pub fn into_inner(self) -> R {
self.obj
}
}
impl<R: BufRead> Read for BrotliDecoder<R> {
fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
if buf.is_empty() { return Ok(0) }
// If the decompressor has failed at some point, this is set.
// Unfortunately we have no idea what status is in the compressor
// was in when it failed so we can't do anything except bail again.
if let Some(ref err) = self.err {
return Err(err.clone().into())
}
loop {
let amt_in;
let amt_out;
let status;
{
let mut input = try!(self.obj.fill_buf());
let avail_in = input.len();
let avail_out = buf.len();
status = match self.data.decompress(&mut input, &mut buf) {
Ok(s) => s,
Err(err) => {
self.err = Some(err.clone().into());
return Err(err.into())
},
};
amt_in = avail_in - input.len();
amt_out = avail_out - buf.len()
}
self.obj.consume(amt_in);
if amt_in == 0 && status == DeStatus::NeedInput {
return Err(io::Error::new(io::ErrorKind::Other,
"corrupted brotli stream"))
}
if amt_out == 0 && status != DeStatus::Finished {
assert!(amt_in != 0);
continue
}
return Ok(amt_out)
}
}
}