use super::{Compression, Predicate}; use crate::compression::predicate::DefaultPredicate; use crate::compression::CompressionLevel; use crate::compression_utils::AcceptEncoding; use tower_layer::Layer; /// Compress response bodies of the underlying service. /// /// This uses the `Accept-Encoding` header to pick an appropriate encoding and adds the /// `Content-Encoding` header to responses. /// /// See the [module docs](crate::compression) for more details. #[derive(Clone, Debug, Default)] pub struct CompressionLayer

{ accept: AcceptEncoding, predicate: P, quality: CompressionLevel, } impl Layer for CompressionLayer

where P: Predicate, { type Service = Compression; fn layer(&self, inner: S) -> Self::Service { Compression { inner, accept: self.accept, predicate: self.predicate.clone(), quality: self.quality, } } } impl CompressionLayer { /// Creates a new [`CompressionLayer`]. pub fn new() -> Self { Self::default() } /// Sets whether to enable the gzip encoding. #[cfg(feature = "compression-gzip")] pub fn gzip(mut self, enable: bool) -> Self { self.accept.set_gzip(enable); self } /// Sets whether to enable the Deflate encoding. #[cfg(feature = "compression-deflate")] pub fn deflate(mut self, enable: bool) -> Self { self.accept.set_deflate(enable); self } /// Sets whether to enable the Brotli encoding. #[cfg(feature = "compression-br")] pub fn br(mut self, enable: bool) -> Self { self.accept.set_br(enable); self } /// Sets whether to enable the Zstd encoding. #[cfg(feature = "compression-zstd")] pub fn zstd(mut self, enable: bool) -> Self { self.accept.set_zstd(enable); self } /// Sets the compression quality. pub fn quality(mut self, quality: CompressionLevel) -> Self { self.quality = quality; self } /// Disables the 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 the 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 the 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 } /// Disables the Zstd encoding. /// /// This method is available even if the `zstd` crate feature is disabled. pub fn no_zstd(mut self) -> Self { self.accept.set_zstd(false); self } /// Replace the current compression predicate. /// /// See [`Compression::compress_when`] for more details. pub fn compress_when(self, predicate: C) -> CompressionLayer where C: Predicate, { CompressionLayer { accept: self.accept, predicate, quality: self.quality, } } } #[cfg(test)] mod tests { use super::*; use crate::test_helpers::Body; use http::{header::ACCEPT_ENCODING, Request, Response}; use http_body_util::BodyExt; use std::convert::Infallible; use tokio::fs::File; use tokio_util::io::ReaderStream; use tower::{Service, ServiceBuilder, ServiceExt}; async fn handle(_req: Request) -> Result, Infallible> { // Open the file. let file = File::open("Cargo.toml").await.expect("file missing"); // Convert the file into a `Stream`. let stream = ReaderStream::new(file); // Convert the `Stream` into a `Body`. let body = Body::from_stream(stream); // Create response. Ok(Response::new(body)) } #[tokio::test] async fn accept_encoding_configuration_works() -> Result<(), crate::BoxError> { let deflate_only_layer = CompressionLayer::new() .quality(CompressionLevel::Best) .no_br() .no_gzip(); let mut service = ServiceBuilder::new() // Compress responses based on the `Accept-Encoding` header. .layer(deflate_only_layer) .service_fn(handle); // Call the service with the deflate only layer let request = Request::builder() .header(ACCEPT_ENCODING, "gzip, deflate, br") .body(Body::empty())?; let response = service.ready().await?.call(request).await?; assert_eq!(response.headers()["content-encoding"], "deflate"); // Read the body let body = response.into_body(); let bytes = body.collect().await.unwrap().to_bytes(); let deflate_bytes_len = bytes.len(); let br_only_layer = CompressionLayer::new() .quality(CompressionLevel::Best) .no_gzip() .no_deflate(); let mut service = ServiceBuilder::new() // Compress responses based on the `Accept-Encoding` header. .layer(br_only_layer) .service_fn(handle); // Call the service with the br only layer let request = Request::builder() .header(ACCEPT_ENCODING, "gzip, deflate, br") .body(Body::empty())?; let response = service.ready().await?.call(request).await?; assert_eq!(response.headers()["content-encoding"], "br"); // Read the body let body = response.into_body(); let bytes = body.collect().await.unwrap().to_bytes(); let br_byte_length = bytes.len(); // check the corresponding algorithms are actually used // br should compresses better than deflate assert!(br_byte_length < deflate_bytes_len * 9 / 10); Ok(()) } /// Test ensuring that zstd compression will not exceed an 8MiB window size; browsers do not /// accept responses using 16MiB+ window sizes. #[tokio::test] async fn zstd_is_web_safe() -> Result<(), crate::BoxError> { async fn zeroes(_req: Request) -> Result, Infallible> { Ok(Response::new(Body::from(vec![0u8; 18_874_368]))) } // zstd will (I believe) lower its window size if a larger one isn't beneficial and // it knows the size of the input; use an 18MiB body to ensure it would want a // >=16MiB window (though it might not be able to see the input size here). let zstd_layer = CompressionLayer::new() .quality(CompressionLevel::Best) .no_br() .no_deflate() .no_gzip(); let mut service = ServiceBuilder::new().layer(zstd_layer).service_fn(zeroes); let request = Request::builder() .header(ACCEPT_ENCODING, "zstd") .body(Body::empty())?; let response = service.ready().await?.call(request).await?; assert_eq!(response.headers()["content-encoding"], "zstd"); let body = response.into_body(); let bytes = body.collect().await?.to_bytes(); let mut dec = zstd::Decoder::new(&*bytes)?; dec.window_log_max(23)?; // Limit window size accepted by decoder to 2 ^ 23 bytes (8MiB) std::io::copy(&mut dec, &mut std::io::sink())?; Ok(()) } }