afire_compress/
lib.rs

1//! Compress outgoing HTTP
2
3#![warn(missing_docs)]
4
5use std::cell::RefCell;
6use std::fmt;
7use std::io::prelude::*;
8
9use afire::{
10    middleware::{MiddleResponse, Middleware},
11    Header, Request, Response, Server,
12};
13
14use brotli2;
15use flate2;
16use libflate::deflate;
17
18/// Compression Methods
19#[derive(Debug, Clone, Copy)]
20pub enum CompressType {
21    /// Gzip Compression
22    ///
23    /// The number is the quality (0-9)
24    Gzip(u32),
25
26    /// Deflate Compression
27    Deflate,
28
29    /// Brotli Compression
30    ///
31    /// The number is the quality (0-9)
32    Brotli(u32),
33}
34
35/// Compression Middleware
36#[derive(Debug, Clone, Copy)]
37pub struct Compress {
38    compression: CompressType,
39    threshold: usize,
40}
41
42impl Middleware for Compress {
43    fn post(&mut self, req: Request, res: Response) -> MiddleResponse {
44        // Dont compress if body is under threshold
45        if res.data.len() <= self.threshold {
46            return MiddleResponse::Continue;
47        }
48
49        // Check if client dosent support compression
50        match req.header("Accept-Encoding") {
51            Some(i) => {
52                if !i
53                    .split(",")
54                    .map(|x| x.trim().to_owned())
55                    .collect::<Vec<_>>()
56                    .contains(&self.compression.to_string())
57                {
58                    return MiddleResponse::Continue;
59                }
60            }
61            None => return MiddleResponse::Continue,
62        }
63
64        // Compress with specified method
65        let new = match self.compression {
66            CompressType::Gzip(level) => {
67                let mut encoder =
68                    flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::new(level));
69                encoder.write_all(&res.data).unwrap();
70                (encoder.finish().unwrap(), "gzip")
71            }
72
73            CompressType::Deflate => {
74                let mut encoder = deflate::Encoder::new(Vec::new());
75                encoder.write_all(&res.data).unwrap();
76                (encoder.finish().into_result().unwrap(), "deflate")
77            }
78
79            CompressType::Brotli(level) => {
80                let encoder = brotli2::read::BrotliEncoder::new(&*res.data, level);
81                (encoder.into_inner().to_vec(), "br")
82            }
83        };
84
85        MiddleResponse::Add(
86            res.bytes(new.0)
87                .header(Header::new("Content-Encoding", new.1))
88                .content(afire::Content::TXT),
89        )
90    }
91}
92
93impl Compress {
94    /// Make a new Compressor
95    pub fn new() -> Self {
96        Compress {
97            compression: CompressType::Gzip(6),
98            threshold: 1024,
99        }
100    }
101
102    /// Set the body size threshold.
103    /// This stops from compressing tiny ammounts of data
104    ///
105    /// Deafult is 1024
106    pub fn threshold(self, threshold: usize) -> Self {
107        Compress { threshold, ..self }
108    }
109
110    /// Set compression method
111    ///
112    /// Defult is Gzip(6)
113    pub fn compression(self, compression: CompressType) -> Self {
114        Compress {
115            compression,
116            ..self
117        }
118    }
119
120    /// Attach Compressor to a server
121    pub fn attach(self, server: &mut Server)
122    where
123        Self: Sized + 'static,
124    {
125        server.middleware.push(Box::new(RefCell::new(self)));
126    }
127}
128
129impl fmt::Display for CompressType {
130    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
131        write!(
132            f,
133            "{}",
134            match self {
135                CompressType::Gzip(_) => "gzip",
136                CompressType::Deflate => "deflate",
137                CompressType::Brotli(_) => "br",
138            }
139        )
140    }
141}
afire_compress/lib.rs

afire_compress/
lib.rs
