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