1use crate::utils::request::Request;
2use crate::utils::response::Response;
3use futures::future::BoxFuture;
4use serde_json::json;
5use std::collections::HashMap;
6use std::future::Future;
7use tokio::io::AsyncWriteExt;
8
9type Handler = Box<dyn Fn(&Request) -> Response + Send + Sync>;
10type AsyncHandler = Box<dyn Fn(&Request) -> BoxFuture<Response> + Send + Sync>;
11
12
13pub struct Router {
14 get_routes: HashMap<String, EitherHandler>,
15 post_routes: HashMap<String, EitherHandler>,
16 put_routes: HashMap<String, EitherHandler>,
17 delete_routes: HashMap<String, EitherHandler>,
18 patch_routes: HashMap<String, EitherHandler>,
19}
20
21#[allow(dead_code)]
22enum EitherHandler {
23 Sync(Handler),
24 Async(AsyncHandler),
25}
26
27#[allow(dead_code)]
28impl Router {
29 pub fn new() -> Self {
30 Self {
31 get_routes: HashMap::new(),
32 post_routes: HashMap::new(),
33 put_routes: HashMap::new(),
34 delete_routes: HashMap::new(),
35 patch_routes: HashMap::new(),
36 }
37 }
38
39 pub fn get<F>(&mut self, route: &str, handler: F)
40 where
41 F: Fn(&Request) -> Response + Send + Sync + 'static,
42 {
43 self.get_routes
44 .insert(route.to_string(), EitherHandler::Sync(Box::new(handler)));
45 }
46
47 pub fn get_async<F, Fut>(&mut self, route: &str, handler: F)
48 where
49 F: Fn(&Request) -> Fut + Send + Sync + 'static,
50 Fut: Future<Output = Response> + Send + 'static,
51 {
52 self.get_routes.insert(
53 route.to_string(),
54 EitherHandler::Async(Box::new(move |req| Box::pin(handler(req)))),
55 );
56 }
57
58 pub fn post<F>(&mut self, route: &str, handler: F)
59 where
60 F: Fn(&Request) -> Response + Send + Sync + 'static,
61 {
62 self.post_routes
63 .insert(route.to_string(), EitherHandler::Sync(Box::new(handler)));
64 }
65
66 pub fn post_async<F, Fut>(&mut self, route: &str, handler: F)
67 where
68 F: Fn(&Request) -> Fut + Send + Sync + 'static,
69 Fut: Future<Output = Response> + Send + 'static,
70 {
71 self.post_routes.insert(
72 route.to_string(),
73 EitherHandler::Async(Box::new(move |req| Box::pin(handler(req)))),
74 );
75 }
76
77 pub fn put<F>(&mut self, route: &str, handler: F)
78 where
79 F: Fn(&Request) -> Response + Send + Sync + 'static,
80 {
81 self.put_routes
82 .insert(route.to_string(), EitherHandler::Sync(Box::new(handler)));
83 }
84
85 pub fn put_async<F, Fut>(&mut self, route: &str, handler: F)
86 where
87 F: Fn(&Request) -> Fut + Send + Sync + 'static,
88 Fut: Future<Output = Response> + Send + 'static,
89 {
90 self.put_routes.insert(
91 route.to_string(),
92 EitherHandler::Async(Box::new(move |req| Box::pin(handler(req)))),
93 );
94 }
95
96 pub fn patch<F>(&mut self, route: &str, handler: F)
97 where
98 F: Fn(&Request) -> Response + Send + Sync + 'static,
99 {
100 self.patch_routes
101 .insert(route.to_string(), EitherHandler::Sync(Box::new(handler)));
102 }
103
104 pub fn patch_async<F, Fut>(&mut self, route: &str, handler: F)
105 where
106 F: Fn(&Request) -> Fut + Send + Sync + 'static,
107 Fut: Future<Output = Response> + Send + 'static,
108 {
109 self.patch_routes.insert(
110 route.to_string(),
111 EitherHandler::Async(Box::new(move |req| Box::pin(handler(req)))),
112 );
113 }
114
115 pub fn delete<F>(&mut self, route: &str, handler: F)
116 where
117 F: Fn(&Request) -> Response + Send + Sync + 'static,
118 {
119 self.delete_routes
120 .insert(route.to_string(), EitherHandler::Sync(Box::new(handler)));
121 }
122
123 pub fn delete_async<F, Fut>(&mut self, route: &str, handler: F)
124 where
125 F: Fn(&Request) -> Fut + Send + Sync + 'static,
126 Fut: Future<Output = Response> + Send + 'static,
127 {
128 self.delete_routes.insert(
129 route.to_string(),
130 EitherHandler::Async(Box::new(move |req| Box::pin(handler(req)))),
131 );
132 }
133
134 pub async fn handle_request(&self, stream: &mut tokio::net::TcpStream, request: &Request) {
135 let response = match request.method.as_str() {
136 "GET" => {
137 if let Some(handler) = self.get_routes.get(&request.route) {
138 self.execute_handler(handler, request).await
139 } else {
140 handle_not_found()
141 }
142 }
143 "POST" => {
144 if let Some(handler) = self.post_routes.get(&request.route) {
145 self.execute_handler(handler, request).await
146 } else {
147 handle_not_found()
148 }
149 }
150 "PUT" => {
151 if let Some(handler) = self.put_routes.get(&request.route) {
152 self.execute_handler(handler, request).await
153 } else {
154 handle_not_found()
155 }
156 }
157 "DELETE" => {
158 if let Some(handler) = self.delete_routes.get(&request.route) {
159 self.execute_handler(handler, request).await
160 } else {
161 handle_not_found()
162 }
163 }
164
165 "PATCH" => {
166 if let Some(handler) = self.patch_routes.get(&request.route) {
167 self.execute_handler(handler, request).await
168 } else {
169 handle_not_found()
170 }
171 }
172 _ => handle_method_not_allowed(),
173 };
174
175 stream
176 .write_all(response.as_bytes())
177 .await
178 .expect("Failed to write response back to the client");
179 }
180
181 async fn execute_handler(&self, handler: &EitherHandler, request: &Request) -> String {
182 match handler {
183 EitherHandler::Sync(sync_handler) => {
184 let mut response = sync_handler(request);
185 handle_ok(&mut response)
186 }
187 EitherHandler::Async(async_handler) => {
188 let mut response = async_handler(request).await;
189 handle_ok(&mut response)
190 }
191 }
192 }
193}
194
195fn handle_ok(response: &mut Response) -> String {
196 let mut cookies = String::new();
197 if !response.cookies.is_empty() {
198 for (k, v) in response.cookies.iter() {
199 cookies += &format!(
200 "Set-Cookie: {}={}; HttpOnly; Path=/; Domain=example.com; Max-Age=3600\r\n",
201 k, v
202 );
203 }
204 }
205
206 let mut content_disposition = String::new();
207 let body_content: String;
208 let content_length;
209
210 if response.is_file {
211 let body_string = response.body.to_string();
212 let mut parts = body_string.splitn(2, ":");
213
214 let filename = parts.next().unwrap_or("file.txt").trim_matches('"');
215
216 let file_data = parts
217 .next()
218 .unwrap_or("")
219 .replace("\\r\\n", "\n")
220 .replace("\\\"", "\"")
221 .trim_end_matches("\"")
222 .into();
223
224 content_disposition = format!(
225 "Content-Disposition: attachment; filename=\"{}\"\r\n",
226 filename
227 );
228 body_content = file_data;
229 content_length = body_content.len();
230 } else {
231 body_content = response.body.as_str().unwrap_or("").to_string();
232 content_length = body_content.len();
233 }
234
235 let new_response = format!(
236 "HTTP/1.1 {} {}\r\n{}Content-Type: {}\r\n{}Content-Length: {}\r\n\r\n{}",
237 response.status,
238 response.status_text,
239 content_disposition,
240 response.content_type,
241 cookies,
242 content_length,
243 body_content
244 );
245 new_response
246}
247
248fn handle_not_found() -> String {
249 format!(
250 "HTTP/1.1 404 Not Found\r\nContent-type: application/json\r\n\r\n{}",
251 serde_json::to_string(&json!({
252 "error": "Page not found"
253 }))
254 .unwrap()
255 )
256}
257
258fn handle_method_not_allowed() -> String {
259 format!(
260 "HTTP/1.1 405 Method Not Allowed\r\nContent-type: application/json\r\n\r\n{}",
261 serde_json::to_string(&json!({
262 "error": "This method is not allowed"
263 }))
264 .unwrap()
265 )
266}