1 use crate::deflate::symbol::{self, HuffmanCodec};
2 use crate::lz77;
3 use crate::non_blocking::transaction::TransactionalBitReader;
4 use std::cmp;
5 use std::io;
6 use std::io::Read;
7 
8 /// DEFLATE decoder which supports non-blocking I/O.
9 #[derive(Debug)]
10 pub struct Decoder<R> {
11     state: DecoderState,
12     eos: bool,
13     bit_reader: TransactionalBitReader<R>,
14     block_decoder: BlockDecoder,
15 }
16 impl<R: Read> Decoder<R> {
17     /// Makes a new decoder instance.
18     ///
19     /// `inner` is to be decoded DEFLATE stream.
20     ///
21     /// # Examples
22     /// ```
23     /// use std::io::{Cursor, Read};
24     /// use libflate::non_blocking::deflate::Decoder;
25     ///
26     /// let encoded_data = [243, 72, 205, 201, 201, 87, 8, 207, 47, 202, 73, 81, 4, 0];
27     /// let mut decoder = Decoder::new(&encoded_data[..]);
28     /// let mut buf = Vec::new();
29     /// decoder.read_to_end(&mut buf).unwrap();
30     ///
31     /// assert_eq!(buf, b"Hello World!");
32     /// ```
new(inner: R) -> Self33     pub fn new(inner: R) -> Self {
34         Decoder {
35             state: DecoderState::ReadBlockHeader,
36             eos: false,
37             bit_reader: TransactionalBitReader::new(inner),
38             block_decoder: BlockDecoder::new(),
39         }
40     }
41 
42     /// Returns the immutable reference to the inner stream.
as_inner_ref(&self) -> &R43     pub fn as_inner_ref(&self) -> &R {
44         self.bit_reader.as_inner_ref()
45     }
46 
47     /// Returns the mutable reference to the inner stream.
as_inner_mut(&mut self) -> &mut R48     pub fn as_inner_mut(&mut self) -> &mut R {
49         self.bit_reader.as_inner_mut()
50     }
51 
52     /// Unwraps this `Decoder`, returning the underlying reader.
53     ///
54     /// # Examples
55     /// ```
56     /// use std::io::Cursor;
57     /// use libflate::non_blocking::deflate::Decoder;
58     ///
59     /// let encoded_data = [243, 72, 205, 201, 201, 87, 8, 207, 47, 202, 73, 81, 4, 0];
60     /// let decoder = Decoder::new(Cursor::new(&encoded_data));
61     /// assert_eq!(decoder.into_inner().into_inner(), &encoded_data);
62     /// ```
into_inner(self) -> R63     pub fn into_inner(self) -> R {
64         self.bit_reader.into_inner()
65     }
66 
bit_reader_mut(&mut self) -> &mut TransactionalBitReader<R>67     pub(crate) fn bit_reader_mut(&mut self) -> &mut TransactionalBitReader<R> {
68         &mut self.bit_reader
69     }
70 }
71 impl<R: Read> Read for Decoder<R> {
read(&mut self, buf: &mut [u8]) -> io::Result<usize>72     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
73         let mut read_size;
74         loop {
75             let next = match self.state {
76                 DecoderState::ReadBlockHeader => {
77                     let (bfinal, btype) = self.bit_reader.transaction(|r| {
78                         let bfinal = r.read_bit()?;
79                         let btype = r.read_bits(2)?;
80                         Ok((bfinal, btype))
81                     })?;
82                     self.eos = bfinal;
83                     self.block_decoder.enter_new_block();
84                     match btype {
85                         0b00 => DecoderState::ReadNonCompressedBlockLen,
86                         0b01 => DecoderState::LoadFixedHuffmanCode,
87                         0b10 => DecoderState::LoadDynamicHuffmanCode,
88                         0b11 => {
89                             return Err(invalid_data_error!(
90                                 "btype 0x11 of DEFLATE is reserved(error) value"
91                             ));
92                         }
93                         _ => unreachable!(),
94                     }
95                 }
96                 DecoderState::ReadNonCompressedBlockLen => {
97                     let len = self.bit_reader.transaction(|r| {
98                         r.reset();
99                         let mut buf = [0; 2];
100                         r.as_inner_mut().read_exact(&mut buf)?;
101                         let len = u16::from_le_bytes(buf);
102                         r.as_inner_mut().read_exact(&mut buf)?;
103                         let nlen = u16::from_le_bytes(buf);
104                         if !len != nlen {
105                             Err(invalid_data_error!(
106                                 "LEN={} is not the one's complement of NLEN={}",
107                                 len,
108                                 nlen
109                             ))
110                         } else {
111                             Ok(len)
112                         }
113                     })?;
114                     DecoderState::ReadNonCompressedBlock { len }
115                 }
116                 DecoderState::ReadNonCompressedBlock { len: 0 } => {
117                     if self.eos {
118                         read_size = 0;
119                         break;
120                     } else {
121                         DecoderState::ReadBlockHeader
122                     }
123                 }
124                 DecoderState::ReadNonCompressedBlock { ref mut len } => {
125                     let buf_len = buf.len();
126                     let buf = &mut buf[..cmp::min(buf_len, *len as usize)];
127                     read_size = self.bit_reader.as_inner_mut().read(buf)?;
128 
129                     self.block_decoder.extend(&buf[..read_size]);
130                     *len -= read_size as u16;
131                     break;
132                 }
133                 DecoderState::LoadFixedHuffmanCode => {
134                     let symbol_decoder = self
135                         .bit_reader
136                         .transaction(|r| symbol::FixedHuffmanCodec.load(r))?;
137                     DecoderState::DecodeBlock(symbol_decoder)
138                 }
139                 DecoderState::LoadDynamicHuffmanCode => {
140                     let symbol_decoder = self
141                         .bit_reader
142                         .transaction(|r| symbol::DynamicHuffmanCodec.load(r))?;
143                     DecoderState::DecodeBlock(symbol_decoder)
144                 }
145                 DecoderState::DecodeBlock(ref mut symbol_decoder) => {
146                     self.block_decoder
147                         .decode(&mut self.bit_reader, symbol_decoder)?;
148                     read_size = self.block_decoder.read(buf)?;
149                     if read_size == 0 && !buf.is_empty() && !self.eos {
150                         DecoderState::ReadBlockHeader
151                     } else {
152                         break;
153                     }
154                 }
155             };
156             self.state = next;
157         }
158         Ok(read_size)
159     }
160 }
161 
162 #[derive(Debug)]
163 enum DecoderState {
164     ReadBlockHeader,
165     ReadNonCompressedBlockLen,
166     ReadNonCompressedBlock { len: u16 },
167     LoadFixedHuffmanCode,
168     LoadDynamicHuffmanCode,
169     DecodeBlock(symbol::Decoder),
170 }
171 
172 #[derive(Debug)]
173 struct BlockDecoder {
174     lz77_decoder: lz77::Lz77Decoder,
175     eob: bool,
176 }
177 impl BlockDecoder {
new() -> Self178     pub fn new() -> Self {
179         BlockDecoder {
180             lz77_decoder: lz77::Lz77Decoder::new(),
181             eob: false,
182         }
183     }
enter_new_block(&mut self)184     pub fn enter_new_block(&mut self) {
185         self.eob = false;
186     }
decode<R: Read>( &mut self, bit_reader: &mut TransactionalBitReader<R>, symbol_decoder: &mut symbol::Decoder, ) -> io::Result<()>187     pub fn decode<R: Read>(
188         &mut self,
189         bit_reader: &mut TransactionalBitReader<R>,
190         symbol_decoder: &mut symbol::Decoder,
191     ) -> io::Result<()> {
192         if self.eob {
193             return Ok(());
194         }
195         while let Some(s) = self.decode_symbol(bit_reader, symbol_decoder)? {
196             match s {
197                 symbol::Symbol::Code(code) => {
198                     self.lz77_decoder.decode(code)?;
199                 }
200                 symbol::Symbol::EndOfBlock => {
201                     self.eob = true;
202                     break;
203                 }
204             }
205         }
206         Ok(())
207     }
208 
extend(&mut self, buf: &[u8])209     fn extend(&mut self, buf: &[u8]) {
210         self.lz77_decoder.extend_from_slice(buf);
211     }
212 
decode_symbol<R: Read>( &mut self, bit_reader: &mut TransactionalBitReader<R>, symbol_decoder: &mut symbol::Decoder, ) -> io::Result<Option<symbol::Symbol>>213     fn decode_symbol<R: Read>(
214         &mut self,
215         bit_reader: &mut TransactionalBitReader<R>,
216         symbol_decoder: &mut symbol::Decoder,
217     ) -> io::Result<Option<symbol::Symbol>> {
218         let result = bit_reader.transaction(|bit_reader| {
219             let s = symbol_decoder.decode_unchecked(bit_reader);
220             bit_reader.check_last_error().map(|()| s)
221         });
222         match result {
223             Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(None),
224             Err(e) => Err(e),
225             Ok(s) => Ok(Some(s)),
226         }
227     }
228 }
229 impl Read for BlockDecoder {
read(&mut self, buf: &mut [u8]) -> io::Result<usize>230     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
231         if !self.lz77_decoder.buffer().is_empty() {
232             self.lz77_decoder.read(buf)
233         } else if self.eob {
234             Ok(0)
235         } else {
236             Err(io::Error::new(io::ErrorKind::WouldBlock, "Would block"))
237         }
238     }
239 }
240 
241 #[cfg(test)]
242 mod tests {
243     use super::*;
244     use crate::deflate::{EncodeOptions, Encoder};
245     use crate::util::{nb_read_to_end, WouldBlockReader};
246     use std::io::{self, Read};
247 
248     #[test]
it_works()249     fn it_works() {
250         let mut encoder = Encoder::new(Vec::new());
251         io::copy(&mut &b"Hello World!"[..], &mut encoder).unwrap();
252         let encoded_data = encoder.finish().into_result().unwrap();
253 
254         let mut decoder = Decoder::new(&encoded_data[..]);
255         let mut decoded_data = Vec::new();
256         decoder.read_to_end(&mut decoded_data).unwrap();
257 
258         assert_eq!(decoded_data, b"Hello World!");
259     }
260 
261     #[test]
non_blocking_io_works()262     fn non_blocking_io_works() {
263         let mut encoder = Encoder::new(Vec::new());
264         io::copy(&mut &b"Hello World!"[..], &mut encoder).unwrap();
265         let encoded_data = encoder.finish().into_result().unwrap();
266 
267         let decoder = Decoder::new(WouldBlockReader::new(&encoded_data[..]));
268         let decoded_data = nb_read_to_end(decoder).unwrap();
269 
270         assert_eq!(decoded_data, b"Hello World!");
271     }
272 
273     #[test]
non_blocking_io_for_large_text_works()274     fn non_blocking_io_for_large_text_works() {
275         let text: String = (0..10000)
276             .into_iter()
277             .map(|i| format!("test {}", i))
278             .collect();
279 
280         let mut encoder = crate::deflate::Encoder::new(Vec::new());
281         io::copy(&mut text.as_bytes(), &mut encoder).unwrap();
282         let encoded_data = encoder.finish().into_result().unwrap();
283 
284         let decoder = Decoder::new(WouldBlockReader::new(&encoded_data[..]));
285         let decoded_data = nb_read_to_end(decoder).unwrap();
286         assert_eq!(decoded_data, text.as_bytes());
287     }
288 
289     #[test]
non_compressed_non_blocking_io_works()290     fn non_compressed_non_blocking_io_works() {
291         let mut encoder = Encoder::with_options(Vec::new(), EncodeOptions::new().no_compression());
292         io::copy(&mut &b"Hello World!"[..], &mut encoder).unwrap();
293         let encoded_data = encoder.finish().into_result().unwrap();
294 
295         let decoder = Decoder::new(WouldBlockReader::new(&encoded_data[..]));
296         let decoded_data = nb_read_to_end(decoder).unwrap();
297 
298         assert_eq!(decoded_data, b"Hello World!");
299     }
300 }
301