iron_pack/
lib.rs

1#![cfg_attr(feature = "unstable", feature(test))]
2//! Compression middleware for Iron. This crate lets you automatically compress iron responses
3//! by providing an AfterMiddleware for your iron server.
4
5extern 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
113/// **Compression Middleware**
114///
115/// Currently either compresses using brotli, gzip or deflate algorithms. The algorithm is
116/// chosen by evaluating the `AcceptEncoding` header sent by the client.
117///
118/// # Example
119/// ```rust,no_run
120/// extern crate iron;
121/// extern crate iron_pack;
122///
123/// use iron::prelude::*;
124/// use iron_pack::CompressionMiddleware;
125///
126/// fn a_lot_of_batman(_: &mut Request) -> IronResult<Response> {
127///     let nana = "Na".repeat(5000);
128///     Ok(Response::with((iron::status::Ok, format!("{}, Batman!", nana))))
129/// }
130///
131/// fn main() {
132///     let mut chain = Chain::new(a_lot_of_batman);
133///     chain.link_after(CompressionMiddleware);
134///     Iron::new(chain).http("localhost:3000").unwrap();
135/// }
136/// ```
137pub struct CompressionMiddleware;
138
139impl AfterMiddleware for CompressionMiddleware {
140
141    /// Implementation of the compression middleware
142    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}