tower-rs · davidpdrsn · Feb 24, 2023 · Jul 6, 2022 · Jul 7, 2022 · Jul 7, 2022
diff --git a/tower-http/src/decompression/mod.rs b/tower-http/src/decompression/mod.rs
@@ -1,111 +1,17 @@
-//! Middleware that decompresses response bodies.
-//!
-//! # Example
-//!
-//! ```rust
-//! use bytes::BytesMut;
-//! use http::{Request, Response};
-//! use http_body::Body as _; // for Body::data
-//! use hyper::Body;
-//! use std::convert::Infallible;
-//! use tower::{Service, ServiceExt, ServiceBuilder, service_fn};
-//! use tower_http::{compression::Compression, decompression::DecompressionLayer, BoxError};
-//! #
-//! # #[tokio::main]
-//! # async fn main() -> Result<(), tower_http::BoxError> {
-//! # async fn handle(req: Request<Body>) -> Result<Response<Body>, Infallible> {
-//! #     let body = Body::from("Hello, World!");
-//! #     Ok(Response::new(body))
-//! # }
-//!
-//! // Some opaque service that applies compression.
-//! let service = Compression::new(service_fn(handle));
-//!
-//! // Our HTTP client.
-//! let mut client = ServiceBuilder::new()
-//!     // Automatically decompress response bodies.
-//!     .layer(DecompressionLayer::new())
-//!     .service(service);
-//!
-//! // Call the service.
-//! //
-//! // `DecompressionLayer` takes care of setting `Accept-Encoding`.
-//! let request = Request::new(Body::empty());
-//!
-//! let response = client
-//!     .ready()
-//!     .await?
-//!     .call(request)
-//!     .await?;
-//!
-//! // Read the body
-//! let mut body = response.into_body();
-//! let mut bytes = BytesMut::new();
-//! while let Some(chunk) = body.data().await {
-//!     let chunk = chunk?;
-//!     bytes.extend_from_slice(&chunk[..]);
-//! }
-//! let body = String::from_utf8(bytes.to_vec()).map_err(Into::<BoxError>::into)?;
-//!
-//! assert_eq!(body, "Hello, World!");
-//! #
-//! # Ok(())
-//! # }
-//! ```
+//! Middleware for decompressing Requests and Responses
 
 mod body;
-mod future;
-mod layer;
-mod service;
+pub mod request;
+mod response;
 
 pub use self::{
-    body::DecompressionBody, future::ResponseFuture, layer::DecompressionLayer,
-    service::Decompression,
+    body::DecompressionBody,
+    request::{
+        layer::RequestDecompressionLayer,
+        service::RequestDecompression
+    },
+    response::{
+        future::ResponseFuture, layer::ResponseDecompressionLayer as DecompressionLayer,
+        service::ResponseDecompression as Decompression,
+    },
 };
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::compression::Compression;
-    use bytes::BytesMut;
-    use http::Response;
-    use http_body::Body as _;
-    use hyper::{Body, Client, Error, Request};
-    use tower::{service_fn, Service, ServiceExt};
-
-    #[tokio::test]
-    async fn works() {
-        let mut client = Decompression::new(Compression::new(service_fn(handle)));
-
-        let req = Request::builder()
-            .header("accept-encoding", "gzip")
-            .body(Body::empty())
-            .unwrap();
-        let res = client.ready().await.unwrap().call(req).await.unwrap();
-
-        // read the body, it will be decompressed automatically
-        let mut body = res.into_body();
-        let mut data = BytesMut::new();
-        while let Some(chunk) = body.data().await {
-            let chunk = chunk.unwrap();
-            data.extend_from_slice(&chunk[..]);
-        }
-        let decompressed_data = String::from_utf8(data.freeze().to_vec()).unwrap();
-
-        assert_eq!(decompressed_data, "Hello, World!");
-    }
-
-    async fn handle(_req: Request<Body>) -> Result<Response<Body>, Error> {
-        Ok(Response::new(Body::from("Hello, World!")))
-    }
-
-    #[allow(dead_code)]
-    async fn is_compatible_with_hyper() {
-        let mut client = Decompression::new(Client::new());
-
-        let req = Request::new(Body::empty());
-
-        let _: Response<DecompressionBody<Body>> =
-            client.ready().await.unwrap().call(req).await.unwrap();
-    }
-}
diff --git a/tower-http/src/decompression/request/future.rs b/tower-http/src/decompression/request/future.rs
@@ -0,0 +1,93 @@
+use crate::compression_utils::AcceptEncoding;
+use crate::BoxError;
+use bytes::Buf;
+use http::{header, Response, StatusCode};
+use http_body::{combinators::UnsyncBoxBody, Body, Empty};
+use pin_project_lite::pin_project;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+pin_project! {
+    #[derive(Debug)]
+    pub struct ResponseFuture<F, B, E>
+    where
+        F: Future<Output = Result<Response<B>, E>>,
+        B: Body
+    {
+        #[pin]
+        kind: Kind<F, B, E>,
+    }
+}
+
+pin_project! {
+    #[derive(Debug)]
+    #[project = StateProj]
+    enum Kind<F, B, E>
+    where
+        F: Future<Output = Result<Response<B>, E>>,
+        B: Body
+    {
+        Inner {
+            #[pin]
+            fut: F
+        },
+        Unsupported {
+            #[pin]
+            accept: AcceptEncoding
+        },
+    }
+}
+
+impl<F, B, E> ResponseFuture<F, B, E>
+where
+    F: Future<Output = Result<Response<B>, E>>,
+    B: Body,
+{
+    #[must_use]
+    pub(super) fn unsupported_encoding(accept: AcceptEncoding) -> Self {
+        Self {
+            kind: Kind::Unsupported { accept },
+        }
+    }
+
+    #[must_use]
+    pub(super) fn inner(fut: F) -> Self {
+        Self {
+            kind: Kind::Inner { fut },
+        }
+    }
+}
+
+impl<F, B, E> Future for ResponseFuture<F, B, E>
+where
+    F: Future<Output = Result<Response<B>, E>>,
+    B: Body + Send + 'static,
+    <B as Body>::Data: Buf + 'static,
+    <B as Body>::Error: Into<BoxError> + 'static,
+    E: Into<BoxError>,
+{
+    type Output = Result<Response<UnsyncBoxBody<<B as Body>::Data, BoxError>>, BoxError>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.project().kind.project() {
+            StateProj::Inner { fut } => fut
+                .poll(cx)
+                .map_ok(|res| res.map(|body| body.map_err(Into::into).boxed_unsync()))
+                .map_err(Into::into),
+            StateProj::Unsupported { accept } => {
+                let builder = if let Some(accept) = accept.to_header_value() {
+                    Response::builder().header(header::ACCEPT_ENCODING, accept)
+                } else {
+                    Response::builder()
+                };
+                let res = builder
+                    .status(StatusCode::UNSUPPORTED_MEDIA_TYPE)
+                    .body(Empty::new().map_err(Into::into).boxed_unsync())
+                    .map_err(Into::into);
+                Poll::Ready(res)
+            }
+        }
+    }
+}
diff --git a/tower-http/src/decompression/request/layer.rs b/tower-http/src/decompression/request/layer.rs
@@ -0,0 +1,82 @@
+use super::service::RequestDecompression;
+use crate::compression_utils::AcceptEncoding;
+use tower_layer::Layer;
+
+#[derive(Debug, Default, Clone)]
+pub struct RequestDecompressionLayer {
+    accept: AcceptEncoding,
+    pass_through_unaccepted: bool,
+}
+
+impl<S> Layer<S> for RequestDecompressionLayer {
+    type Service = RequestDecompression<S>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        RequestDecompression {
+            inner: service,
+            accept: self.accept,
+            pass_through_unaccepted: self.pass_through_unaccepted,
+        }
+    }
+}
+
+impl RequestDecompressionLayer {
+    /// Creates a new `RequestDecompressionLayer`.
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    /// Sets whether to support
+    /// gzip encoding.
+    #[cfg(feature = "decompression-gzip")]
+    pub fn gzip(mut self, enable: bool) -> Self {
+        self.accept.set_gzip(enable);
+        self
+    }
+
+    /// Sets whether to support
+    /// Deflate encoding.
+    #[cfg(feature = "decompression-deflate")]
+    pub fn deflate(mut self, enable: bool) -> Self {
+        self.accept.set_deflate(enable);
+        self
+    }
+
+    /// Sets whether to support
+    /// Brotli encoding.
+    #[cfg(feature = "decompression-br")]
+    pub fn br(mut self, enable: bool) -> Self {
+        self.accept.set_br(enable);
+        self
+    }
+
+    /// Disables support for gzip encoding.
+    ///
+    /// This method is available even if the `gzip` crate feature is disabled.
+    pub fn no_gzip(mut self) -> Self {
+        self.accept.set_gzip(false);
+        self
+    }
+
+    /// Disables support for Deflate encoding.
+    ///
+    /// This method is available even if the `deflate` crate feature is disabled.
+    pub fn no_deflate(mut self) -> Self {
+        self.accept.set_deflate(false);
+        self
+    }
+
+    /// Disables support for Brotli encoding.
+    ///
+    /// This method is available even if the `br` crate feature is disabled.
+    pub fn no_br(mut self) -> Self {
+        self.accept.set_br(false);
+        self
+    }
+
+    /// Passes through unaccepted encodings
+    pub fn pass_through_unaccepted(mut self) -> Self {
+        self.pass_through_unaccepted = true;
+        self
+    }
+}