scratch_server/
http_response.rs

1use std::{
2    fs::File,
3    io::{self, BufReader, Write},
4};
5
6use flate2::write::GzEncoder;
7use flate2::Compression;
8
9use crate::ReadWrite;
10
11#[derive(Debug)]
12pub enum Body {
13    Text(String),
14    Json(serde_json::Value),
15    FileStream(File, String),
16    StaticFile(&'static [u8], String),
17}
18
19#[derive(Debug)]
20pub struct HttpResponse {
21    pub content_type: String,
22    pub body: Option<Body>,
23    pub status_code: u16,
24    pub headers: Vec<(String, String)>,
25}
26
27impl HttpResponse {
28    pub fn new(body: Option<Body>, content_type: Option<String>, status_code: u16) -> Self {
29        HttpResponse {
30            content_type: content_type.unwrap_or_else(|| "application/json".to_string()),
31            body,
32            status_code,
33            headers: Vec::new(),
34        }
35    }
36    pub fn write_response(
37        self,
38        stream: &mut Box<dyn ReadWrite>,
39        compress: bool,
40    ) -> Result<(), Box<dyn std::error::Error>> {
41        let mut base_headers = format!(
42            "HTTP/1.1 {}\r\n\
43            Content-Type: {}\r\n\
44            Connection: keep-alive\r\n\
45            Server: RustHttpServer/1.0\r\n\
46            ",
47            self.status_code, self.content_type
48        );
49
50        self.headers.iter().for_each(|(key, value)| {
51            base_headers.push_str(&format!("{}: {}\r\n", key, value));
52        });
53
54        if let Some(body) = self.body {
55            if compress {
56                base_headers.push_str("Content-Encoding: gzip\r\n");
57                base_headers.push_str("Vary: Accept-Encoding\r\n");
58
59                let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
60
61                match body {
62                    Body::Text(text) => encoder.write_all(text.as_bytes())?,
63                    Body::Json(json) => encoder.write_all(json.to_string().as_bytes())?,
64                    Body::StaticFile(file, _) => {
65                        encoder.write_all(file)?;
66                        let encoded = encoder.finish()?;
67                        base_headers.push_str(&format!("Content-Length: {}\r\n", encoded.len()));
68                        base_headers.push_str("\r\n");
69                        stream.write_all(base_headers.as_bytes())?;
70                        stream.write_all(&encoded)?;
71                        return Ok(());
72                    }
73                    Body::FileStream(file, name) => {
74                        return handle_file_stream(file, name, base_headers, stream);
75                    }
76                }
77
78                let encoded = encoder.finish()?;
79                base_headers.push_str(&format!(
80                    "Content-Length: {}\r\n\
81                \r\n\
82                ",
83                    encoded.len(),
84                ));
85                stream.write_all(base_headers.as_bytes())?;
86                stream.write_all(&encoded)?;
87            } else {
88                let body = match body {
89                    Body::Text(text) => text.clone(),
90                    Body::Json(json) => json.to_string(),
91                    Body::StaticFile(file, _) => {
92                        base_headers.push_str(&format!("Content-Length: {}\r\n", file.len()));
93                        base_headers.push_str("\r\n");
94                        stream.write_all(base_headers.as_bytes())?;
95                        stream.write_all(file)?;
96                        return Ok(());
97                    }
98                    Body::FileStream(file, name) => {
99                        return handle_file_stream(file, name, base_headers, stream);
100                    }
101                };
102                base_headers.push_str(&format!(
103                    "Content-Length: {}\r\n\
104                \r\n\
105                {}",
106                    body.len(),
107                    body
108                ));
109                stream.write_all(base_headers.as_bytes())?;
110            }
111        }
112
113        Ok(())
114    }
115    pub fn add_response_header(mut self, key: &str, value: &str) -> Self {
116        self.headers.push((key.to_string(), value.to_string()));
117        self
118    }
119}
120
121fn handle_file_stream(
122    file: File,
123    name: String,
124    mut headers: String,
125    mut stream: &mut Box<dyn ReadWrite>,
126) -> Result<(), Box<dyn std::error::Error>> {
127    headers.push_str(&format!(
128        "Content-Disposition: attachment; filename=\"{}\"\r\n\
129    \rn",
130        name
131    ));
132    stream.write_all(headers.as_bytes())?;
133    let mut reader = BufReader::new(file);
134    io::copy(&mut reader, &mut stream)?;
135    Ok(())
136}