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}