rusty_web/response/
mod.rs

1use std::collections::HashMap;
2use std::io::{BufWriter, Write};
3use std::net::{Shutdown};
4use crate::headers::Headers;
5use crate::request::Request;
6use crate::status::{Status, StatusCode, StatusMethods};
7
8pub struct Response {
9    pub request: Request,
10    // Response headers
11    pub headers: Option<Headers>,
12    pub status: Option<usize>,
13    pub fixed_content: Option<String>,
14}
15
16impl Response {
17    pub fn new(request: Request) -> Self {
18        return Self {
19            request,
20            headers: None,
21            status: None,
22            fixed_content: None,
23        };
24    }
25
26    fn init_headers(&mut self) {
27        if !self.headers.is_some() {
28            self.headers = Some(HashMap::new());
29        }
30    }
31
32    pub fn set_content_type(&mut self, text: &str) {
33        self.init_headers();
34
35        // Unwrap header hashmap
36        if let Some(ref mut headers) = self.headers {
37            if headers.contains_key("Content-Type") {
38                let content_types = headers.get_mut("Content-Type").unwrap();
39
40                // Content type can only have one instance
41                content_types[0] = text.to_string();
42            } else {
43                let content_types = vec![text.to_string()];
44                let _ = headers.insert("Content-Type".to_string(), content_types);
45            }
46        }
47    }
48
49    /// Headers will be keep appending to the list if already exists
50    pub fn add_header(&mut self, name: &str, value: &str) -> &mut Self {
51        self.init_headers();
52
53        // Unwrap header hashmap
54        if let Some(ref mut headers) = self.headers {
55            if headers.contains_key(name) {
56                let values = headers.get_mut(name).unwrap();
57                values.push(value.to_string());
58            } else {
59                let values = vec![value.to_string()];
60                headers.insert(name.to_string(), values);
61            }
62        }
63
64        return self;
65    }
66
67    pub fn html<T: StatusCode>(&mut self, status: T, text: String) -> &mut Self {
68        self.set_content(status.to_usize(), text);
69        self.set_content_type("text/html");
70        return self;
71    }
72
73    pub fn json<T: StatusCode>(&mut self, status: T, text: String) -> &mut Self {
74        self.set_content(status.to_usize(), text);
75        self.set_content_type("application/json");
76        return self;
77    }
78
79    pub fn set_content(&mut self, status: usize, text: String) -> &mut Self {
80        self.status = Some(status);
81        self.fixed_content = Some(text);
82        return self;
83    }
84
85    pub fn send(&mut self) {
86        if self.status.is_some() {
87            let request = &self.request;
88            let access_from: String;
89
90            match request.stream.peer_addr() {
91                Ok(addr) => {
92                    access_from = addr.to_string();
93                }
94
95                Err(_) => {
96                    access_from = "UnKnown".to_string()
97                }
98            }
99
100            println!("{} - \"{} {} {}\"", access_from, request.method, request.pathname,
101                     self.status.unwrap());
102            self.write_http();
103        }
104    }
105
106    fn write_http(&mut self) {
107        let should_close = self.request.should_close_connection();
108
109        let headers = self.headers.as_mut().expect("Response headers missing.");
110        let content_length = format!("{}", self.fixed_content.as_ref()
111            .expect("Fixed content is missing.").len());
112        headers.insert("Content-Length".to_string(), vec![content_length]);
113
114        if !should_close {
115            headers.insert("Connection".to_string(), vec!["keep-alive".to_string()]);
116        }
117
118        // Write repose headers
119        let headers = self.prepare_raw_headers();
120
121        let cloned_stream = self.request.stream.try_clone();
122        if !cloned_stream.is_ok() {
123            println!("Connection closed");
124            self.request.context.dont_wait();
125            return;
126        }
127
128        let mut buf_writer = BufWriter::new(cloned_stream.unwrap());
129        match buf_writer.write_all(headers.as_bytes()) {
130            Ok(_) => {}
131            Err(_) => {
132                println!("Connection closed");
133                self.request.context.dont_wait();
134                return;
135            }
136        }
137
138        // Write response body
139        if self.request.method != "HEAD" {
140            if let Some(content) = &self.fixed_content {
141                buf_writer.write_all(content.as_bytes()).unwrap();
142            }
143        }
144
145        // Flush the buffer
146        if !buf_writer.flush().is_ok() {
147            print!("Connection closed");
148            self.request.context.dont_wait();
149        };
150
151        if should_close {
152            let _ = self.request.stream.shutdown(Shutdown::Both);
153            self.request.context.dont_wait();
154        }
155    }
156
157    fn prepare_raw_headers(&mut self) -> String {
158        let status_code = self.status.expect("Status code not set.");
159
160        let mut status_text = Status::status_text(status_code);
161        if !status_text.is_some() {
162            status_text = Some("Custom Status".to_string());
163        }
164
165        // Header start
166        let mut raw_headers = format!("HTTP/1.1 {} {}\r\n", self.status.unwrap(), status_text.unwrap());
167
168        if let Some(headers) = &self.headers {
169            for header_name in headers.keys() {
170                let values = headers.get(header_name).unwrap();
171
172                for value in values {
173                    let header_line = format!("{}: {}\r\n", header_name, value);
174                    raw_headers.push_str(&header_line);
175                }
176            }
177        }
178
179        // Header end
180        raw_headers.push_str("\r\n");
181        return raw_headers;
182    }
183}