Merge pull request #10 from fairingrey/decompression

Decompressed Stream Support

Author: Nemo157

src/stream/brotli.rs

@@ -2,9 +2,9 @@ use core::{
     pin::Pin,
     task::{Context, Poll},
 };
-use std::io::Result;
+use std::io::{Error, ErrorKind, Result};
 
-use brotli2::raw::{CoStatus, CompressOp};
+use brotli2::raw::{CoStatus, CompressOp, DeStatus, Decompress};
 pub use brotli2::{raw::Compress, CompressParams};
 use bytes::{Bytes, BytesMut};
 use futures::{ready, stream::Stream};
@@ -72,3 +72,67 @@ impl<S: Stream<Item = Result<Bytes>>> BrotliStream<S> {
         }
     }
 }
+
+#[unsafe_project(Unpin)]
+pub struct DecompressedBrotliStream<S: Stream<Item = Result<Bytes>>> {
+    #[pin]
+    inner: S,
+    flush: bool,
+    decompress: Decompress,
+}
+
+impl<S: Stream<Item = Result<Bytes>>> Stream for DecompressedBrotliStream<S> {
+    type Item = Result<Bytes>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes>>> {
+        const OUTPUT_BUFFER_SIZE: usize = 8_000;
+
+        let this = self.project();
+
+        if *this.flush {
+            return Poll::Ready(None);
+        }
+
+        let input_buffer = if let Some(bytes) = ready!(this.inner.poll_next(cx)) {
+            bytes?
+        } else {
+            *this.flush = true;
+            Bytes::new()
+        };
+
+        let mut decompressed_output = BytesMut::with_capacity(OUTPUT_BUFFER_SIZE);
+        let input_ref = &mut input_buffer.as_ref();
+        let output_ref = &mut &mut [][..];
+        loop {
+            let status = this.decompress.decompress(input_ref, output_ref)?;
+            while let Some(buf) = this.decompress.take_output(None) {
+                decompressed_output.extend_from_slice(buf);
+            }
+            match status {
+                DeStatus::Finished => break,
+                DeStatus::NeedInput => {
+                    if *this.flush {
+                        return Poll::Ready(Some(Err(Error::new(
+                            ErrorKind::UnexpectedEof,
+                            "reached unexpected EOF",
+                        ))));
+                    }
+                    break;
+                }
+                DeStatus::NeedOutput => (),
+            }
+        }
+
+        Poll::Ready(Some(Ok(decompressed_output.freeze())))
+    }
+}
+
+impl<S: Stream<Item = Result<Bytes>>> DecompressedBrotliStream<S> {
+    pub fn new(stream: S) -> DecompressedBrotliStream<S> {
+        DecompressedBrotliStream {
+            inner: stream,
+            flush: false,
+            decompress: Decompress::new(),
+        }
+    }
+}

src/stream/deflate.rs

@@ -5,10 +5,10 @@ use core::{
 use std::io::Result;
 use std::marker::Unpin;
 
-use super::flate::CompressedStream;
+use super::flate::{CompressedStream, DecompressedStream};
 use bytes::Bytes;
-use flate2::Compress;
 pub use flate2::Compression;
+use flate2::{Compress, Decompress};
 use futures::{stream::Stream, stream::StreamExt};
 
 pub struct DeflateStream<S: Stream<Item = Result<Bytes>> + Unpin> {
@@ -30,3 +30,23 @@ impl<S: Stream<Item = Result<Bytes>> + Unpin> DeflateStream<S> {
         }
     }
 }
+
+pub struct DecompressedDeflateStream<S: Stream<Item = Result<Bytes>> + Unpin> {
+    inner: DecompressedStream<S>,
+}
+
+impl<S: Stream<Item = Result<Bytes>> + Unpin> Stream for DecompressedDeflateStream<S> {
+    type Item = Result<Bytes>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes>>> {
+        self.inner.poll_next_unpin(cx)
+    }
+}
+
+impl<S: Stream<Item = Result<Bytes>> + Unpin> DecompressedDeflateStream<S> {
+    pub fn new(stream: S) -> DecompressedDeflateStream<S> {
+        DecompressedDeflateStream {
+            inner: DecompressedStream::new(stream, Decompress::new(false)),
+        }
+    }
+}

src/stream/flate.rs

@@ -6,8 +6,8 @@ use std::{
 };
 
 use bytes::{Bytes, BytesMut};
-pub(crate) use flate2::Compress;
-use flate2::{FlushCompress, Status};
+pub(crate) use flate2::{Compress, Decompress};
+use flate2::{FlushCompress, FlushDecompress, Status};
 use futures::{ready, stream::Stream};
 use pin_project::unsafe_project;
 
@@ -125,3 +125,109 @@ impl<S: Stream<Item = Result<Bytes>>> CompressedStream<S> {
         }
     }
 }
+
+#[unsafe_project(Unpin)]
+pub(crate) struct DecompressedStream<S: Stream<Item = Result<Bytes>>> {
+    #[pin]
+    inner: S,
+    state: State,
+    output: BytesMut,
+    decompress: Decompress,
+}
+
+impl<S: Stream<Item = Result<Bytes>>> Stream for DecompressedStream<S> {
+    type Item = Result<Bytes>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes>>> {
+        let mut this = self.project();
+
+        fn decompress(
+            decompress: &mut Decompress,
+            input: &mut Bytes,
+            output: &mut BytesMut,
+            flush: FlushDecompress,
+        ) -> Result<(Status, Bytes)> {
+            const OUTPUT_BUFFER_SIZE: usize = 8_000;
+
+            if output.len() < OUTPUT_BUFFER_SIZE {
+                output.resize(OUTPUT_BUFFER_SIZE, 0);
+            }
+
+            let (prior_in, prior_out) = (decompress.total_in(), decompress.total_out());
+            let status = decompress.decompress(input, output, flush)?;
+            let input_len = decompress.total_in() - prior_in;
+            let output_len = decompress.total_out() - prior_out;
+
+            input.advance(input_len as usize);
+            Ok((status, output.split_to(output_len as usize).freeze()))
+        }
+
+        #[allow(clippy::never_loop)] // https://github.com/rust-lang/rust-clippy/issues/4058
+        loop {
+            break match mem::replace(this.state, State::Invalid) {
+                State::Reading => {
+                    *this.state = State::Reading;
+                    *this.state = match ready!(this.inner.as_mut().poll_next(cx)) {
+                        Some(chunk) => State::Writing(chunk?),
+                        None => State::Flushing,
+                    };
+                    continue;
+                }
+
+                State::Writing(mut input) => {
+                    if input.is_empty() {
+                        *this.state = State::Reading;
+                        continue;
+                    }
+
+                    let (status, chunk) = decompress(
+                        &mut this.decompress,
+                        &mut input,
+                        &mut this.output,
+                        FlushDecompress::None,
+                    )?;
+
+                    *this.state = match status {
+                        Status::Ok => State::Writing(input),
+                        Status::StreamEnd => State::Reading,
+                        Status::BufError => panic!("unexpected BufError"),
+                    };
+
+                    Poll::Ready(Some(Ok(chunk)))
+                }
+
+                State::Flushing => {
+                    let (status, chunk) = decompress(
+                        &mut this.decompress,
+                        &mut Bytes::new(),
+                        &mut this.output,
+                        FlushDecompress::Finish,
+                    )?;
+
+                    *this.state = match status {
+                        Status::Ok => State::Flushing,
+                        Status::StreamEnd => State::Done,
+                        Status::BufError => panic!("unexpected BufError"),
+                    };
+
+                    Poll::Ready(Some(Ok(chunk)))
+                }
+
+                State::Done => Poll::Ready(None),
+
+                State::Invalid => panic!("CompressedStream reached invalid state"),
+            };
+        }
+    }
+}
+
+impl<S: Stream<Item = Result<Bytes>>> DecompressedStream<S> {
+    pub(crate) fn new(stream: S, decompress: Decompress) -> DecompressedStream<S> {
+        DecompressedStream {
+            inner: stream,
+            state: State::Reading,
+            output: BytesMut::new(),
+            decompress,
+        }
+    }
+}