krustie/middleware/
gzip.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
//! A middleware for compressing response body using gzip

use flate2::{write::GzEncoder, Compression};
use std::io::Write;

use super::Middleware;
use crate::{request::Request, response::Response, server::route_handler::HandlerResult};

/// A middleware for compressing response body using gzip.
///
/// # ExampleF
///
/// ```rust
/// use krustie::{server::Server, middleware::gzip::GzipEncoder};
///
/// let mut server = Server::create();
///
/// server.use_handler(GzipEncoder);
///
#[derive(Clone)]
pub struct GzipEncoder;

impl GzipEncoder {
    fn encode(body: &Vec<u8>) -> Result<Vec<u8>, String> {
        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());

        if encoder.write_all(body.as_slice()).is_err() {
            return Err("Error while writing to encoder".to_string());
        }

        match encoder.finish() {
            Ok(compressed_bytes) => {
                return Ok(compressed_bytes);
            }
            Err(err) => {
                return Err(format!("{err}"));
            }
        }
    }
}

impl Middleware for GzipEncoder {
    fn middleware(&mut self, request: &Request, response: &mut Response) -> HandlerResult {
        let body = response.get_body_mut();

        if body.is_empty() {
            return HandlerResult::Next;
        }

        if let Some(str_encodings) = request.get_header("accept-encoding") {
            let encodings = str_encodings
                .split(',')
                .map(|item| item.trim())
                .collect::<Vec<&str>>();

            if !encodings.contains(&"gzip") {
                return HandlerResult::Next;
            }

            match Self::encode(body) {
                Ok(compressed_bytes) => {
                    response.set_header("Content-Encoding", "gzip");

                    let _ = response.set_body(compressed_bytes);
                }
                Err(err) => {
                    eprintln!("Error while compressing: {}", err);
                }
            }
        }

        return HandlerResult::Next;
    }
}