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