1use std::io;
2use std::io::prelude::*;
3
4use crate::zio;
5use crate::{Compress, Decompress};
6
7/// A ZLIB encoder, or compressor.
8///
9/// This structure implements a [`Write`] interface and takes a stream of
10/// uncompressed data, writing the compressed data to the wrapped writer.
11///
12/// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
13///
14/// # Examples
15///
16/// ```
17/// use std::io::prelude::*;
18/// use flate2::Compression;
19/// use flate2::write::ZlibEncoder;
20///
21/// // Vec<u8> implements Write, assigning the compressed bytes of sample string
22///
23/// # fn zlib_encoding() -> std::io::Result<()> {
24/// let mut e = ZlibEncoder::new(Vec::new(), Compression::default());
25/// e.write_all(b"Hello World")?;
26/// let compressed = e.finish()?;
27/// # Ok(())
28/// # }
29/// ```
30#[derive(Debug)]
31pub struct ZlibEncoder<W: Write> {
32 inner: zio::Writer<W, Compress>,
33}
34
35impl<W: Write> ZlibEncoder<W> {
36 /// Creates a new encoder which will write compressed data to the stream
37 /// given at the given compression level.
38 ///
39 /// When this encoder is dropped or unwrapped the final pieces of data will
40 /// be flushed.
41 pub fn new(w: W, level: crate::Compression) -> ZlibEncoder<W> {
42 ZlibEncoder {
43 inner: zio::Writer::new(w, Compress::new(level, true)),
44 }
45 }
46
47 /// Creates a new encoder which will write compressed data to the stream
48 /// `w` with the given `compression` settings.
49 pub fn new_with_compress(w: W, compression: Compress) -> ZlibEncoder<W> {
50 ZlibEncoder {
51 inner: zio::Writer::new(w, compression),
52 }
53 }
54
55 /// Acquires a reference to the underlying writer.
56 pub fn get_ref(&self) -> &W {
57 self.inner.get_ref()
58 }
59
60 /// Acquires a mutable reference to the underlying writer.
61 ///
62 /// Note that mutating the output/input state of the stream may corrupt this
63 /// object, so care must be taken when using this method.
64 pub fn get_mut(&mut self) -> &mut W {
65 self.inner.get_mut()
66 }
67
68 /// Resets the state of this encoder entirely, swapping out the output
69 /// stream for another.
70 ///
71 /// This function will finish encoding the current stream into the current
72 /// output stream before swapping out the two output streams.
73 ///
74 /// After the current stream has been finished, this will reset the internal
75 /// state of this encoder and replace the output stream with the one
76 /// provided, returning the previous output stream. Future data written to
77 /// this encoder will be the compressed into the stream `w` provided.
78 ///
79 /// # Errors
80 ///
81 /// This function will perform I/O to complete this stream, and any I/O
82 /// errors which occur will be returned from this function.
83 pub fn reset(&mut self, w: W) -> io::Result<W> {
84 self.inner.finish()?;
85 self.inner.data.reset();
86 Ok(self.inner.replace(w))
87 }
88
89 /// Attempt to finish this output stream, writing out final chunks of data.
90 ///
91 /// Note that this function can only be used once data has finished being
92 /// written to the output stream. After this function is called then further
93 /// calls to `write` may result in a panic.
94 ///
95 /// # Panics
96 ///
97 /// Attempts to write data to this stream may result in a panic after this
98 /// function is called.
99 ///
100 /// # Errors
101 ///
102 /// This function will perform I/O to complete this stream, and any I/O
103 /// errors which occur will be returned from this function.
104 pub fn try_finish(&mut self) -> io::Result<()> {
105 self.inner.finish()
106 }
107
108 /// Consumes this encoder, flushing the output stream.
109 ///
110 /// This will flush the underlying data stream, close off the compressed
111 /// stream and, if successful, return the contained writer.
112 ///
113 /// Note that this function may not be suitable to call in a situation where
114 /// the underlying stream is an asynchronous I/O stream. To finish a stream
115 /// the `try_finish` (or `shutdown`) method should be used instead. To
116 /// re-acquire ownership of a stream it is safe to call this method after
117 /// `try_finish` or `shutdown` has returned `Ok`.
118 ///
119 /// # Errors
120 ///
121 /// This function will perform I/O to complete this stream, and any I/O
122 /// errors which occur will be returned from this function.
123 pub fn finish(mut self) -> io::Result<W> {
124 self.inner.finish()?;
125 Ok(self.inner.take_inner())
126 }
127
128 /// Consumes this encoder, flushing the output stream.
129 ///
130 /// This will flush the underlying data stream and then return the contained
131 /// writer if the flush succeeded.
132 /// The compressed stream will not closed but only flushed. This
133 /// means that obtained byte array can by extended by another deflated
134 /// stream. To close the stream add the two bytes 0x3 and 0x0.
135 ///
136 /// # Errors
137 ///
138 /// This function will perform I/O to complete this stream, and any I/O
139 /// errors which occur will be returned from this function.
140 pub fn flush_finish(mut self) -> io::Result<W> {
141 self.inner.flush()?;
142 Ok(self.inner.take_inner())
143 }
144
145 /// Returns the number of bytes that have been written to this compressor.
146 ///
147 /// Note that not all bytes written to this object may be accounted for,
148 /// there may still be some active buffering.
149 pub fn total_in(&self) -> u64 {
150 self.inner.data.total_in()
151 }
152
153 /// Returns the number of bytes that the compressor has produced.
154 ///
155 /// Note that not all bytes may have been written yet, some may still be
156 /// buffered.
157 pub fn total_out(&self) -> u64 {
158 self.inner.data.total_out()
159 }
160}
161
162impl<W: Write> Write for ZlibEncoder<W> {
163 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
164 self.inner.write(buf)
165 }
166
167 fn flush(&mut self) -> io::Result<()> {
168 self.inner.flush()
169 }
170}
171
172impl<W: Read + Write> Read for ZlibEncoder<W> {
173 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
174 self.get_mut().read(buf)
175 }
176}
177
178/// A ZLIB decoder, or decompressor.
179///
180/// This structure implements a [`Write`] and will emit a stream of decompressed
181/// data when fed a stream of compressed data.
182///
183/// After decoding a single member of the ZLIB data this writer will return the number of bytes up to
184/// to the end of the ZLIB member and subsequent writes will return Ok(0) allowing the caller to
185/// handle any data following the ZLIB member.
186///
187/// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
188///
189/// # Examples
190///
191/// ```
192/// use std::io::prelude::*;
193/// use std::io;
194/// # use flate2::Compression;
195/// # use flate2::write::ZlibEncoder;
196/// use flate2::write::ZlibDecoder;
197///
198/// # fn main() {
199/// # let mut e = ZlibEncoder::new(Vec::new(), Compression::default());
200/// # e.write_all(b"Hello World").unwrap();
201/// # let bytes = e.finish().unwrap();
202/// # println!("{}", decode_reader(bytes).unwrap());
203/// # }
204/// #
205/// // Uncompresses a Zlib Encoded vector of bytes and returns a string or error
206/// // Here Vec<u8> implements Write
207///
208/// fn decode_reader(bytes: Vec<u8>) -> io::Result<String> {
209/// let mut writer = Vec::new();
210/// let mut z = ZlibDecoder::new(writer);
211/// z.write_all(&bytes[..])?;
212/// writer = z.finish()?;
213/// let return_string = String::from_utf8(writer).expect("String parsing error");
214/// Ok(return_string)
215/// }
216/// ```
217#[derive(Debug)]
218pub struct ZlibDecoder<W: Write> {
219 inner: zio::Writer<W, Decompress>,
220}
221
222impl<W: Write> ZlibDecoder<W> {
223 /// Creates a new decoder which will write uncompressed data to the stream.
224 ///
225 /// When this decoder is dropped or unwrapped the final pieces of data will
226 /// be flushed.
227 pub fn new(w: W) -> ZlibDecoder<W> {
228 ZlibDecoder {
229 inner: zio::Writer::new(w, Decompress::new(true)),
230 }
231 }
232
233 /// Creates a new decoder which will write uncompressed data to the stream `w`
234 /// using the given `decompression` settings.
235 ///
236 /// When this decoder is dropped or unwrapped the final pieces of data will
237 /// be flushed.
238 pub fn new_with_decompress(w: W, decompression: Decompress) -> ZlibDecoder<W> {
239 ZlibDecoder {
240 inner: zio::Writer::new(w, decompression),
241 }
242 }
243
244 /// Acquires a reference to the underlying writer.
245 pub fn get_ref(&self) -> &W {
246 self.inner.get_ref()
247 }
248
249 /// Acquires a mutable reference to the underlying writer.
250 ///
251 /// Note that mutating the output/input state of the stream may corrupt this
252 /// object, so care must be taken when using this method.
253 pub fn get_mut(&mut self) -> &mut W {
254 self.inner.get_mut()
255 }
256
257 /// Resets the state of this decoder entirely, swapping out the output
258 /// stream for another.
259 ///
260 /// This will reset the internal state of this decoder and replace the
261 /// output stream with the one provided, returning the previous output
262 /// stream. Future data written to this decoder will be decompressed into
263 /// the output stream `w`.
264 ///
265 /// # Errors
266 ///
267 /// This function will perform I/O to complete this stream, and any I/O
268 /// errors which occur will be returned from this function.
269 pub fn reset(&mut self, w: W) -> io::Result<W> {
270 self.inner.finish()?;
271 self.inner.data = Decompress::new(true);
272 Ok(self.inner.replace(w))
273 }
274
275 /// Attempt to finish this output stream, writing out final chunks of data.
276 ///
277 /// Note that this function can only be used once data has finished being
278 /// written to the output stream. After this function is called then further
279 /// calls to `write` may result in a panic.
280 ///
281 /// # Panics
282 ///
283 /// Attempts to write data to this stream may result in a panic after this
284 /// function is called.
285 ///
286 /// # Errors
287 ///
288 /// This function will perform I/O to complete this stream, and any I/O
289 /// errors which occur will be returned from this function.
290 pub fn try_finish(&mut self) -> io::Result<()> {
291 self.inner.finish()
292 }
293
294 /// Consumes this encoder, flushing the output stream.
295 ///
296 /// This will flush the underlying data stream and then return the contained
297 /// writer if the flush succeeded.
298 ///
299 /// Note that this function may not be suitable to call in a situation where
300 /// the underlying stream is an asynchronous I/O stream. To finish a stream
301 /// the `try_finish` (or `shutdown`) method should be used instead. To
302 /// re-acquire ownership of a stream it is safe to call this method after
303 /// `try_finish` or `shutdown` has returned `Ok`.
304 ///
305 /// # Errors
306 ///
307 /// This function will perform I/O to complete this stream, and any I/O
308 /// errors which occur will be returned from this function.
309 pub fn finish(mut self) -> io::Result<W> {
310 self.inner.finish()?;
311 Ok(self.inner.take_inner())
312 }
313
314 /// Returns the number of bytes that the decompressor has consumed for
315 /// decompression.
316 ///
317 /// Note that this will likely be smaller than the number of bytes
318 /// successfully written to this stream due to internal buffering.
319 pub fn total_in(&self) -> u64 {
320 self.inner.data.total_in()
321 }
322
323 /// Returns the number of bytes that the decompressor has written to its
324 /// output stream.
325 pub fn total_out(&self) -> u64 {
326 self.inner.data.total_out()
327 }
328}
329
330impl<W: Write> Write for ZlibDecoder<W> {
331 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
332 self.inner.write(buf)
333 }
334
335 fn flush(&mut self) -> io::Result<()> {
336 self.inner.flush()
337 }
338}
339
340impl<W: Read + Write> Read for ZlibDecoder<W> {
341 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
342 self.inner.get_mut().read(buf)
343 }
344}
345
346#[cfg(test)]
347mod tests {
348 use super::*;
349 use crate::Compression;
350
351 const STR: &str = "Hello World Hello World Hello World Hello World Hello World \
352 Hello World Hello World Hello World Hello World Hello World \
353 Hello World Hello World Hello World Hello World Hello World \
354 Hello World Hello World Hello World Hello World Hello World \
355 Hello World Hello World Hello World Hello World Hello World";
356
357 // ZlibDecoder consumes one zlib archive and then returns 0 for subsequent writes, allowing any
358 // additional data to be consumed by the caller.
359 #[test]
360 fn decode_extra_data() {
361 let compressed = {
362 let mut e = ZlibEncoder::new(Vec::new(), Compression::default());
363 e.write(STR.as_ref()).unwrap();
364 let mut b = e.finish().unwrap();
365 b.push(b'x');
366 b
367 };
368
369 let mut writer = Vec::new();
370 let mut decoder = ZlibDecoder::new(writer);
371 let mut consumed_bytes = 0;
372 loop {
373 let n = decoder.write(&compressed[consumed_bytes..]).unwrap();
374 if n == 0 {
375 break;
376 }
377 consumed_bytes += n;
378 }
379 writer = decoder.finish().unwrap();
380 let actual = String::from_utf8(writer).expect("String parsing error");
381 assert_eq!(actual, STR);
382 assert_eq!(&compressed[consumed_bytes..], b"x");
383 }
384}
385