Skip to main content

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            return match (body, compress) {
57                (Body::DownloadStream(file, name), _) => {
58                    handle_file_stream(file, Some(name), base_headers, stream, true)
59                }
60                (Body::FileStream(file), true) => {
61                    handle_compressed_file_stream(file, base_headers, stream)
62                }
63                (Body::FileStream(file), false) => {
64                    handle_file_stream(file, None, base_headers, stream, false)
65                }
66                (Body::Text(text), should_compress) => {
67                    write_buffered_body(base_headers, text.as_bytes(), should_compress, stream)
68                }
69                (Body::Json(json), should_compress) => {
70                    let serialized = json.to_string();
71                    write_buffered_body(
72                        base_headers,
73                        serialized.as_bytes(),
74                        should_compress,
75                        stream,
76                    )
77                }
78                (Body::StaticFile(file, _), should_compress) => {
79                    write_buffered_body(base_headers, file, should_compress, stream)
80                }
81            };
82        }
83
84        Ok(())
85    }
86    pub fn add_response_header(mut self, key: &str, value: &str) -> Self {
87        self.headers.push((key.to_string(), value.to_string()));
88        self
89    }
90}
91
92fn handle_file_stream(
93    file: File,
94    mut name: Option<String>,
95    mut headers: String,
96    mut stream: &mut Box<dyn ReadWrite>,
97    is_attachment: bool,
98) -> Result<(), Box<dyn std::error::Error>> {
99    let metadata = file.metadata()?;
100    let file_size = metadata.len();
101
102    headers.push_str(&format!("Content-Length: {}\r\n", file_size));
103
104    if is_attachment {
105        headers.push_str(&format!(
106            "Content-Disposition: attachment; filename=\"{}\"\r\n",
107            name.take().unwrap()
108        ));
109    }
110    headers.push_str("\r\n");
111
112    stream.write_all(headers.as_bytes())?;
113    let mut reader = BufReader::new(file);
114    io::copy(&mut reader, &mut stream)?;
115    Ok(())
116}
117
118fn handle_compressed_file_stream(
119    file: File,
120    mut headers: String,
121    stream: &mut Box<dyn ReadWrite>,
122) -> Result<(), Box<dyn std::error::Error>> {
123    if headers.contains("Connection: keep-alive") {
124        headers = headers.replace("Connection: keep-alive", "Connection: close");
125    }
126    headers.push_str("Content-Encoding: gzip\r\n");
127    headers.push_str("Vary: Accept-Encoding\r\n");
128    headers.push_str("\r\n");
129
130    stream.write_all(headers.as_bytes())?;
131
132    let mut encoder = GzEncoder::new(stream, Compression::default());
133    let mut reader = BufReader::new(file);
134    println!("headers: {}", headers);
135    io::copy(&mut reader, &mut encoder)?;
136    encoder.finish()?;
137
138    Ok(())
139}
140
141fn write_buffered_body(
142    mut headers: String,
143    body: &[u8],
144    compress: bool,
145    stream: &mut Box<dyn ReadWrite>,
146) -> Result<(), Box<dyn std::error::Error>> {
147    if compress {
148        headers.push_str("Content-Encoding: gzip\r\n");
149        headers.push_str("Vary: Accept-Encoding\r\n");
150
151        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
152        encoder.write_all(body)?;
153        let encoded = encoder.finish()?;
154
155        headers.push_str(&format!("Content-Length: {}\r\n", encoded.len()));
156        headers.push_str("\r\n");
157        stream.write_all(headers.as_bytes())?;
158        stream.write_all(&encoded)?;
159        return Ok(());
160    }
161
162    headers.push_str(&format!("Content-Length: {}\r\n", body.len()));
163    headers.push_str("\r\n");
164    stream.write_all(headers.as_bytes())?;
165    stream.write_all(body)?;
166
167    Ok(())
168}