Skip to main content

ullun/
compression.rs

1//! Gzip response compression middleware (feature `compression`)
2//!
3//! Compresses text-like responses (`text/*`, JSON, JavaScript, XML, SVG) of at
4//! least [`Compression::min_size`] bytes when the client sends
5//! `Accept-Encoding: gzip`. Already encoded responses, `HEAD` requests and
6//! bodiless statuses are left alone. `Vary: Accept-Encoding` is added whenever a
7//! response could have been compressed, so caches keep both variants apart.
8//!
9//! ```no_run
10//! use ullun::compression::Compression;
11//! use ullun::prelude::*;
12//!
13//! # async fn run() {
14//! App::new()
15//!     .middleware(Compression::new())
16//!     .get("/", |_req: Request| async { Ok(Response::text("hello ".repeat(500))) })
17//!     .run("127.0.0.1:3000")
18//!     .await
19//!     .unwrap();
20//! # }
21//! ```
22
23use crate::error::Result;
24use crate::middleware::{Middleware, Next};
25use crate::request::Request;
26use crate::response::Response;
27use async_trait::async_trait;
28use bytes::Bytes;
29use flate2::write::GzEncoder;
30use flate2::Compression as Level;
31use std::io::Write;
32
33/// Gzip compression middleware
34pub struct Compression {
35    min_size: usize,
36    level: u32,
37}
38
39impl Compression {
40    /// Defaults: bodies of 1 KiB and more, compression level 6
41    pub fn new() -> Self {
42        Self {
43            min_size: 1024,
44            level: 6,
45        }
46    }
47
48    /// Smallest body (bytes) worth compressing
49    pub fn min_size(mut self, bytes: usize) -> Self {
50        self.min_size = bytes;
51        self
52    }
53
54    /// gzip level 0 (none) to 9 (best); values above 9 are clamped
55    pub fn level(mut self, level: u32) -> Self {
56        self.level = level.min(9);
57        self
58    }
59}
60
61impl Default for Compression {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67#[async_trait]
68impl Middleware for Compression {
69    async fn handle(&self, req: Request, next: Next) -> Result<Response> {
70        let wants_gzip = accepts_gzip(&req);
71        let is_head = req.method == http::Method::HEAD;
72        let mut res = match next.run(req).await {
73            Ok(res) => res,
74            Err(err) => Response::from_error(err),
75        };
76
77        if is_head
78            || res.status.is_informational()
79            || res.status == http::StatusCode::NO_CONTENT
80            || res.status == http::StatusCode::NOT_MODIFIED
81            || res.headers.contains_key(http::header::CONTENT_ENCODING)
82            || !is_compressible(&res)
83        {
84            return Ok(res);
85        }
86
87        add_vary(&mut res);
88        if !wants_gzip || res.body.len() < self.min_size {
89            return Ok(res);
90        }
91
92        let mut encoder = GzEncoder::new(
93            Vec::with_capacity(res.body.len() / 2),
94            Level::new(self.level),
95        );
96        if encoder.write_all(&res.body).is_err() {
97            return Ok(res);
98        }
99        let Ok(compressed) = encoder.finish() else {
100            return Ok(res);
101        };
102        if compressed.len() >= res.body.len() {
103            return Ok(res);
104        }
105
106        res.body = Bytes::from(compressed);
107        res.headers.insert(
108            http::header::CONTENT_ENCODING,
109            http::HeaderValue::from_static("gzip"),
110        );
111        res.headers.remove(http::header::CONTENT_LENGTH);
112        Ok(res)
113    }
114}
115
116/// `gzip` (or `*`) listed with a non-zero q-value
117fn accepts_gzip(req: &Request) -> bool {
118    let Some(value) = req
119        .headers
120        .get(http::header::ACCEPT_ENCODING)
121        .and_then(|v| v.to_str().ok())
122    else {
123        return false;
124    };
125    value.split(',').any(|item| {
126        let mut parts = item.split(';');
127        let coding = parts.next().unwrap_or("").trim();
128        if !(coding.eq_ignore_ascii_case("gzip") || coding == "*") {
129            return false;
130        }
131        parts
132            .filter_map(|p| p.trim().strip_prefix("q="))
133            .next()
134            .and_then(|q| q.trim().parse::<f32>().ok())
135            .map(|q| q > 0.0)
136            .unwrap_or(true)
137    })
138}
139
140fn is_compressible(res: &Response) -> bool {
141    let Some(ct) = res
142        .headers
143        .get(http::header::CONTENT_TYPE)
144        .and_then(|v| v.to_str().ok())
145    else {
146        return false;
147    };
148    let ct = ct
149        .split(';')
150        .next()
151        .unwrap_or("")
152        .trim()
153        .to_ascii_lowercase();
154    ct.starts_with("text/")
155        || ct.ends_with("json")
156        || ct.ends_with("javascript")
157        || ct.ends_with("xml")
158        || ct == "image/svg+xml"
159}
160
161fn add_vary(res: &mut Response) {
162    let existing = res
163        .headers
164        .get_all(http::header::VARY)
165        .iter()
166        .filter_map(|v| v.to_str().ok())
167        .collect::<Vec<_>>()
168        .join(", ");
169    if existing.split(',').any(|v| {
170        let v = v.trim();
171        v == "*" || v.eq_ignore_ascii_case("accept-encoding")
172    }) {
173        return;
174    }
175    let value = if existing.is_empty() {
176        "Accept-Encoding".to_string()
177    } else {
178        format!("{}, Accept-Encoding", existing)
179    };
180    if let Ok(v) = http::HeaderValue::from_str(&value) {
181        res.headers.insert(http::header::VARY, v);
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    fn req_with(accept: Option<&str>) -> Request {
190        let mut headers = http::HeaderMap::new();
191        if let Some(a) = accept {
192            headers.insert(http::header::ACCEPT_ENCODING, a.parse().unwrap());
193        }
194        Request::new(
195            http::Method::GET,
196            "/".parse().unwrap(),
197            headers,
198            Bytes::new(),
199        )
200    }
201
202    #[test]
203    fn accept_encoding_parsing() {
204        assert!(accepts_gzip(&req_with(Some("gzip"))));
205        assert!(accepts_gzip(&req_with(Some("br, GZIP;q=0.5"))));
206        assert!(accepts_gzip(&req_with(Some("*"))));
207        assert!(!accepts_gzip(&req_with(Some("gzip;q=0"))));
208        assert!(!accepts_gzip(&req_with(Some("br, deflate"))));
209        assert!(!accepts_gzip(&req_with(None)));
210    }
211
212    #[test]
213    fn vary_is_merged_not_duplicated() {
214        let mut res = Response::text("x");
215        add_vary(&mut res);
216        assert_eq!(res.headers["vary"], "Accept-Encoding");
217        add_vary(&mut res);
218        assert_eq!(res.headers.get_all("vary").iter().count(), 1);
219
220        let mut res = Response::text("x").with_header("vary", "Origin");
221        add_vary(&mut res);
222        assert_eq!(res.headers["vary"], "Origin, Accept-Encoding");
223    }
224
225    #[test]
226    fn compressible_types() {
227        let t = |ct: &str| is_compressible(&Response::ok().with_header("content-type", ct));
228        assert!(t("text/html; charset=utf-8") && t("application/json") && t("image/svg+xml"));
229        assert!(t("application/javascript") && t("application/ld+json") && t("application/xml"));
230        assert!(!t("image/png") && !t("application/zip") && !t("application/octet-stream"));
231        assert!(!is_compressible(&Response::ok()));
232    }
233}