1#![cfg_attr(feature = "unstable", feature(test))]
2extern crate iron;
6extern crate libflate;
7extern crate brotli;
8
9use std::io;
10use std::io::Write;
11use iron::prelude::*;
12use iron::headers::*;
13use iron::{AfterMiddleware};
14
15use iron::headers::Encoding;
16use iron::response::WriteBody;
17
18const DEFAULT_MIN_BYTES_FOR_COMPRESSION: u64 = 860;
19
20#[derive(PartialEq, Clone, Debug)]
21enum CompressionEncoding {
22 Brotli,
23 Deflate,
24 Gzip,
25}
26
27struct BrotliBody(Box<WriteBody>);
28
29impl WriteBody for BrotliBody {
30 fn write_body(&mut self, w: &mut Write) -> io::Result<()> {
31 const BUFFER_SIZE: usize = 4096;
32 const QUALITY: u32 = 8;
33 const LG_WINDOW_SIZE: u32 = 20;
34 let mut encoder = brotli::CompressorWriter::new(w, BUFFER_SIZE, QUALITY, LG_WINDOW_SIZE);
35 self.0.write_body(&mut encoder)?;
36 Ok(())
37 }
38}
39
40struct GzipBody(Box<WriteBody>);
41
42impl WriteBody for GzipBody {
43 fn write_body(&mut self, w: &mut Write) -> io::Result<()> {
44 let mut encoder = libflate::gzip::Encoder::new(w)?;
45 self.0.write_body(&mut encoder)?;
46 encoder.finish().into_result().map(|_| ())
47 }
48}
49
50struct DeflateBody(Box<WriteBody>);
51
52impl WriteBody for DeflateBody {
53 fn write_body(&mut self, w: &mut Write) -> io::Result<()> {
54 let mut encoder = libflate::deflate::Encoder::new(w);
55 self.0.write_body(&mut encoder)?;
56 encoder.finish().into_result().map(|_| ())
57 }
58}
59
60fn encoding_matches_header(encoding: &CompressionEncoding, header: &Encoding) -> bool {
61 match encoding {
62 &CompressionEncoding::Brotli => *header == Encoding::EncodingExt(String::from("br")),
63 &CompressionEncoding::Deflate => *header == Encoding::Deflate,
64 &CompressionEncoding::Gzip => *header == Encoding::Gzip || *header == Encoding::EncodingExt(String::from("*")),
65 }
66}
67
68fn get_body(encoding: &CompressionEncoding, wrapped_body: Box<WriteBody>) -> Box<WriteBody> {
69 match encoding {
70 &CompressionEncoding::Brotli => Box::new(BrotliBody(wrapped_body)),
71 &CompressionEncoding::Deflate => Box::new(DeflateBody(wrapped_body)),
72 &CompressionEncoding::Gzip => Box::new(GzipBody(wrapped_body)),
73 }
74}
75
76fn get_header(encoding: &CompressionEncoding) -> Encoding {
77 match encoding {
78 &CompressionEncoding::Brotli => Encoding::EncodingExt(String::from("br")),
79 &CompressionEncoding::Deflate => Encoding::Deflate,
80 &CompressionEncoding::Gzip => Encoding::Gzip,
81 }
82}
83
84fn which_compression<'a, 'b>(req: &'b Request, res: &'b Response, priority: &Vec<CompressionEncoding>) -> Option<CompressionEncoding> {
85 return match (res.headers.get::<iron::headers::ContentEncoding>(), res.headers.get::<ContentLength>(), req.headers.get::<AcceptEncoding>()) {
86 (None, Some(content_length), Some(&AcceptEncoding(ref quality_items))) => {
87 if (content_length as &u64) < &DEFAULT_MIN_BYTES_FOR_COMPRESSION {
88 return None;
89 }
90
91 let max_quality = quality_items.iter().map(|qi| qi.quality).max();
92
93 if let Some(max_quality) = max_quality {
94 let quality_items: Vec<&QualityItem<Encoding>> = quality_items
95 .iter()
96 .filter(|qi| qi.quality != Quality(0) && qi.quality == max_quality)
97 .collect();
98
99 return priority
100 .iter()
101 .filter(|ce| quality_items.iter().find(|qi| {
102 encoding_matches_header(ce, &qi.item)
103 }).is_some())
104 .nth(0)
105 .map(|ce| ce.clone());
106 }
107 None
108 }
109 _ => None
110 };
111}
112
113pub struct CompressionMiddleware;
138
139impl AfterMiddleware for CompressionMiddleware {
140
141 fn after(&self, req: &mut Request, mut res: Response) -> IronResult<Response> {
143 let brotli = CompressionEncoding::Brotli;
144 let deflate = CompressionEncoding::Deflate;
145 let gzip = CompressionEncoding::Gzip;
146 let default_priorities = vec!(brotli, gzip, deflate);
147
148 if res.body.is_some() {
149 if let Some(compression) = which_compression(&req, &res, &default_priorities) {
150 res.headers.set(ContentEncoding(vec![get_header(&compression)]));
151 res.headers.remove::<ContentLength>();
152 res.body = Some(get_body(&compression, res.body.take().unwrap()));
153 }
154 }
155
156 Ok(res)
157 }
158}
159
160#[cfg(test)]
161mod test_common {
162 extern crate iron_test;
163
164 use std::io::Read;
165 use iron::prelude::*;
166 use iron::headers::*;
167 use iron::{Chain, status};
168 use iron::modifiers::Header;
169 use self::iron_test::{request};
170
171 use super::CompressionMiddleware;
172
173 pub fn build_compressed_echo_chain(with_encoding: bool) -> Chain {
174 let mut chain = Chain::new(move |req: &mut Request| {
175 let mut body: Vec<u8> = vec!();
176 req.body.read_to_end(&mut body).unwrap();
177
178 if !with_encoding {
179 Ok(Response::with((status::Ok, body)))
180 } else {
181 Ok(Response::with((status::Ok, Header(ContentEncoding(vec![Encoding::Chunked])), body)))
182 }
183 });
184 chain.link_after(CompressionMiddleware);
185 return chain;
186 }
187
188 pub fn post_data_with_accept_encoding(data: &str, accept_encoding: Option<AcceptEncoding>, chain: &Chain) -> Response {
189 let mut headers = Headers::new();
190 if let Some(value) = accept_encoding {
191 headers.set(value);
192 }
193
194 return request::post("http://localhost:3000/",
195 headers,
196 data,
197 chain).unwrap();
198 }
199}
200
201#[cfg(test)]
202mod uncompressable_tests {
203 extern crate iron_test;
204
205 use iron::headers::*;
206 use self::iron_test::{response};
207
208 use super::test_common::*;
209
210 #[test]
211 fn it_should_not_compress_response_when_client_does_not_send_accept_encoding_header() {
212 let chain = build_compressed_echo_chain(false);
213 let value = "a".repeat(1000);
214 let res = post_data_with_accept_encoding(&value, None, &chain);
215
216 assert_eq!(res.headers.get::<ContentEncoding>(), None);
217 assert_eq!(response::extract_body_to_string(res), value);
218 }
219
220 #[test]
221 fn it_should_not_compress_response_when_client_does_not_send_supported_encoding() {
222 let chain = build_compressed_echo_chain(false);
223 let value = "a".repeat(1000);
224 let res = post_data_with_accept_encoding(&value,
225 Some(AcceptEncoding(vec![qitem(Encoding::Chunked)])),
226 &chain);
227
228 assert_eq!(res.headers.get::<ContentEncoding>(), None);
229 assert_eq!(response::extract_body_to_bytes(res), value.into_bytes());
230 }
231
232 #[test]
233 fn it_should_not_compress_small_response() {
234 let value = "a".repeat(10);
235 let chain = build_compressed_echo_chain(false);
236 let res = post_data_with_accept_encoding(&value,
237 Some(AcceptEncoding(vec![qitem(Encoding::Gzip)])),
238 &chain);
239
240 assert_eq!(res.headers.get::<ContentEncoding>(), None);
241 assert_eq!(response::extract_body_to_bytes(res), value.into_bytes());
242 }
243
244 #[test]
245 fn it_should_not_compress_already_encoded_response() {
246 let value = "a".repeat(1000);
247 let chain = build_compressed_echo_chain(true);
248 let res = post_data_with_accept_encoding(&value,
249 Some(AcceptEncoding(vec![qitem(Encoding::Gzip)])),
250 &chain);
251
252 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::Chunked])));
253 assert_eq!(response::extract_body_to_bytes(res), value.into_bytes());
254 }
255}
256
257#[cfg(test)]
258mod gzip_tests {
259 extern crate iron_test;
260
261 use std::io::Read;
262 use iron::headers::*;
263 use self::iron_test::{response};
264 use libflate::gzip;
265
266 use super::test_common::*;
267
268 #[test]
269 fn it_should_compress_response_body_correctly_using_gzip_and_set_header() {
270 let value = "a".repeat(1000);
271 let chain = build_compressed_echo_chain(false);
272 let res = post_data_with_accept_encoding(&value,
273 Some(AcceptEncoding(vec![qitem(Encoding::Gzip)])),
274 &chain);
275
276 assert_eq!(res.headers.get::<ContentLength>(), None);
277 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::Gzip])));
278
279 let compressed_bytes = response::extract_body_to_bytes(res);
280 let mut decoder = gzip::Decoder::new(&compressed_bytes[..]).unwrap();
281 let mut decoded_data = Vec::new();
282 decoder.read_to_end(&mut decoded_data).unwrap();
283 assert_eq!(decoded_data, value.into_bytes());
284 }
285}
286
287#[cfg(test)]
288mod deflate_tests {
289 extern crate iron_test;
290
291 use std::io::Read;
292 use iron::headers::*;
293 use self::iron_test::{response};
294 use libflate::deflate;
295
296 use super::test_common::*;
297
298 #[test]
299 fn it_should_compress_response_body_correctly_using_deflate_and_set_header() {
300 let value = "a".repeat(1000);
301 let chain = build_compressed_echo_chain(false);
302 let res = post_data_with_accept_encoding(&value,
303 Some(AcceptEncoding(vec![qitem(Encoding::Deflate)])),
304 &chain);
305
306 assert_eq!(res.headers.get::<ContentLength>(), None);
307 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::Deflate])));
308
309 let compressed_bytes = response::extract_body_to_bytes(res);
310 let mut decoder = deflate::Decoder::new(&compressed_bytes[..]);
311 let mut decoded_data = Vec::new();
312 decoder.read_to_end(&mut decoded_data).unwrap();
313 assert_eq!(decoded_data, value.into_bytes());
314 }
315}
316
317#[cfg(test)]
318mod brotli_tests {
319 extern crate iron_test;
320
321 use std::io::Read;
322 use iron::headers::*;
323 use self::iron_test::{response};
324 use brotli;
325
326 use super::test_common::*;
327
328 #[test]
329 fn it_should_compress_response_body_correctly_using_brotli_and_set_header() {
330 let value = "a".repeat(1000);
331 let chain = build_compressed_echo_chain(false);
332 let res = post_data_with_accept_encoding(&value,
333 Some(AcceptEncoding(vec![
334 qitem(Encoding::EncodingExt(String::from("br")))
335 ])),
336 &chain);
337
338 assert_eq!(res.headers.get::<ContentLength>(), None);
339 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::EncodingExt(String::from("br"))])));
340
341 let compressed_bytes = response::extract_body_to_bytes(res);
342 let mut decoder = brotli::Decompressor::new(&compressed_bytes[..], 4096);
343 let mut decoded_data = Vec::new();
344 decoder.read_to_end(&mut decoded_data).unwrap();
345 assert_eq!(decoded_data, value.into_bytes());
346 }
347}
348
349#[cfg(test)]
350mod priority_tests {
351 use iron::headers::*;
352
353 use super::test_common::*;
354
355 #[test]
356 fn it_should_use_the_more_prior_compression_based_on_quality_for_gzip() {
357 let value = "a".repeat(1000);
358 let chain = build_compressed_echo_chain(false);
359 let res = post_data_with_accept_encoding(&value,
360 Some(AcceptEncoding(vec![
361 QualityItem { item: Encoding::Gzip, quality: q(0.5) },
362 QualityItem { item: Encoding::Deflate, quality: q(1.0) }
363 ])),
364 &chain);
365
366 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::Deflate])));
367 }
368
369 #[test]
370 fn it_should_use_the_more_prior_compression_based_on_quality_for_deflate() {
371 let value = "a".repeat(1000);
372 let chain = build_compressed_echo_chain(false);
373 let res = post_data_with_accept_encoding(&value,
374 Some(AcceptEncoding(vec![
375 QualityItem { item: Encoding::Deflate, quality: q(1.0) },
376 QualityItem { item: Encoding::Gzip, quality: q(0.5) }
377 ])),
378 &chain);
379
380 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::Deflate])));
381 }
382
383 #[test]
384 fn it_should_not_use_a_compression_with_quality_0() {
385 let value = "a".repeat(1000);
386 let chain = build_compressed_echo_chain(false);
387 let res = post_data_with_accept_encoding(&value,
388 Some(AcceptEncoding(vec![
389 QualityItem { item: Encoding::Gzip, quality: q(0.0) }
390 ])),
391 &chain);
392
393 assert_eq!(res.headers.get::<ContentEncoding>(), None);
394 }
395
396 #[test]
397 fn it_should_use_the_brotli_compression_preferably_when_explicitly_sent() {
398 let value = "a".repeat(1000);
399 let chain = build_compressed_echo_chain(false);
400 let res = post_data_with_accept_encoding(&value,
401 Some(AcceptEncoding(vec![
402 qitem(Encoding::EncodingExt(String::from("*"))),
403 qitem(Encoding::Gzip),
404 qitem(Encoding::EncodingExt(String::from("br"))),
405 qitem(Encoding::Deflate),
406 ])),
407 &chain);
408
409 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::EncodingExt(String::from("br"))])));
410 }
411
412 #[test]
413 fn it_should_use_the_gzip_compression_if_the_any_encoding_is_sent() {
414 let value = "a".repeat(1000);
415 let chain = build_compressed_echo_chain(false);
416 let res = post_data_with_accept_encoding(&value,
417 Some(AcceptEncoding(vec![
418 qitem(Encoding::EncodingExt(String::from("*"))),
419 qitem(Encoding::Deflate),
420 ])),
421 &chain);
422
423 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::Gzip])));
424 }
425
426 #[test]
427 fn it_should_use_the_gzip_compression_as_second_preference() {
428 let value = "a".repeat(1000);
429 let chain = build_compressed_echo_chain(false);
430 let res = post_data_with_accept_encoding(&value,
431 Some(AcceptEncoding(vec![
432 qitem(Encoding::Deflate),
433 qitem(Encoding::Gzip),
434 ])),
435 &chain);
436
437 assert_eq!(res.headers.get::<ContentEncoding>(), Some(&ContentEncoding(vec![Encoding::Gzip])));
438 }
439}
440
441#[cfg(all(feature = "unstable", test))]
442mod middleware_benchmarks {
443 macro_rules! bench_chain_with_header_and_size {
444 ($name:ident, $chain:expr, $header:expr, $response_size:expr) => {
445 #[bench]
446 fn $name(b: &mut Bencher) {
447 let chain = $chain;
448 let mut rng = rand::IsaacRng::new_unseeded();
449
450 b.iter(|| {
451 let data: String = rng.gen_ascii_chars().take($response_size).collect();
452 let res = post_data_with_accept_encoding(&data,
453 $header,
454 &chain);
455 let compressed_bytes = response::extract_body_to_bytes(res);
456
457 assert!(compressed_bytes.len() > 0);
458 })
459 }
460 };
461 }
462
463 macro_rules! bench_chains_with_size {
464 ($mod_name:ident, $size:expr) => {
465 mod $mod_name {
466 extern crate iron_test;
467 extern crate test;
468 extern crate rand;
469
470 use std::io::Read;
471 use iron::prelude::*;
472 use iron::{Chain, status};
473 use iron::headers::*;
474 use self::test::Bencher;
475 use self::iron_test::{response};
476 use self::rand::Rng;
477 use super::super::test_common::*;
478
479 fn build_echo_chain() -> Chain {
480 let chain = Chain::new(|req: &mut Request| {
481 let mut body: Vec<u8> = vec!();
482 req.body.read_to_end(&mut body).unwrap();
483 Ok(Response::with((status::Ok, body)))
484 });
485 return chain;
486 }
487
488 bench_chain_with_header_and_size!(without_middleware, build_echo_chain(), None, $size);
489 bench_chain_with_header_and_size!(with_middleware_no_accept_header, build_compressed_echo_chain(false), None, $size);
490 bench_chain_with_header_and_size!(with_middleware_gzip,
491 build_compressed_echo_chain(false),
492 Some(AcceptEncoding(vec![qitem(Encoding::Gzip)])),
493 $size);
494 bench_chain_with_header_and_size!(with_middleware_deflate,
495 build_compressed_echo_chain(false),
496 Some(AcceptEncoding(vec![qitem(Encoding::Deflate)])),
497 $size);
498 bench_chain_with_header_and_size!(with_middleware_brotli,
499 build_compressed_echo_chain(false),
500 Some(AcceptEncoding(vec![qitem(Encoding::EncodingExt(String::from("br")))])),
501 $size);
502 }
503 };
504 }
505
506 bench_chains_with_size!(response_1kb, 1024);
507 bench_chains_with_size!(response_128kb, 128 * 1024);
508 bench_chains_with_size!(response_1mb, 1024 * 1024);
509}