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.is_streaming()
79            || res.status.is_informational()
80            || res.status == http::StatusCode::NO_CONTENT
81            || res.status == http::StatusCode::NOT_MODIFIED
82            || res.headers.contains_key(http::header::CONTENT_ENCODING)
83            || !is_compressible(&res)
84        {
85            return Ok(res);
86        }
87
88        add_vary(&mut res);
89        if !wants_gzip || res.body.len() < self.min_size {
90            return Ok(res);
91        }
92
93        let mut encoder = GzEncoder::new(
94            Vec::with_capacity(res.body.len() / 2),
95            Level::new(self.level),
96        );
97        if encoder.write_all(&res.body).is_err() {
98            return Ok(res);
99        }
100        let Ok(compressed) = encoder.finish() else {
101            return Ok(res);
102        };
103        if compressed.len() >= res.body.len() {
104            return Ok(res);
105        }
106
107        res.body = Bytes::from(compressed);
108        res.headers.insert(
109            http::header::CONTENT_ENCODING,
110            http::HeaderValue::from_static("gzip"),
111        );
112        res.headers.remove(http::header::CONTENT_LENGTH);
113        Ok(res)
114    }
115}
116
117/// `gzip` (or `*`) listed with a non-zero q-value
118fn accepts_gzip(req: &Request) -> bool {
119    let Some(value) = req
120        .headers
121        .get(http::header::ACCEPT_ENCODING)
122        .and_then(|v| v.to_str().ok())
123    else {
124        return false;
125    };
126    value.split(',').any(|item| {
127        let mut parts = item.split(';');
128        let coding = parts.next().unwrap_or("").trim();
129        if !(coding.eq_ignore_ascii_case("gzip") || coding == "*") {
130            return false;
131        }
132        parts
133            .filter_map(|p| p.trim().strip_prefix("q="))
134            .next()
135            .and_then(|q| q.trim().parse::<f32>().ok())
136            .map(|q| q > 0.0)
137            .unwrap_or(true)
138    })
139}
140
141fn is_compressible(res: &Response) -> bool {
142    let Some(ct) = res
143        .headers
144        .get(http::header::CONTENT_TYPE)
145        .and_then(|v| v.to_str().ok())
146    else {
147        return false;
148    };
149    let ct = ct
150        .split(';')
151        .next()
152        .unwrap_or("")
153        .trim()
154        .to_ascii_lowercase();
155    ct.starts_with("text/")
156        || ct.ends_with("json")
157        || ct.ends_with("javascript")
158        || ct.ends_with("xml")
159        || ct == "image/svg+xml"
160}
161
162fn add_vary(res: &mut Response) {
163    let existing = res
164        .headers
165        .get_all(http::header::VARY)
166        .iter()
167        .filter_map(|v| v.to_str().ok())
168        .collect::<Vec<_>>()
169        .join(", ");
170    if existing.split(',').any(|v| {
171        let v = v.trim();
172        v == "*" || v.eq_ignore_ascii_case("accept-encoding")
173    }) {
174        return;
175    }
176    let value = if existing.is_empty() {
177        "Accept-Encoding".to_string()
178    } else {
179        format!("{}, Accept-Encoding", existing)
180    };
181    if let Ok(v) = http::HeaderValue::from_str(&value) {
182        res.headers.insert(http::header::VARY, v);
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    fn req_with(accept: Option<&str>) -> Request {
191        let mut headers = http::HeaderMap::new();
192        if let Some(a) = accept {
193            headers.insert(http::header::ACCEPT_ENCODING, a.parse().unwrap());
194        }
195        Request::new(
196            http::Method::GET,
197            "/".parse().unwrap(),
198            headers,
199            Bytes::new(),
200        )
201    }
202
203    #[test]
204    fn accept_encoding_parsing() {
205        assert!(accepts_gzip(&req_with(Some("gzip"))));
206        assert!(accepts_gzip(&req_with(Some("br, GZIP;q=0.5"))));
207        assert!(accepts_gzip(&req_with(Some("*"))));
208        assert!(!accepts_gzip(&req_with(Some("gzip;q=0"))));
209        assert!(!accepts_gzip(&req_with(Some("br, deflate"))));
210        assert!(!accepts_gzip(&req_with(None)));
211    }
212
213    #[test]
214    fn vary_is_merged_not_duplicated() {
215        let mut res = Response::text("x");
216        add_vary(&mut res);
217        assert_eq!(res.headers["vary"], "Accept-Encoding");
218        add_vary(&mut res);
219        assert_eq!(res.headers.get_all("vary").iter().count(), 1);
220
221        let mut res = Response::text("x").with_header("vary", "Origin");
222        add_vary(&mut res);
223        assert_eq!(res.headers["vary"], "Origin, Accept-Encoding");
224    }
225
226    #[test]
227    fn compressible_types() {
228        let t = |ct: &str| is_compressible(&Response::ok().with_header("content-type", ct));
229        assert!(t("text/html; charset=utf-8") && t("application/json") && t("image/svg+xml"));
230        assert!(t("application/javascript") && t("application/ld+json") && t("application/xml"));
231        assert!(!t("image/png") && !t("application/zip") && !t("application/octet-stream"));
232        assert!(!is_compressible(&Response::ok()));
233    }
234}