]> git.proxmox.com Git - rustc.git/blob - vendor/flate2/src/gz/write.rs
New upstream version 1.41.1+dfsg1
[rustc.git] / vendor / flate2 / src / gz / write.rs
1 use std::cmp;
2 use std::io;
3 use std::io::prelude::*;
4
5 #[cfg(feature = "tokio")]
6 use futures::Poll;
7 #[cfg(feature = "tokio")]
8 use tokio_io::{AsyncRead, AsyncWrite};
9
10 use super::bufread::{corrupt, read_gz_header};
11 use super::{GzBuilder, GzHeader};
12 use crate::crc::{Crc, CrcWriter};
13 use crate::zio;
14 use crate::{Compress, Compression, Decompress, Status};
15
16 /// A gzip streaming encoder
17 ///
18 /// This structure exposes a [`Write`] interface that will emit compressed data
19 /// to the underlying writer `W`.
20 ///
21 /// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
22 ///
23 /// # Examples
24 ///
25 /// ```
26 /// use std::io::prelude::*;
27 /// use flate2::Compression;
28 /// use flate2::write::GzEncoder;
29 ///
30 /// // Vec<u8> implements Write to print the compressed bytes of sample string
31 /// # fn main() {
32 ///
33 /// let mut e = GzEncoder::new(Vec::new(), Compression::default());
34 /// e.write_all(b"Hello World").unwrap();
35 /// println!("{:?}", e.finish().unwrap());
36 /// # }
37 /// ```
38 #[derive(Debug)]
39 pub struct GzEncoder<W: Write> {
40 inner: zio::Writer<W, Compress>,
41 crc: Crc,
42 crc_bytes_written: usize,
43 header: Vec<u8>,
44 }
45
46 pub fn gz_encoder<W: Write>(header: Vec<u8>, w: W, lvl: Compression) -> GzEncoder<W> {
47 GzEncoder {
48 inner: zio::Writer::new(w, Compress::new(lvl, false)),
49 crc: Crc::new(),
50 header: header,
51 crc_bytes_written: 0,
52 }
53 }
54
55 impl<W: Write> GzEncoder<W> {
56 /// Creates a new encoder which will use the given compression level.
57 ///
58 /// The encoder is not configured specially for the emitted header. For
59 /// header configuration, see the `GzBuilder` type.
60 ///
61 /// The data written to the returned encoder will be compressed and then
62 /// written to the stream `w`.
63 pub fn new(w: W, level: Compression) -> GzEncoder<W> {
64 GzBuilder::new().write(w, level)
65 }
66
67 /// Acquires a reference to the underlying writer.
68 pub fn get_ref(&self) -> &W {
69 self.inner.get_ref()
70 }
71
72 /// Acquires a mutable reference to the underlying writer.
73 ///
74 /// Note that mutation of the writer may result in surprising results if
75 /// this encoder is continued to be used.
76 pub fn get_mut(&mut self) -> &mut W {
77 self.inner.get_mut()
78 }
79
80 /// Attempt to finish this output stream, writing out final chunks of data.
81 ///
82 /// Note that this function can only be used once data has finished being
83 /// written to the output stream. After this function is called then further
84 /// calls to `write` may result in a panic.
85 ///
86 /// # Panics
87 ///
88 /// Attempts to write data to this stream may result in a panic after this
89 /// function is called.
90 ///
91 /// # Errors
92 ///
93 /// This function will perform I/O to complete this stream, and any I/O
94 /// errors which occur will be returned from this function.
95 pub fn try_finish(&mut self) -> io::Result<()> {
96 self.write_header()?;
97 self.inner.finish()?;
98
99 while self.crc_bytes_written < 8 {
100 let (sum, amt) = (self.crc.sum() as u32, self.crc.amount());
101 let buf = [
102 (sum >> 0) as u8,
103 (sum >> 8) as u8,
104 (sum >> 16) as u8,
105 (sum >> 24) as u8,
106 (amt >> 0) as u8,
107 (amt >> 8) as u8,
108 (amt >> 16) as u8,
109 (amt >> 24) as u8,
110 ];
111 let inner = self.inner.get_mut();
112 let n = inner.write(&buf[self.crc_bytes_written..])?;
113 self.crc_bytes_written += n;
114 }
115 Ok(())
116 }
117
118 /// Finish encoding this stream, returning the underlying writer once the
119 /// encoding is done.
120 ///
121 /// Note that this function may not be suitable to call in a situation where
122 /// the underlying stream is an asynchronous I/O stream. To finish a stream
123 /// the `try_finish` (or `shutdown`) method should be used instead. To
124 /// re-acquire ownership of a stream it is safe to call this method after
125 /// `try_finish` or `shutdown` has returned `Ok`.
126 ///
127 /// # Errors
128 ///
129 /// This function will perform I/O to complete this stream, and any I/O
130 /// errors which occur will be returned from this function.
131 pub fn finish(mut self) -> io::Result<W> {
132 self.try_finish()?;
133 Ok(self.inner.take_inner())
134 }
135
136 fn write_header(&mut self) -> io::Result<()> {
137 while self.header.len() > 0 {
138 let n = self.inner.get_mut().write(&self.header)?;
139 self.header.drain(..n);
140 }
141 Ok(())
142 }
143 }
144
145 impl<W: Write> Write for GzEncoder<W> {
146 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
147 assert_eq!(self.crc_bytes_written, 0);
148 self.write_header()?;
149 let n = self.inner.write(buf)?;
150 self.crc.update(&buf[..n]);
151 Ok(n)
152 }
153
154 fn flush(&mut self) -> io::Result<()> {
155 assert_eq!(self.crc_bytes_written, 0);
156 self.write_header()?;
157 self.inner.flush()
158 }
159 }
160
161 #[cfg(feature = "tokio")]
162 impl<W: AsyncWrite> AsyncWrite for GzEncoder<W> {
163 fn shutdown(&mut self) -> Poll<(), io::Error> {
164 self.try_finish()?;
165 self.get_mut().shutdown()
166 }
167 }
168
169 impl<R: Read + Write> Read for GzEncoder<R> {
170 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
171 self.get_mut().read(buf)
172 }
173 }
174
175 #[cfg(feature = "tokio")]
176 impl<R: AsyncRead + AsyncWrite> AsyncRead for GzEncoder<R> {}
177
178 impl<W: Write> Drop for GzEncoder<W> {
179 fn drop(&mut self) {
180 if self.inner.is_present() {
181 let _ = self.try_finish();
182 }
183 }
184 }
185
186 /// A gzip streaming decoder
187 ///
188 /// This structure exposes a [`Write`] interface that will emit compressed data
189 /// to the underlying writer `W`.
190 ///
191 /// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
192 ///
193 /// # Examples
194 ///
195 /// ```
196 /// use std::io::prelude::*;
197 /// use std::io;
198 /// use flate2::Compression;
199 /// use flate2::write::{GzEncoder, GzDecoder};
200 ///
201 /// # fn main() {
202 /// # let mut e = GzEncoder::new(Vec::new(), Compression::default());
203 /// # e.write(b"Hello World").unwrap();
204 /// # let bytes = e.finish().unwrap();
205 /// # assert_eq!("Hello World", decode_writer(bytes).unwrap());
206 /// # }
207 /// // Uncompresses a gzip encoded vector of bytes and returns a string or error
208 /// // Here Vec<u8> implements Write
209 /// fn decode_writer(bytes: Vec<u8>) -> io::Result<String> {
210 /// let mut writer = Vec::new();
211 /// let mut decoder = GzDecoder::new(writer);
212 /// decoder.write_all(&bytes[..])?;
213 /// writer = decoder.finish()?;
214 /// let return_string = String::from_utf8(writer).expect("String parsing error");
215 /// Ok(return_string)
216 /// }
217 /// ```
218 #[derive(Debug)]
219 pub struct GzDecoder<W: Write> {
220 inner: zio::Writer<CrcWriter<W>, Decompress>,
221 crc_bytes: Vec<u8>,
222 header: Option<GzHeader>,
223 header_buf: Vec<u8>,
224 }
225
226 const CRC_BYTES_LEN: usize = 8;
227
228 impl<W: Write> GzDecoder<W> {
229 /// Creates a new decoder which will write uncompressed data to the stream.
230 ///
231 /// When this encoder is dropped or unwrapped the final pieces of data will
232 /// be flushed.
233 pub fn new(w: W) -> GzDecoder<W> {
234 GzDecoder {
235 inner: zio::Writer::new(CrcWriter::new(w), Decompress::new(false)),
236 crc_bytes: Vec::with_capacity(CRC_BYTES_LEN),
237 header: None,
238 header_buf: Vec::new(),
239 }
240 }
241
242 /// Returns the header associated with this stream.
243 pub fn header(&self) -> Option<&GzHeader> {
244 self.header.as_ref()
245 }
246
247 /// Acquires a reference to the underlying writer.
248 pub fn get_ref(&self) -> &W {
249 self.inner.get_ref().get_ref()
250 }
251
252 /// Acquires a mutable reference to the underlying writer.
253 ///
254 /// Note that mutating the output/input state of the stream may corrupt this
255 /// object, so care must be taken when using this method.
256 pub fn get_mut(&mut self) -> &mut W {
257 self.inner.get_mut().get_mut()
258 }
259
260 /// Attempt to finish this output stream, writing out final chunks of data.
261 ///
262 /// Note that this function can only be used once data has finished being
263 /// written to the output stream. After this function is called then further
264 /// calls to `write` may result in a panic.
265 ///
266 /// # Panics
267 ///
268 /// Attempts to write data to this stream may result in a panic after this
269 /// function is called.
270 ///
271 /// # Errors
272 ///
273 /// This function will perform I/O to finish the stream, returning any
274 /// errors which happen.
275 pub fn try_finish(&mut self) -> io::Result<()> {
276 self.finish_and_check_crc()?;
277 Ok(())
278 }
279
280 /// Consumes this decoder, flushing the output stream.
281 ///
282 /// This will flush the underlying data stream and then return the contained
283 /// writer if the flush succeeded.
284 ///
285 /// Note that this function may not be suitable to call in a situation where
286 /// the underlying stream is an asynchronous I/O stream. To finish a stream
287 /// the `try_finish` (or `shutdown`) method should be used instead. To
288 /// re-acquire ownership of a stream it is safe to call this method after
289 /// `try_finish` or `shutdown` has returned `Ok`.
290 ///
291 /// # Errors
292 ///
293 /// This function will perform I/O to complete this stream, and any I/O
294 /// errors which occur will be returned from this function.
295 pub fn finish(mut self) -> io::Result<W> {
296 self.finish_and_check_crc()?;
297 Ok(self.inner.take_inner().into_inner())
298 }
299
300 fn finish_and_check_crc(&mut self) -> io::Result<()> {
301 self.inner.finish()?;
302
303 if self.crc_bytes.len() != 8 {
304 return Err(corrupt());
305 }
306
307 let crc = ((self.crc_bytes[0] as u32) << 0)
308 | ((self.crc_bytes[1] as u32) << 8)
309 | ((self.crc_bytes[2] as u32) << 16)
310 | ((self.crc_bytes[3] as u32) << 24);
311 let amt = ((self.crc_bytes[4] as u32) << 0)
312 | ((self.crc_bytes[5] as u32) << 8)
313 | ((self.crc_bytes[6] as u32) << 16)
314 | ((self.crc_bytes[7] as u32) << 24);
315 if crc != self.inner.get_ref().crc().sum() as u32 {
316 return Err(corrupt());
317 }
318 if amt != self.inner.get_ref().crc().amount() {
319 return Err(corrupt());
320 }
321 Ok(())
322 }
323 }
324
325 struct Counter<T: Read> {
326 inner: T,
327 pos: usize,
328 }
329
330 impl<T: Read> Read for Counter<T> {
331 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
332 let pos = self.inner.read(buf)?;
333 self.pos += pos;
334 Ok(pos)
335 }
336 }
337
338 impl<W: Write> Write for GzDecoder<W> {
339 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
340 if self.header.is_none() {
341 // trying to avoid buffer usage
342 let (res, pos) = {
343 let mut counter = Counter {
344 inner: self.header_buf.chain(buf),
345 pos: 0,
346 };
347 let res = read_gz_header(&mut counter);
348 (res, counter.pos)
349 };
350
351 match res {
352 Err(err) => {
353 if err.kind() == io::ErrorKind::UnexpectedEof {
354 // not enough data for header, save to the buffer
355 self.header_buf.extend(buf);
356 Ok(buf.len())
357 } else {
358 Err(err)
359 }
360 }
361 Ok(header) => {
362 self.header = Some(header);
363 let pos = pos - self.header_buf.len();
364 self.header_buf.truncate(0);
365 Ok(pos)
366 }
367 }
368 } else {
369 let (n, status) = self.inner.write_with_status(buf)?;
370
371 if status == Status::StreamEnd {
372 if n < buf.len() && self.crc_bytes.len() < 8 {
373 let remaining = buf.len() - n;
374 let crc_bytes = cmp::min(remaining, CRC_BYTES_LEN - self.crc_bytes.len());
375 self.crc_bytes.extend(&buf[n..n + crc_bytes]);
376 return Ok(n + crc_bytes);
377 }
378 }
379 Ok(n)
380 }
381 }
382
383 fn flush(&mut self) -> io::Result<()> {
384 self.inner.flush()
385 }
386 }
387
388 #[cfg(feature = "tokio")]
389 impl<W: AsyncWrite> AsyncWrite for GzDecoder<W> {
390 fn shutdown(&mut self) -> Poll<(), io::Error> {
391 self.try_finish()?;
392 self.inner.get_mut().get_mut().shutdown()
393 }
394 }
395
396 impl<W: Read + Write> Read for GzDecoder<W> {
397 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
398 self.inner.get_mut().get_mut().read(buf)
399 }
400 }
401
402 #[cfg(feature = "tokio")]
403 impl<W: AsyncRead + AsyncWrite> AsyncRead for GzDecoder<W> {}
404
405 #[cfg(test)]
406 mod tests {
407 use super::*;
408
409 const STR: &'static str = "Hello World Hello World Hello World Hello World Hello World \
410 Hello World Hello World Hello World Hello World Hello World \
411 Hello World Hello World Hello World Hello World Hello World \
412 Hello World Hello World Hello World Hello World Hello World \
413 Hello World Hello World Hello World Hello World Hello World";
414
415 #[test]
416 fn decode_writer_one_chunk() {
417 let mut e = GzEncoder::new(Vec::new(), Compression::default());
418 e.write(STR.as_ref()).unwrap();
419 let bytes = e.finish().unwrap();
420
421 let mut writer = Vec::new();
422 let mut decoder = GzDecoder::new(writer);
423 let n = decoder.write(&bytes[..]).unwrap();
424 decoder.write(&bytes[n..]).unwrap();
425 decoder.try_finish().unwrap();
426 writer = decoder.finish().unwrap();
427 let return_string = String::from_utf8(writer).expect("String parsing error");
428 assert_eq!(return_string, STR);
429 }
430
431 #[test]
432 fn decode_writer_partial_header() {
433 let mut e = GzEncoder::new(Vec::new(), Compression::default());
434 e.write(STR.as_ref()).unwrap();
435 let bytes = e.finish().unwrap();
436
437 let mut writer = Vec::new();
438 let mut decoder = GzDecoder::new(writer);
439 assert_eq!(decoder.write(&bytes[..5]).unwrap(), 5);
440 let n = decoder.write(&bytes[5..]).unwrap();
441 if n < bytes.len() - 5 {
442 decoder.write(&bytes[n + 5..]).unwrap();
443 }
444 writer = decoder.finish().unwrap();
445 let return_string = String::from_utf8(writer).expect("String parsing error");
446 assert_eq!(return_string, STR);
447 }
448
449 #[test]
450 fn decode_writer_exact_header() {
451 let mut e = GzEncoder::new(Vec::new(), Compression::default());
452 e.write(STR.as_ref()).unwrap();
453 let bytes = e.finish().unwrap();
454
455 let mut writer = Vec::new();
456 let mut decoder = GzDecoder::new(writer);
457 assert_eq!(decoder.write(&bytes[..10]).unwrap(), 10);
458 decoder.write(&bytes[10..]).unwrap();
459 writer = decoder.finish().unwrap();
460 let return_string = String::from_utf8(writer).expect("String parsing error");
461 assert_eq!(return_string, STR);
462 }
463
464 #[test]
465 fn decode_writer_partial_crc() {
466 let mut e = GzEncoder::new(Vec::new(), Compression::default());
467 e.write(STR.as_ref()).unwrap();
468 let bytes = e.finish().unwrap();
469
470 let mut writer = Vec::new();
471 let mut decoder = GzDecoder::new(writer);
472 let l = bytes.len() - 5;
473 let n = decoder.write(&bytes[..l]).unwrap();
474 decoder.write(&bytes[n..]).unwrap();
475 writer = decoder.finish().unwrap();
476 let return_string = String::from_utf8(writer).expect("String parsing error");
477 assert_eq!(return_string, STR);
478 }
479 }