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.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
116fn 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}