1use 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
33pub struct Compression {
35 min_size: usize,
36 level: u32,
37}
38
39impl Compression {
40 pub fn new() -> Self {
42 Self {
43 min_size: 1024,
44 level: 6,
45 }
46 }
47
48 pub fn min_size(mut self, bytes: usize) -> Self {
50 self.min_size = bytes;
51 self
52 }
53
54 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
117fn 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}