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    DownloadStream(File, String),
16    FileStream(File),
17    StaticFile(&'static [u8], String),
18}
19
20#[derive(Debug)]
21pub struct HttpResponse {
22    pub content_type: String,
23    pub body: Option<Body>,
24    pub status_code: u16,
25    pub headers: Vec<(String, String)>,
26}
27
28impl HttpResponse {
29    pub fn new(body: Option<Body>, content_type: Option<String>, status_code: u16) -> Self {
30        HttpResponse {
31            content_type: content_type.unwrap_or_else(|| "application/json".to_string()),
32            body,
33            status_code,
34            headers: Vec::new(),
35        }
36    }
37    pub fn write_response(
38        self,
39        stream: &mut Box<dyn ReadWrite>,
40        compress: bool,
41    ) -> Result<(), Box<dyn std::error::Error>> {
42        let mut base_headers = format!(
43            "HTTP/1.1 {}\r\n\
44            Content-Type: {}\r\n\
45            Connection: keep-alive\r\n\
46            Server: RustHttpServer/1.0\r\n\
47            ",
48            self.status_code, self.content_type
49        );
50
51        self.headers.iter().for_each(|(key, value)| {
52            base_headers.push_str(&format!("{}: {}\r\n", key, value));
53        });
54
55        if let Some(body) = self.body {
56            match body {
57                Body::DownloadStream(file, name) => {
58                    return handle_file_stream(file, Some(name), base_headers, stream, true);
59                }
60                Body::FileStream(file) => {
61                    if compress {
62                        return handle_compressed_file_stream(file, base_headers, stream);
63                    }
64                    return handle_file_stream(file, None, base_headers, stream, false);
65                }
66                _ => {}
67            }
68
69            if compress {
70                base_headers.push_str("Content-Encoding: gzip\r\n");
71                base_headers.push_str("Vary: Accept-Encoding\r\n");
72
73                let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
74
75                match body {
76                    Body::Text(text) => encoder.write_all(text.as_bytes())?,
77                    Body::Json(json) => encoder.write_all(json.to_string().as_bytes())?,
78                    Body::StaticFile(file, _) => {
79                        encoder.write_all(file)?;
80                        let encoded = encoder.finish()?;
81                        base_headers.push_str(&format!("Content-Length: {}\r\n", encoded.len()));
82                        base_headers.push_str("\r\n");
83                        stream.write_all(base_headers.as_bytes())?;
84                        stream.write_all(&encoded)?;
85                        return Ok(());
86                    }
87                    Body::DownloadStream(_, _) | Body::FileStream(_) => unreachable!(),
88                }
89
90                let encoded = encoder.finish()?;
91                base_headers.push_str(&format!(
92                    "Content-Length: {}\r\n\
93                \r\n\
94                ",
95                    encoded.len(),
96                ));
97                stream.write_all(base_headers.as_bytes())?;
98                stream.write_all(&encoded)?;
99            } else {
100                let body = match body {
101                    Body::Text(text) => text.clone(),
102                    Body::Json(json) => json.to_string(),
103                    Body::StaticFile(file, _) => {
104                        base_headers.push_str(&format!("Content-Length: {}\r\n", file.len()));
105                        base_headers.push_str("\r\n");
106                        stream.write_all(base_headers.as_bytes())?;
107                        stream.write_all(file)?;
108                        return Ok(());
109                    }
110                    Body::DownloadStream(file, name) => {
111                        return handle_file_stream(file, Some(name), base_headers, stream, true);
112                    }
113                    Body::FileStream(file) => {
114                        return handle_file_stream(file, None, base_headers, stream, false);
115                    }
116                };
117                base_headers.push_str(&format!(
118                    "Content-Length: {}\r\n\
119                \r\n\
120                {}",
121                    body.len(),
122                    body
123                ));
124                stream.write_all(base_headers.as_bytes())?;
125            }
126        }
127
128        Ok(())
129    }
130    pub fn add_response_header(mut self, key: &str, value: &str) -> Self {
131        self.headers.push((key.to_string(), value.to_string()));
132        self
133    }
134}
135
136fn handle_file_stream(
137    file: File,
138    mut name: Option<String>,
139    mut headers: String,
140    mut stream: &mut Box<dyn ReadWrite>,
141    is_attachment: bool,
142) -> Result<(), Box<dyn std::error::Error>> {
143    let metadata = file.metadata()?;
144    let file_size = metadata.len();
145
146    headers.push_str(&format!("Content-Length: {}\r\n", file_size));
147
148    if is_attachment {
149        headers.push_str(&format!(
150            "Content-Disposition: attachment; filename=\"{}\"\r\n",
151            name.take().unwrap()
152        ));
153    }
154    headers.push_str("\r\n");
155
156    stream.write_all(headers.as_bytes())?;
157    let mut reader = BufReader::new(file);
158    io::copy(&mut reader, &mut stream)?;
159    Ok(())
160}
161
162fn handle_compressed_file_stream(
163    file: File,
164    mut headers: String,
165    stream: &mut Box<dyn ReadWrite>,
166) -> Result<(), Box<dyn std::error::Error>> {
167    if headers.contains("Connection: keep-alive") {
168        headers = headers.replace("Connection: keep-alive", "Connection: close");
169    }
170    headers.push_str("Content-Encoding: gzip\r\n");
171    headers.push_str("Vary: Accept-Encoding\r\n");
172    headers.push_str("\r\n");
173
174    stream.write_all(headers.as_bytes())?;
175
176    let mut encoder = GzEncoder::new(stream, Compression::default());
177    let mut reader = BufReader::new(file);
178    println!("headers: {}", headers);
179    io::copy(&mut reader, &mut encoder)?;
180    encoder.finish()?;
181
182    Ok(())
183}
scratch_server/http_response.rs

scratch_server/
http_response.rs
