libflate 2.3.2

use super::symbol;
use crate::bit;
use crate::lz77;
use no_std_io2::io::{self, Read};

/// The maximum number of decoded-but-unread bytes buffered internally before
/// the decoder yields output to the caller.
///
/// DEFLATE compressed blocks have no maximum expanded size, so without this
/// bound a single highly-compressible block could grow the internal buffer
/// without limit while decoding untrusted input. A single decoded symbol
/// produces at most 258 bytes, so the buffer can only overshoot this threshold
/// by at most 257 bytes before output is yielded.
const MAX_INTERNAL_BUFFER: usize = 64 * 1024;

/// DEFLATE decoder.
#[derive(Debug)]
pub struct Decoder<R> {
    bit_reader: bit::BitReader<R>,
    lz77_decoder: lz77::Lz77Decoder,
    /// Active Huffman decoder when the decoder is suspended in the middle of a
    /// compressed block (after `MAX_INTERNAL_BUFFER` was reached). `None` while
    /// between blocks.
    block_decoder: Option<symbol::Decoder>,
    eos: bool,
}
impl<R> Decoder<R>
where
    R: Read,
{
    /// Makes a new decoder instance.
    ///
    /// `inner` is to be decoded DEFLATE stream.
    ///
    /// # Examples
    /// ```
    /// # extern crate alloc;
    /// # use alloc::vec::Vec;
    /// use no_std_io2::io::{Cursor, Read};
    /// use libflate::deflate::Decoder;
    ///
    /// let encoded_data = [243, 72, 205, 201, 201, 87, 8, 207, 47, 202, 73, 81, 4, 0];
    /// let mut decoder = Decoder::new(&encoded_data[..]);
    /// let mut buf = Vec::new();
    /// decoder.read_to_end(&mut buf).unwrap();
    ///
    /// assert_eq!(buf, b"Hello World!");
    /// ```
    pub fn new(inner: R) -> Self {
        Decoder {
            bit_reader: bit::BitReader::new(inner),
            lz77_decoder: lz77::Lz77Decoder::new(),
            block_decoder: None,
            eos: false,
        }
    }

    /// Returns the immutable reference to the inner stream.
    pub fn as_inner_ref(&self) -> &R {
        self.bit_reader.as_inner_ref()
    }

    /// Returns the mutable reference to the inner stream.
    pub fn as_inner_mut(&mut self) -> &mut R {
        self.bit_reader.as_inner_mut()
    }

    /// Unwraps this `Decoder`, returning the underlying reader.
    ///
    /// # Examples
    /// ```
    /// use no_std_io2::io::Cursor;
    /// use libflate::deflate::Decoder;
    ///
    /// let encoded_data = [243, 72, 205, 201, 201, 87, 8, 207, 47, 202, 73, 81, 4, 0];
    /// let decoder = Decoder::new(Cursor::new(&encoded_data));
    /// assert_eq!(decoder.into_inner().into_inner(), &encoded_data);
    /// ```
    pub fn into_inner(self) -> R {
        self.bit_reader.into_inner()
    }

    /// Returns the data that has been decoded but has not yet been read.
    ///
    /// This method is useful to retrieve partial decoded data when the decoding process is failed.
    pub fn unread_decoded_data(&self) -> &[u8] {
        self.lz77_decoder.buffer()
    }

    pub(crate) fn reset(&mut self) {
        self.bit_reader.reset();
        self.lz77_decoder.clear();
        self.block_decoder = None;
        self.eos = false
    }

    fn read_non_compressed_block(&mut self) -> io::Result<()> {
        self.bit_reader.reset();
        let mut buf = [0; 2];
        self.bit_reader.as_inner_mut().read_exact(&mut buf)?;
        let len = u16::from_le_bytes(buf);
        self.bit_reader.as_inner_mut().read_exact(&mut buf)?;
        let nlen = u16::from_le_bytes(buf);
        if !len != nlen {
            Err(invalid_data_error!(
                "LEN={} is not the one's complement of NLEN={}",
                len,
                nlen
            ))
        } else {
            self.lz77_decoder
                .extend_from_reader(self.bit_reader.as_inner_mut().take(len.into()))
                .and_then(|used| {
                    if used != len.into() {
                        Err(io::Error::new(
                            io::ErrorKind::UnexpectedEof,
                            #[cfg(feature = "std")]
                            format!("The reader has incorrect length: expected {len}, read {used}"),
                            #[cfg(not(feature = "std"))]
                            "The reader has incorrect length",
                        ))
                    } else {
                        Ok(())
                    }
                })
        }
    }
    fn enter_compressed_block<H>(&mut self, huffman: &H) -> io::Result<()>
    where
        H: symbol::HuffmanCodec,
    {
        self.block_decoder = Some(huffman.load(&mut self.bit_reader)?);
        Ok(())
    }

    fn read_compressed_block(&mut self) -> io::Result<()> {
        debug_assert!(self.block_decoder.is_some());
        loop {
            let symbol_decoder = self
                .block_decoder
                .as_mut()
                .expect("block_decoder must be present when reading a block");
            let s = symbol_decoder.decode_unchecked(&mut self.bit_reader);
            self.bit_reader.check_last_error()?;
            match s {
                symbol::Symbol::Code(code) => {
                    self.lz77_decoder.decode(code)?;
                    if self.lz77_decoder.buffer().len() >= MAX_INTERNAL_BUFFER {
                        break;
                    }
                }
                symbol::Symbol::EndOfBlock => {
                    self.block_decoder = None;
                    break;
                }
            }
        }
        Ok(())
    }
}
impl<R> Read for Decoder<R>
where
    R: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        loop {
            if !self.lz77_decoder.buffer().is_empty() {
                return self.lz77_decoder.read(buf);
            }
            // Resume decoding a compressed block that was suspended after
            // `MAX_INTERNAL_BUFFER` was reached.
            if self.block_decoder.is_some() {
                self.read_compressed_block()?;
                continue;
            }
            if self.eos {
                return Ok(0);
            }
            let bfinal = self.bit_reader.read_bit()?;
            let btype = self.bit_reader.read_bits(2)?;
            self.eos = bfinal;
            match btype {
                0b00 => self.read_non_compressed_block()?,
                0b01 => {
                    self.enter_compressed_block(&symbol::FixedHuffmanCodec)?;
                    self.read_compressed_block()?;
                }
                0b10 => {
                    self.enter_compressed_block(&symbol::DynamicHuffmanCodec)?;
                    self.read_compressed_block()?;
                }
                0b11 => {
                    return Err(invalid_data_error!(
                        "btype 0x11 of DEFLATE is reserved(error) value"
                    ));
                }
                _ => unreachable!(),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "std")]
    use super::*;
    use crate::deflate::symbol::{DynamicHuffmanCodec, HuffmanCodec};
    #[cfg(feature = "std")]
    use std::io;

    #[test]
    fn test_issues_3() {
        // see: https://github.com/sile/libflate/issues/3
        let input = [
            180, 253, 73, 143, 28, 201, 150, 46, 8, 254, 150, 184, 139, 75, 18, 69, 247, 32, 157,
            51, 27, 141, 132, 207, 78, 210, 167, 116, 243, 160, 223, 136, 141, 66, 205, 76, 221,
            76, 195, 213, 84, 236, 234, 224, 78, 227, 34, 145, 221, 139, 126, 232, 69, 173, 170,
            208, 192, 219, 245, 67, 3, 15, 149, 120, 171, 70, 53, 106, 213, 175, 23, 21, 153, 139,
            254, 27, 249, 75, 234, 124, 71, 116, 56, 71, 68, 212, 204, 121, 115, 64, 222, 160, 203,
            119, 142, 170, 169, 138, 202, 112, 228, 140, 38,
        ];
        let mut bit_reader = crate::bit::BitReader::new(&input[..]);
        assert_eq!(bit_reader.read_bit().unwrap(), false); // not final block
        assert_eq!(bit_reader.read_bits(2).unwrap(), 0b10); // DynamicHuffmanCodec
        DynamicHuffmanCodec.load(&mut bit_reader).unwrap();
    }

    #[test]
    #[cfg(feature = "std")]
    fn it_works() {
        let input = [
            180, 253, 73, 143, 28, 201, 150, 46, 8, 254, 150, 184, 139, 75, 18, 69, 247, 32, 157,
            51, 27, 141, 132, 207, 78, 210, 167, 116, 243, 160, 223, 136, 141, 66, 205, 76, 221,
            76, 195, 213, 84, 236, 234, 224, 78, 227, 34, 145, 221, 139, 126, 232, 69, 173, 170,
            208, 192, 219, 245, 67, 3, 15, 149, 120, 171, 70, 53, 106, 213, 175, 23, 21, 153, 139,
            254, 27, 249, 75, 234, 124, 71, 116, 56, 71, 68, 212, 204, 121, 115, 64, 222, 160, 203,
            119, 142, 170, 169, 138, 202, 112, 228, 140, 38, 171, 162, 88, 212, 235, 56, 136, 231,
            233, 239, 113, 249, 163, 252, 16, 42, 138, 49, 226, 108, 73, 28, 153,
        ];
        let mut decoder = Decoder::new(&input[..]);

        let result = io::copy(&mut decoder, &mut io::sink());
        assert!(result.is_err());

        let error = result.err().unwrap();
        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
        assert!(error.to_string().starts_with("Too long backword reference"));
    }

    #[test]
    #[cfg(feature = "std")]
    fn test_issue_64() {
        let input = b"\x04\x04\x04\x05:\x1az*\xfc\x06\x01\x90\x01\x06\x01";
        let mut decoder = Decoder::new(&input[..]);
        assert!(io::copy(&mut decoder, &mut io::sink()).is_err());
    }

    /// The minimal valid WebAssembly module — used as the payload of the regression
    /// test below just so the decoded bytes are recognizable.
    #[cfg(feature = "std")]
    const WASM: [u8; 8] = [0x00, b'a', b's', b'm', 0x01, 0x00, 0x00, 0x00];

    // Regression test for https://github.com/sile/libflate/issues/88 :
    // decoding a stream carrying many DEFLATE blocks used to blow the stack
    // because `Read for Decoder` was implemented with self-recursive tail calls.
    #[test]
    #[cfg(feature = "std")]
    fn test_issue_88() {
        let gzip = make_large_deflate_stream(250_000);
        let mut decoder = crate::gzip::Decoder::new(&gzip[..]).unwrap();
        let mut decoded = Vec::new();
        decoder.read_to_end(&mut decoded).unwrap();
        assert_eq!(decoded, WASM);
    }

    /// Build a gzip stream that decompresses to `WASM` but is padded with
    /// `blocks - 1` empty non-final DEFLATE stored blocks in front of the
    /// final payload-carrying block. The empty blocks decompress to nothing,
    /// so the point of a large `blocks` count is stress: each empty block
    /// used to add one stack frame to `deflate::Decoder::read` and would
    /// eventually overflow the thread stack (see `test_issue_88`).
    #[cfg(feature = "std")]
    fn make_large_deflate_stream(blocks: usize) -> Vec<u8> {
        debug_assert!(
            blocks >= 1,
            "at least one block is required for the final payload"
        );
        /// Gzip header. CM=deflate, OS=unknown.
        const HEADER: [u8; 10] = [0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03];
        /// A non-final DEFLATE stored block of length zero: BFINAL=0, LEN=0, NLEN=0xffff.
        const EMPTY_NONFINAL_STORED_BLOCK: [u8; 5] = [0x00, 0x00, 0x00, 0xff, 0xff];
        /// Compute the IEEE CRC-32 (as used by gzip) of `data`.
        fn crc32(data: &[u8]) -> u32 {
            let mut crc: u32 = 0xffff_ffff;
            for &byte in data {
                crc ^= byte as u32;
                for _ in 0..8 {
                    let mask = (crc & 1).wrapping_neg();
                    crc = (crc >> 1) ^ (0xedb8_8320 & mask);
                }
            }
            !crc
        }

        let len = WASM.len() as u16;
        let mut payload = Vec::with_capacity(
            HEADER.len() + (blocks - 1) * EMPTY_NONFINAL_STORED_BLOCK.len() + 21,
        );
        payload.extend_from_slice(&HEADER);
        for _ in 0..(blocks - 1) {
            payload.extend_from_slice(&EMPTY_NONFINAL_STORED_BLOCK);
        }
        // Final stored block: BFINAL byte, then LEN and its ones-complement NLEN
        // then the raw stored bytes.
        payload.push(1);
        payload.extend_from_slice(&len.to_le_bytes());
        payload.extend_from_slice(&(!len).to_le_bytes());
        payload.extend_from_slice(&WASM);
        // gzip trailer: CRC32 of the uncompressed data, then ISIZE mod 2^32.
        payload.extend_from_slice(&crc32(&WASM).to_le_bytes());
        payload.extend_from_slice(&(WASM.len() as u32).to_le_bytes());
        payload
    }

    // Regression test for https://github.com/sile/libflate/issues/90 :
    // decoding a single, highly-compressible DEFLATE block used to buffer the
    // entire expanded block before the first `read` returned, which allowed the
    // internal buffer to grow without bound. The decoder must now yield output
    // in bounded chunks (at most `MAX_INTERNAL_BUFFER + 258` unread bytes).
    #[test]
    #[cfg(feature = "std")]
    fn test_issue_90_bounded_buffering() {
        use no_std_io2::io::Write as _;

        let text = b"abcdefgh".repeat(100_000); // 800 KiB, a single DEFLATE block
        let mut encoder = crate::deflate::Encoder::new(Vec::new());
        encoder.write_all(&text).unwrap();
        let encoded = encoder.finish().into_result().unwrap();

        let mut decoder = Decoder::new(&encoded[..]);
        let mut output = Vec::new();
        let mut chunk = [0u8; 1024];
        loop {
            let n = decoder.read(&mut chunk).unwrap();
            if n == 0 {
                break;
            }
            output.extend_from_slice(&chunk[..n]);
            let buffered = decoder.unread_decoded_data().len();
            assert!(
                buffered <= MAX_INTERNAL_BUFFER + 258,
                "internal buffer exceeded bound: {buffered} bytes"
            );
        }
        assert_eq!(output, text);
    }
}