1 //! Extra streaming decompression functionality.
2 //!
3 //! As of now this is mainly inteded for use to build a higher-level wrapper.
4 use std::io::Cursor;
5 use std::{cmp, mem};
6
7 use crate::inflate::core::{decompress, inflate_flags, DecompressorOxide, TINFL_LZ_DICT_SIZE};
8 use crate::inflate::TINFLStatus;
9 use crate::{DataFormat, MZError, MZFlush, MZResult, MZStatus, StreamResult};
10
11 /// A struct that compbines a decompressor with extra data for streaming decompression.
12 ///
13 pub struct InflateState {
14 /// Inner decompressor struct
15 decomp: DecompressorOxide,
16
17 /// Buffer of input bytes for matches.
18 /// TODO: Could probably do this a bit cleaner with some
19 /// Cursor-like class.
20 /// We may also look into whether we need to keep a buffer here, or just one in the
21 /// decompressor struct.
22 dict: [u8; TINFL_LZ_DICT_SIZE],
23 /// Where in the buffer are we currently at?
24 dict_ofs: usize,
25 /// How many bytes of data to be flushed is there currently in the buffer?
26 dict_avail: usize,
27
28 first_call: bool,
29 has_flushed: bool,
30
31 /// Whether the input data is wrapped in a zlib header and checksum.
32 /// TODO: This should be stored in the decompressor.
33 data_format: DataFormat,
34 last_status: TINFLStatus,
35 }
36
37 impl Default for InflateState {
default() -> Self38 fn default() -> Self {
39 InflateState {
40 decomp: DecompressorOxide::default(),
41 dict: [0; TINFL_LZ_DICT_SIZE],
42 dict_ofs: 0,
43 dict_avail: 0,
44 first_call: true,
45 has_flushed: false,
46 data_format: DataFormat::Raw,
47 last_status: TINFLStatus::NeedsMoreInput,
48 }
49 }
50 }
51 impl InflateState {
52 /// Create a new state.
53 ///
54 /// Note that this struct is quite large due to internal buffers, and as such storing it on
55 /// the stack is not recommended.
56 ///
57 /// # Parameters
58 /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
59 /// metadata.
new(data_format: DataFormat) -> InflateState60 pub fn new(data_format: DataFormat) -> InflateState {
61 let mut b = InflateState::default();
62 b.data_format = data_format;
63 b
64 }
65
66 /// Create a new state on the heap.
67 ///
68 /// # Parameters
69 /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
70 /// metadata.
new_boxed(data_format: DataFormat) -> Box<InflateState>71 pub fn new_boxed(data_format: DataFormat) -> Box<InflateState> {
72 let mut b: Box<InflateState> = Box::default();
73 b.data_format = data_format;
74 b
75 }
76
77 /// Access the innner decompressor.
decompressor(&mut self) -> &mut DecompressorOxide78 pub fn decompressor(&mut self) -> &mut DecompressorOxide {
79 &mut self.decomp
80 }
81
82 /// Return the status of the last call to `inflate` with this `InflateState`.
last_status(&self) -> TINFLStatus83 pub fn last_status(&self) -> TINFLStatus {
84 self.last_status
85 }
86
87 /// Create a new state using miniz/zlib style window bits parameter.
88 ///
89 /// The decompressor does not support different window sizes. As such,
90 /// any positive (>0) value will set the zlib header flag, while a negative one
91 /// will not.
new_boxed_with_window_bits(window_bits: i32) -> Box<InflateState>92 pub fn new_boxed_with_window_bits(window_bits: i32) -> Box<InflateState> {
93 let mut b: Box<InflateState> = Box::default();
94 b.data_format = DataFormat::from_window_bits(window_bits);
95 b
96 }
97
98 /// Reset the decompressor without re-allocating memory, using the given
99 /// data format.
reset(&mut self, data_format: DataFormat)100 pub fn reset(&mut self, data_format: DataFormat) {
101 self.decompressor().init();
102 self.dict = [0; TINFL_LZ_DICT_SIZE];
103 self.dict_ofs = 0;
104 self.dict_avail = 0;
105 self.first_call = true;
106 self.has_flushed = false;
107 self.data_format = data_format;
108 self.last_status = TINFLStatus::NeedsMoreInput;
109 }
110 }
111
112 /// Try to decompress from `input` to `output` with the given `InflateState`
113 ///
114 /// # Errors
115 ///
116 /// Returns `MZError::Buf` If the size of the `output` slice is empty or no progress was made due to
117 /// lack of expected input data or called after the decompression was
118 /// finished without MZFlush::Finish.
119 ///
120 /// Returns `MZError::Param` if the compressor parameters are set wrong.
inflate( state: &mut InflateState, input: &[u8], output: &mut [u8], flush: MZFlush, ) -> StreamResult121 pub fn inflate(
122 state: &mut InflateState,
123 input: &[u8],
124 output: &mut [u8],
125 flush: MZFlush,
126 ) -> StreamResult {
127 let mut bytes_consumed = 0;
128 let mut bytes_written = 0;
129 let mut next_in = input;
130 let mut next_out = output;
131
132 if flush == MZFlush::Full {
133 return StreamResult::error(MZError::Stream);
134 }
135
136 let mut decomp_flags = inflate_flags::TINFL_FLAG_COMPUTE_ADLER32;
137 if state.data_format == DataFormat::Zlib {
138 decomp_flags |= inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER;
139 }
140
141 let first_call = state.first_call;
142 state.first_call = false;
143 if (state.last_status as i32) < 0 {
144 return StreamResult::error(MZError::Data);
145 }
146
147 if state.has_flushed && (flush != MZFlush::Finish) {
148 return StreamResult::error(MZError::Stream);
149 }
150 state.has_flushed |= flush == MZFlush::Finish;
151
152 if (flush == MZFlush::Finish) && first_call {
153 decomp_flags |= inflate_flags::TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF;
154
155 let status = decompress(
156 &mut state.decomp,
157 next_in,
158 &mut Cursor::new(next_out),
159 decomp_flags,
160 );
161 let in_bytes = status.1;
162 let out_bytes = status.2;
163 let status = status.0;
164
165 state.last_status = status;
166
167 bytes_consumed += in_bytes;
168 bytes_written += out_bytes;
169
170 let ret_status = {
171 if (status as i32) < 0 {
172 Err(MZError::Data)
173 } else if status != TINFLStatus::Done {
174 state.last_status = TINFLStatus::Failed;
175 Err(MZError::Buf)
176 } else {
177 Ok(MZStatus::StreamEnd)
178 }
179 };
180 return StreamResult {
181 bytes_consumed,
182 bytes_written,
183 status: ret_status,
184 };
185 }
186
187 if flush != MZFlush::Finish {
188 decomp_flags |= inflate_flags::TINFL_FLAG_HAS_MORE_INPUT;
189 }
190
191 if state.dict_avail != 0 {
192 bytes_written += push_dict_out(state, &mut next_out);
193 return StreamResult {
194 bytes_consumed,
195 bytes_written,
196 status: Ok(
197 if (state.last_status == TINFLStatus::Done) && (state.dict_avail == 0) {
198 MZStatus::StreamEnd
199 } else {
200 MZStatus::Ok
201 },
202 ),
203 };
204 }
205
206 let status = inflate_loop(
207 state,
208 &mut next_in,
209 &mut next_out,
210 &mut bytes_consumed,
211 &mut bytes_written,
212 decomp_flags,
213 flush,
214 );
215 StreamResult {
216 bytes_consumed,
217 bytes_written,
218 status,
219 }
220 }
221
inflate_loop( state: &mut InflateState, next_in: &mut &[u8], next_out: &mut &mut [u8], total_in: &mut usize, total_out: &mut usize, decomp_flags: u32, flush: MZFlush, ) -> MZResult222 fn inflate_loop(
223 state: &mut InflateState,
224 next_in: &mut &[u8],
225 next_out: &mut &mut [u8],
226 total_in: &mut usize,
227 total_out: &mut usize,
228 decomp_flags: u32,
229 flush: MZFlush,
230 ) -> MZResult {
231 let orig_in_len = next_in.len();
232 loop {
233 let status = {
234 let mut cursor = Cursor::new(&mut state.dict[..]);
235 cursor.set_position(state.dict_ofs as u64);
236 decompress(&mut state.decomp, *next_in, &mut cursor, decomp_flags)
237 };
238
239 let in_bytes = status.1;
240 let out_bytes = status.2;
241 let status = status.0;
242
243 state.last_status = status;
244
245 *next_in = &next_in[in_bytes..];
246 *total_in += in_bytes;
247
248 state.dict_avail = out_bytes;
249 *total_out += push_dict_out(state, next_out);
250
251 // The stream was corrupted, and decompression failed.
252 if (status as i32) < 0 {
253 return Err(MZError::Data);
254 }
255
256 // The decompressor has flushed all it's data and is waiting for more input, but
257 // there was no more input provided.
258 if (status == TINFLStatus::NeedsMoreInput) && orig_in_len == 0 {
259 return Err(MZError::Buf);
260 }
261
262 if flush == MZFlush::Finish {
263 if status == TINFLStatus::Done {
264 // There is not enough space in the output buffer to flush the remaining
265 // decompressed data in the internal buffer.
266 return if state.dict_avail != 0 {
267 Err(MZError::Buf)
268 } else {
269 Ok(MZStatus::StreamEnd)
270 };
271 // No more space in the output buffer, but we're not done.
272 } else if next_out.is_empty() {
273 return Err(MZError::Buf);
274 }
275 } else {
276 // We're not expected to finish, so it's fine if we can't flush everything yet.
277 let empty_buf = next_in.is_empty() || next_out.is_empty();
278 if (status == TINFLStatus::Done) || empty_buf || (state.dict_avail != 0) {
279 return if (status == TINFLStatus::Done) && (state.dict_avail == 0) {
280 // No more data left, we're done.
281 Ok(MZStatus::StreamEnd)
282 } else {
283 // Ok for now, still waiting for more input data or output space.
284 Ok(MZStatus::Ok)
285 };
286 }
287 }
288 }
289 }
290
push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize291 fn push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize {
292 let n = cmp::min(state.dict_avail as usize, next_out.len());
293 (next_out[..n]).copy_from_slice(&state.dict[state.dict_ofs..state.dict_ofs + n]);
294 *next_out = &mut mem::replace(next_out, &mut [])[n..];
295 state.dict_avail -= n;
296 state.dict_ofs = (state.dict_ofs + (n)) & (TINFL_LZ_DICT_SIZE - 1);
297 n
298 }
299
300 #[cfg(test)]
301 mod test {
302 use super::{inflate, InflateState};
303 use crate::{DataFormat, MZFlush, MZStatus};
304 #[test]
test_state()305 fn test_state() {
306 let encoded = [
307 120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
308 19,
309 ];
310 let mut out = vec![0; 50];
311 let mut state = InflateState::new_boxed(DataFormat::Zlib);
312 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
313 let status = res.status.expect("Failed to decompress!");
314 assert_eq!(status, MZStatus::StreamEnd);
315 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
316 assert_eq!(res.bytes_consumed, encoded.len());
317
318 state.reset(DataFormat::Zlib);
319 let status = res.status.expect("Failed to decompress!");
320 assert_eq!(status, MZStatus::StreamEnd);
321 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
322 assert_eq!(res.bytes_consumed, encoded.len());
323 }
324 }
325