brotli2 0.3.0

//! 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 mut 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)
        }
    }
}