tako/plugins/
compression.rs

1#![cfg_attr(docsrs, doc(cfg(feature = "plugins")))]
2//! HTTP response compression plugin supporting multiple algorithms and streaming.
3//!
4//! This module provides comprehensive HTTP response compression functionality for Tako
5//! applications. It supports multiple compression algorithms including Gzip, Brotli, DEFLATE,
6//! and optionally Zstandard, with configurable compression levels and streaming capabilities.
7//! The plugin automatically negotiates compression based on client Accept-Encoding headers
8//! and applies compression selectively based on content type, response size, and status code.
9//!
10//! The compression plugin can be applied at both router-level (all routes) and route-level
11//! (specific routes), allowing different compression settings for different endpoints.
12//!
13//! # Examples
14//!
15//! ```rust
16//! use tako::plugins::compression::CompressionBuilder;
17//! use tako::plugins::TakoPlugin;
18//! use tako::router::Router;
19//! use tako::Method;
20//!
21//! async fn handler(_req: tako::types::Request) -> &'static str {
22//!     "Response data"
23//! }
24//!
25//! async fn api_handler(_req: tako::types::Request) -> &'static str {
26//!     "Large API response"
27//! }
28//!
29//! let mut router = Router::new();
30//!
31//! // Router-level: Basic compression setup (applied to all routes)
32//! let compression = CompressionBuilder::new()
33//!     .enable_gzip(true)
34//!     .enable_brotli(true)
35//!     .min_size(1024)
36//!     .build();
37//! router.plugin(compression);
38//!
39//! // Route-level: Advanced compression for specific API endpoint
40//! let api_route = router.route(Method::GET, "/api/large-data", api_handler);
41//! let advanced = CompressionBuilder::new()
42//!     .enable_gzip(true)
43//!     .gzip_level(9)
44//!     .enable_brotli(true)
45//!     .brotli_level(11)
46//!     .enable_stream(true)
47//!     .min_size(512)
48//!     .build();
49//! api_route.plugin(advanced);
50//! ```
51
52use anyhow::Result;
53use async_trait::async_trait;
54use bytes::Bytes;
55use flate2::{
56  Compression as GzLevel,
57  write::{DeflateEncoder, GzEncoder},
58};
59use http::{
60  HeaderValue, StatusCode,
61  header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, VARY},
62};
63use http_body_util::BodyExt;
64use std::io::{Read, Write};
65
66pub mod brotli_stream;
67pub mod deflate_stream;
68pub mod gzip_stream;
69pub mod zstd_stream;
70
71#[cfg(feature = "zstd")]
72use zstd::stream::encode_all as zstd_encode;
73
74#[cfg(feature = "zstd")]
75use crate::plugins::compression::zstd_stream::stream_zstd;
76use crate::{
77  body::TakoBody,
78  middleware::Next,
79  plugins::{
80    TakoPlugin,
81    compression::{
82      brotli_stream::stream_brotli, deflate_stream::stream_deflate, gzip_stream::stream_gzip,
83    },
84  },
85  responder::Responder,
86  router::Router,
87  types::{Request, Response},
88};
89
90/// Supported HTTP compression encoding algorithms.
91#[derive(Clone, Copy, Debug, PartialEq, Eq)]
92pub enum Encoding {
93  /// Gzip compression (RFC 1952) - widely supported, good compression ratio.
94  Gzip,
95  /// Brotli compression (RFC 7932) - excellent compression ratio, modern browsers.
96  Brotli,
97  /// DEFLATE compression (RFC 1951) - fast compression, good compatibility.
98  Deflate,
99  /// Zstandard compression - high performance, excellent ratio (requires zstd feature).
100  #[cfg(feature = "zstd")]
101  #[cfg_attr(docsrs, doc(cfg(feature = "zstd")))]
102  Zstd,
103}
104
105impl Encoding {
106  /// Returns the HTTP Content-Encoding header value for this compression algorithm.
107  fn as_str(&self) -> &'static str {
108    match self {
109      Encoding::Gzip => "gzip",
110      Encoding::Brotli => "br",
111      Encoding::Deflate => "deflate",
112      #[cfg(feature = "zstd")]
113      Encoding::Zstd => "zstd",
114    }
115  }
116}
117
118/// Configuration settings for HTTP response compression.
119#[derive(Clone)]
120pub struct Config {
121  /// List of enabled compression encodings in preference order.
122  pub enabled: Vec<Encoding>,
123  /// Minimum response size in bytes required for compression to be applied.
124  pub min_size: usize,
125  /// Gzip compression level (1-9, where 9 is maximum compression).
126  pub gzip_level: u32,
127  /// Brotli compression level (1-11, where 11 is maximum compression).
128  pub brotli_level: u32,
129  /// DEFLATE compression level (1-9, where 9 is maximum compression).
130  pub deflate_level: u32,
131  /// Zstandard compression level (1-22, where 22 is maximum compression).
132  #[cfg(feature = "zstd")]
133  pub zstd_level: i32,
134  /// Whether to use streaming compression instead of buffering entire responses.
135  pub stream: bool,
136}
137
138impl Default for Config {
139  /// Provides sensible default compression configuration.
140  fn default() -> Self {
141    Self {
142      enabled: vec![Encoding::Gzip, Encoding::Brotli, Encoding::Deflate],
143      min_size: 1024,
144      gzip_level: 5,
145      brotli_level: 5,
146      deflate_level: 5,
147      #[cfg(feature = "zstd")]
148      zstd_level: 3,
149      stream: false,
150    }
151  }
152}
153
154/// Builder for configuring HTTP response compression settings.
155///
156/// `CompressionBuilder` provides a fluent API for constructing compression plugin
157/// configurations. It allows selective enabling/disabling of compression algorithms,
158/// setting compression levels, and configuring behavior options like streaming and
159/// minimum response size thresholds.
160///
161/// # Examples
162///
163/// ```rust
164/// use tako::plugins::compression::CompressionBuilder;
165///
166/// // Basic setup with default settings
167/// let basic = CompressionBuilder::new().build();
168///
169/// // Custom configuration
170/// let custom = CompressionBuilder::new()
171///     .enable_gzip(true)
172///     .gzip_level(8)
173///     .enable_brotli(true)
174///     .brotli_level(6)
175///     .enable_deflate(false)
176///     .min_size(2048)
177///     .enable_stream(true)
178///     .build();
179/// ```
180pub struct CompressionBuilder(Config);
181
182impl CompressionBuilder {
183  /// Creates a new compression configuration builder with default settings.
184  pub fn new() -> Self {
185    Self(Config::default())
186  }
187
188  /// Enables or disables Gzip compression.
189  pub fn enable_gzip(mut self, yes: bool) -> Self {
190    if yes && !self.0.enabled.contains(&Encoding::Gzip) {
191      self.0.enabled.push(Encoding::Gzip)
192    }
193    if !yes {
194      self.0.enabled.retain(|e| *e != Encoding::Gzip)
195    }
196    self
197  }
198
199  /// Enables or disables Brotli compression.
200  pub fn enable_brotli(mut self, yes: bool) -> Self {
201    if yes && !self.0.enabled.contains(&Encoding::Brotli) {
202      self.0.enabled.push(Encoding::Brotli)
203    }
204    if !yes {
205      self.0.enabled.retain(|e| *e != Encoding::Brotli)
206    }
207    self
208  }
209
210  /// Enables or disables DEFLATE compression.
211  pub fn enable_deflate(mut self, yes: bool) -> Self {
212    if yes && !self.0.enabled.contains(&Encoding::Deflate) {
213      self.0.enabled.push(Encoding::Deflate)
214    }
215    if !yes {
216      self.0.enabled.retain(|e| *e != Encoding::Deflate)
217    }
218    self
219  }
220
221  /// Enables or disables Zstandard compression (requires zstd feature).
222  #[cfg(feature = "zstd")]
223  #[cfg_attr(docsrs, doc(cfg(feature = "zstd")))]
224  pub fn enable_zstd(mut self, yes: bool) -> Self {
225    if yes && !self.0.enabled.contains(&Encoding::Zstd) {
226      self.0.enabled.push(Encoding::Zstd)
227    }
228    if !yes {
229      self.0.enabled.retain(|e| *e != Encoding::Zstd)
230    }
231    self
232  }
233
234  /// Enables or disables streaming compression mode.
235  pub fn enable_stream(mut self, stream: bool) -> Self {
236    self.0.stream = stream;
237    self
238  }
239
240  /// Sets the minimum response size threshold for compression.
241  pub fn min_size(mut self, bytes: usize) -> Self {
242    self.0.min_size = bytes;
243    self
244  }
245
246  /// Sets the Gzip compression level (1-9).
247  pub fn gzip_level(mut self, lvl: u32) -> Self {
248    self.0.gzip_level = lvl.min(9);
249    self
250  }
251
252  /// Sets the Brotli compression level (1-11).
253  pub fn brotli_level(mut self, lvl: u32) -> Self {
254    self.0.brotli_level = lvl.min(11);
255    self
256  }
257
258  /// Sets the DEFLATE compression level (1-9).
259  pub fn deflate_level(mut self, lvl: u32) -> Self {
260    self.0.deflate_level = lvl.min(9);
261    self
262  }
263
264  /// Sets the Zstandard compression level (1-22, requires zstd feature).
265  #[cfg(feature = "zstd")]
266  #[cfg_attr(docsrs, doc(cfg(feature = "zstd")))]
267  pub fn zstd_level(mut self, lvl: i32) -> Self {
268    self.0.zstd_level = lvl.clamp(1, 22);
269    self
270  }
271
272  /// Builds the compression plugin with the configured settings.
273  pub fn build(self) -> CompressionPlugin {
274    CompressionPlugin { cfg: self.0 }
275  }
276}
277
278pub enum CompressionResponse<R>
279where
280  R: Responder,
281{
282  /// Plain, uncompressed response.
283  Plain(R),
284  /// Compressed or streaming response.
285  Stream(R),
286}
287
288impl<R> Responder for CompressionResponse<R>
289where
290  R: Responder,
291{
292  fn into_response(self) -> Response {
293    match self {
294      CompressionResponse::Plain(r) => r.into_response(),
295      CompressionResponse::Stream(r) => r.into_response(),
296    }
297  }
298}
299
300/// HTTP response compression plugin for Tako applications.
301///
302/// `CompressionPlugin` provides automatic response compression based on client
303/// Accept-Encoding headers and configurable compression algorithms. It supports
304/// multiple compression formats, streaming compression, and intelligent content
305/// type detection to optimize bandwidth usage and response times.
306///
307/// # Examples
308///
309/// ```rust
310/// use tako::plugins::compression::{CompressionPlugin, CompressionBuilder};
311/// use tako::plugins::TakoPlugin;
312/// use tako::router::Router;
313///
314/// // Use default settings
315/// let compression = CompressionPlugin::default();
316/// let mut router = Router::new();
317/// router.plugin(compression);
318///
319/// // Custom configuration
320/// let custom = CompressionBuilder::new()
321///     .enable_gzip(true)
322///     .enable_brotli(true)
323///     .min_size(2048)
324///     .build();
325/// router.plugin(custom);
326/// ```
327#[derive(Clone)]
328#[doc(alias = "compression")]
329#[doc(alias = "gzip")]
330#[doc(alias = "brotli")]
331#[doc(alias = "deflate")]
332pub struct CompressionPlugin {
333  cfg: Config,
334}
335
336impl Default for CompressionPlugin {
337  /// Creates a compression plugin with default configuration settings.
338  fn default() -> Self {
339    Self {
340      cfg: Config::default(),
341    }
342  }
343}
344
345#[async_trait]
346impl TakoPlugin for CompressionPlugin {
347  /// Returns the plugin name for identification and debugging.
348  fn name(&self) -> &'static str {
349    "CompressionPlugin"
350  }
351
352  /// Sets up the compression plugin by registering middleware with the router.
353  fn setup(&self, router: &Router) -> Result<()> {
354    let cfg = self.cfg.clone();
355    router.middleware(move |req, next| {
356      let cfg = cfg.clone();
357      let stream = cfg.stream.clone();
358      async move {
359        if stream == false {
360          return CompressionResponse::Plain(
361            compress_middleware(req, next, cfg).await.into_response(),
362          );
363        } else {
364          return CompressionResponse::Stream(
365            compress_stream_middleware(req, next, cfg)
366              .await
367              .into_response(),
368          );
369        }
370      }
371    });
372    Ok(())
373  }
374}
375
376/// Middleware function for buffered response compression.
377///
378/// This middleware compresses entire response bodies in memory before sending them
379/// to clients. It's more memory-intensive than streaming compression but may have
380/// better compression ratios for smaller responses.
381async fn compress_middleware(req: Request, next: Next, cfg: Config) -> impl Responder {
382  // Parse the `Accept-Encoding` header to determine supported encodings.
383  let accepted = req
384    .headers()
385    .get(ACCEPT_ENCODING)
386    .and_then(|v| v.to_str().ok())
387    .unwrap_or("")
388    .to_ascii_lowercase();
389
390  // Process the request and get the response.
391  let mut resp = next.run(req).await;
392  let chosen = choose_encoding(&accepted, &cfg.enabled);
393
394  // Skip compression for non-successful responses or if already encoded.
395  let status = resp.status();
396  if !(status.is_success() || status == StatusCode::NOT_MODIFIED) {
397    return resp.into_response();
398  }
399
400  if resp.headers().contains_key(CONTENT_ENCODING) {
401    return resp.into_response();
402  }
403
404  // Skip compression for unsupported content types.
405  if let Some(ct) = resp.headers().get(CONTENT_TYPE) {
406    let ct = ct.to_str().unwrap_or("");
407    if !(ct.starts_with("text/")
408      || ct.contains("json")
409      || ct.contains("javascript")
410      || ct.contains("xml"))
411    {
412      return resp.into_response();
413    }
414  }
415
416  // Collect the response body and check its size.
417  let body_bytes = resp.body_mut().collect().await.unwrap().to_bytes();
418  if body_bytes.len() < cfg.min_size {
419    *resp.body_mut() = TakoBody::from(Bytes::from(body_bytes));
420    return resp.into_response();
421  }
422
423  // Compress the response body if a suitable encoding is chosen.
424  if let Some(enc) = chosen {
425    let compressed = match enc {
426      Encoding::Gzip => {
427        compress_gzip(&body_bytes, cfg.gzip_level).unwrap_or_else(|_| body_bytes.to_vec())
428      }
429      Encoding::Brotli => {
430        compress_brotli(&body_bytes, cfg.brotli_level).unwrap_or_else(|_| body_bytes.to_vec())
431      }
432      Encoding::Deflate => {
433        compress_deflate(&body_bytes, cfg.deflate_level).unwrap_or_else(|_| body_bytes.to_vec())
434      }
435      #[cfg(feature = "zstd")]
436      Encoding::Zstd => {
437        compress_zstd(&body_bytes, cfg.zstd_level).unwrap_or_else(|_| body_bytes.to_vec())
438      }
439    };
440    *resp.body_mut() = TakoBody::from(Bytes::from(compressed));
441    resp
442      .headers_mut()
443      .insert(CONTENT_ENCODING, HeaderValue::from_static(enc.as_str()));
444    resp.headers_mut().remove(CONTENT_LENGTH);
445    resp
446      .headers_mut()
447      .insert(VARY, HeaderValue::from_static("Accept-Encoding"));
448  } else {
449    *resp.body_mut() = TakoBody::from(Bytes::from(body_bytes));
450  }
451
452  resp.into_response()
453}
454
455/// Middleware function for streaming response compression.
456///
457/// This middleware compresses response bodies on-the-fly as they stream to clients.
458/// It's more memory-efficient than buffered compression but requires compatible
459/// response body types that support streaming.
460pub async fn compress_stream_middleware(req: Request, next: Next, cfg: Config) -> impl Responder {
461  // Parse the `Accept-Encoding` header to determine supported encodings.
462  let accepted = req
463    .headers()
464    .get(ACCEPT_ENCODING)
465    .and_then(|v| v.to_str().ok())
466    .unwrap_or("")
467    .to_ascii_lowercase();
468
469  // Process the request and get the response.
470  let mut resp = next.run(req).await;
471  let chosen = choose_encoding(&accepted, &cfg.enabled);
472
473  // Skip compression for non-successful responses or if already encoded.
474  let status = resp.status();
475  if !(status.is_success() || status == StatusCode::NOT_MODIFIED) {
476    return resp.into_response();
477  }
478
479  if resp.headers().contains_key(CONTENT_ENCODING) {
480    return resp.into_response();
481  }
482
483  // Skip compression for unsupported content types.
484  if let Some(ct) = resp.headers().get(CONTENT_TYPE) {
485    let ct = ct.to_str().unwrap_or("");
486    if !(ct.starts_with("text/")
487      || ct.contains("json")
488      || ct.contains("javascript")
489      || ct.contains("xml"))
490    {
491      return resp.into_response();
492    }
493  }
494
495  // Estimate size from `Content-Length`.
496  if let Some(len) = resp
497    .headers()
498    .get(CONTENT_LENGTH)
499    .and_then(|v| v.to_str().ok())
500    .and_then(|v| v.parse::<usize>().ok())
501  {
502    if len < cfg.min_size {
503      return resp.into_response();
504    }
505  }
506
507  if let Some(enc) = chosen {
508    let body = std::mem::replace(resp.body_mut(), TakoBody::empty());
509    let new_body = match enc {
510      Encoding::Gzip => stream_gzip(body, cfg.gzip_level),
511      Encoding::Brotli => stream_brotli(body, cfg.brotli_level),
512      Encoding::Deflate => stream_deflate(body, cfg.deflate_level),
513      #[cfg(feature = "zstd")]
514      Encoding::Zstd => stream_zstd(body, cfg.zstd_level),
515    };
516    *resp.body_mut() = new_body;
517    resp
518      .headers_mut()
519      .insert(CONTENT_ENCODING, HeaderValue::from_static(enc.as_str()));
520    resp.headers_mut().remove(CONTENT_LENGTH);
521    resp
522      .headers_mut()
523      .insert(VARY, HeaderValue::from_static("Accept-Encoding"));
524  }
525
526  resp.into_response()
527}
528
529/// Selects the best compression encoding based on client preferences and server capabilities.
530///
531/// This function parses the Accept-Encoding header and chooses the most preferred
532/// compression algorithm that is both supported by the client and enabled on the server.
533/// The selection prioritizes compression quality while respecting client preferences.
534fn choose_encoding(header: &str, enabled: &[Encoding]) -> Option<Encoding> {
535  let header = header.to_ascii_lowercase();
536  let test = |e: Encoding| header.contains(e.as_str()) && enabled.contains(&e);
537  if test(Encoding::Brotli) {
538    Some(Encoding::Brotli)
539  } else if test(Encoding::Gzip) {
540    Some(Encoding::Gzip)
541  } else if test(Encoding::Deflate) {
542    Some(Encoding::Deflate)
543  } else {
544    #[cfg(feature = "zstd")]
545    {
546      if test(Encoding::Zstd) {
547        return Some(Encoding::Zstd);
548      }
549    }
550    None
551  }
552}
553
554/// Compresses data using Gzip algorithm.
555fn compress_gzip(data: &[u8], lvl: u32) -> std::io::Result<Vec<u8>> {
556  let mut enc = GzEncoder::new(Vec::new(), GzLevel::new(lvl));
557  enc.write_all(data)?;
558  enc.finish()
559}
560
561/// Compresses data using Brotli algorithm.
562fn compress_brotli(data: &[u8], lvl: u32) -> std::io::Result<Vec<u8>> {
563  let mut out = Vec::new();
564  brotli::CompressorReader::new(data, 4096, lvl, 22)
565    .read_to_end(&mut out)
566    .map_err(|_| std::io::Error::new(std::io::ErrorKind::Other, "Failed to compress data"))?;
567  Ok(out)
568}
569
570/// Compresses data using DEFLATE algorithm.
571fn compress_deflate(data: &[u8], lvl: u32) -> std::io::Result<Vec<u8>> {
572  let mut enc = DeflateEncoder::new(Vec::new(), flate2::Compression::new(lvl));
573  enc.write_all(data)?;
574  enc.finish()
575}
576
577/// Compresses data using Zstandard algorithm (requires zstd feature).
578#[cfg(feature = "zstd")]
579#[cfg_attr(docsrs, doc(cfg(feature = "zstd")))]
580fn compress_zstd(data: &[u8], lvl: i32) -> std::io::Result<Vec<u8>> {
581  zstd_encode(data, lvl)
582}
tako/plugins/compression.rs

tako/plugins/
compression.rs
