Skip to main content

scratch_server/
lib.rs

1use api_error::ApiError;
2use http_parse_error::HttpParseError;
3use include_dir::{include_dir, Dir};
4use logger::Logger;
5use native_tls::{Identity, TlsAcceptor};
6use std::borrow::Cow;
7use std::collections::HashMap;
8use std::fs::{self, File};
9use std::io::{self, BufRead, BufReader, Read, Write};
10use std::net::{IpAddr, SocketAddr, TcpListener, TcpStream};
11use std::path::PathBuf;
12use std::sync::Arc;
13use termcolor::Color;
14use utils::get_option;
15
16mod errors;
17mod http_response;
18mod logger;
19mod router;
20mod thread_pool;
21mod utils;
22
23pub use errors::*;
24pub use http_response::*;
25pub use router::*;
26
27pub static STATIC_FILES: Dir<'_> = include_dir!("src/dist");
28
29pub trait ReadWrite: Read + Write + Send + 'static {}
30
31impl<T: Read + Write + Send + 'static> ReadWrite for T {}
32
33struct NetworkStream {
34    delegate: Option<Box<dyn ReadWrite>>,
35    tls_acceptor: Option<TlsAcceptor>,
36}
37
38impl NetworkStream {
39    pub fn new(
40        cert_path: Option<&PathBuf>,
41        cert_pass: Option<&String>,
42    ) -> Result<NetworkStream, Box<dyn std::error::Error>> {
43        match &cert_path {
44            Some(path) => {
45                let identity_bytes = fs::read(path)?;
46
47                let identity = Identity::from_pkcs12(&identity_bytes, cert_pass.unwrap())?;
48
49                let tls_acceptor = TlsAcceptor::new(identity)?;
50
51                Ok(NetworkStream {
52                    tls_acceptor: Some(tls_acceptor),
53                    delegate: None,
54                })
55            }
56            None => Ok(NetworkStream {
57                tls_acceptor: None,
58                delegate: None,
59            }),
60        }
61    }
62    pub fn get_stream(
63        &mut self,
64        stream: TcpStream,
65    ) -> Result<&mut NetworkStream, Box<dyn std::error::Error>> {
66        match &self.tls_acceptor {
67            Some(acceptor) => {
68                let tls_stream = acceptor.accept(stream)?;
69                self.delegate = Some(Box::new(tls_stream));
70                Ok(self)
71            }
72            None => {
73                self.delegate = Some(Box::new(stream));
74                Ok(self)
75            }
76        }
77    }
78}
79
80pub struct HttpServer {
81    port: u16,
82    threads: usize,
83    cert_path: Option<PathBuf>,
84    cert_pass: Option<String>,
85    router: Router,
86    logger: Option<Arc<Logger>>,
87    bind_address: IpAddr,
88    compression: bool,
89}
90
91impl HttpServer {
92    pub fn build(
93        port: u16,
94        threads: usize,
95        cert_path: Option<PathBuf>,
96        cert_pass: Option<String>,
97        bind_address: IpAddr,
98        compression: bool,
99    ) -> HttpServer {
100        HttpServer {
101            port,
102            threads,
103            cert_path,
104            cert_pass,
105            router: Router::new(),
106            logger: None,
107            bind_address,
108            compression,
109        }
110    }
111    pub fn with_logger(mut self) -> Self {
112        self.logger = Some(Arc::new(Logger::new()));
113        self.router = self
114            .router
115            .with_logger(Some(Arc::clone(self.logger.as_ref().unwrap())));
116        self
117    }
118
119    pub fn with_credentials(mut self, password: &str, username: &str) -> Self {
120        self.router = self.router.with_credentials(username, password);
121        self
122    }
123
124    pub fn add_routes<F>(mut self, routes: F) -> Self
125    where
126        F: Fn(&mut Router) + Send + Sync + 'static,
127    {
128        routes(&mut self.router);
129        self
130    }
131
132    pub fn with_cors_policy(mut self, policy: Cors) -> Self {
133        self.router = self.router.with_cors(policy);
134        self
135    }
136    pub fn run(self) -> Result<(), Box<dyn std::error::Error>> {
137        self.print_server_info();
138        let listener = TcpListener::bind(SocketAddr::from((self.bind_address, self.port)))?;
139        let pool = thread_pool::ThreadPool::build(self.threads)?;
140
141        let arc_router = Arc::new(self.router);
142        let mut network_stream =
143            NetworkStream::new(self.cert_path.as_ref(), self.cert_pass.as_ref())?;
144        for stream in listener.incoming() {
145            let stream = stream?;
146            let peer_addr = stream.peer_addr()?;
147            let Ok(stream) = network_stream.get_stream(stream) else {
148                continue;
149            };
150            let mut stream = stream.delegate.take().unwrap();
151
152            let router_clone = Arc::clone(&arc_router);
153            let logger_clone = self.logger.clone();
154
155            pool.execute(move || {
156                let (response, req_headers) =
157                    handle_connection(&mut stream, &router_clone, peer_addr.ip()).unwrap_or_else(
158                        |err| {
159                            if let (Some(method), Some(path)) = (&err.method, &err.path) {
160                                router_clone
161                                    .log_response(
162                                        err.error_response.status_code,
163                                        path,
164                                        method,
165                                        peer_addr.ip(),
166                                    )
167                                    .unwrap();
168                            }
169
170                            (err.into_response(), HashMap::new())
171                        },
172                    );
173
174                let compress = self.compression
175                    && req_headers.iter().any(|(key, value)| {
176                        key.eq_ignore_ascii_case("accept-encoding")
177                            && value.to_ascii_lowercase().contains("gzip")
178                    });
179
180                response
181                    .write_response(&mut stream, compress)
182                    .unwrap_or_else(|err| {
183                        if let Some(logger) = logger_clone {
184                            logger
185                                .log_stderr("Error: {}", vec![(err.to_string(), Some(Color::Red))])
186                                .unwrap();
187                        }
188                    });
189            })?;
190        }
191        Ok(())
192    }
193    fn print_server_info(&self) {
194        if let Some(logger) = &self.logger {
195            logger.log_stdout(
196                r#"
197
198 ========================================================================================================
199 
200   _____ _                 _        _    _ _______ _______ _____     _____                          
201  / ____(_)               | |      | |  | |__   __|__   __|  __ \   / ____|                         
202 | (___  _ _ __ ___  _ __ | | ___  | |__| |  | |     | |  | |__) | | (___   ___ _ ____   _____ _ __ 
203  \___ \| | '_ ` _ \| '_ \| |/ _ \ |  __  |  | |     | |  |  ___/   \___ \ / _ \ '__\ \ / / _ \ '__|
204  ____) | | | | | | | |_) | |  __/ | |  | |  | |     | |  | |       ____) |  __/ |   \ V /  __/ |   
205 |_____/|_|_| |_| |_| .__/|_|\___| |_|  |_|  |_|     |_|  |_|      |_____/ \___|_|    \_/ \___|_|   
206                    | |                                                                             
207                    |_|                                                                             
208
209=========================================================================================================
210
211Port: {}
212Threads: {}
213HTTPS: {}
214CORS: {}
215Auth: {}
216Compression: {}
217
218====================
219Logs:"#,
220                vec![
221                    (self.port.to_string(), Some(Color::Blue)),
222                    (self.threads.to_string(), Some(Color::Blue)),
223                    get_option(&self.cert_path),
224                    get_option(&self.router.cors),
225                    get_option(&self.router.credentials),
226                    if self.compression { ("Enabled".to_string(), Some(Color::Green)) } else { ("Disabled".to_string(), Some(Color::Yellow)) },
227                ],
228            )
229            .unwrap();
230        }
231    }
232}
233
234fn parse_http<'a>(
235    reader: &mut BufReader<&mut Box<dyn ReadWrite>>,
236    request_string: &'a mut String,
237) -> Result<(&'a str, &'a str, HashMap<&'a str, &'a str>), HttpParseError> {
238    loop {
239        let mut line = String::new();
240        reader.read_line(&mut line)?;
241        request_string.push_str(&line);
242        if line == "\r\n" {
243            break;
244        }
245    }
246    let http_parts: Vec<&str> = request_string.split("\r\n\r\n").collect();
247    let request_lines: Vec<&str> = http_parts
248        .first()
249        .ok_or(HttpParseError::default())?
250        .lines()
251        .collect();
252
253    let http_method: Vec<&str> = request_lines
254        .first()
255        .ok_or(HttpParseError::default())?
256        .split_whitespace()
257        .collect();
258
259    if http_method.len() < 3 {
260        return Err(HttpParseError::default());
261    }
262
263    let (method, path, _version) = (http_method[0], http_method[1], http_method[2]);
264
265    let mut headers = std::collections::HashMap::new();
266    for line in &request_lines[1..] {
267        let parts: Vec<&str> = line.splitn(2, ':').collect();
268        if parts.len() == 2 {
269            headers.insert(
270                *parts.first().ok_or(HttpParseError::default())?,
271                parts.get(1).ok_or(HttpParseError::default())?.trim(),
272            );
273        }
274    }
275
276    Ok((method, path, headers))
277}
278
279fn handle_connection(
280    stream: &mut Box<dyn ReadWrite>,
281    router: &Arc<Router>,
282    peer_addr: IpAddr,
283) -> Result<(HttpResponse, HashMap<String, String>), ApiError> {
284    let mut reader = BufReader::new(&mut *stream);
285
286    let mut request = String::new();
287    let (method, path, headers) = parse_http(&mut reader, &mut request)?;
288
289    //find a better way without converting to owned headers
290    let owned_headers: HashMap<String, String> = headers
291        .iter()
292        .map(|(k, v)| (k.to_string(), v.to_string()))
293        .collect();
294
295    let mut buffer = Vec::new();
296
297    let body = match headers.get("Content-Type") {
298        Some(content_type) if content_type.contains("multipart/form-data") => {
299            let path = headers.get("Path").unwrap();
300            let response = handle_multipart_file_upload(content_type, &headers, &mut reader, path)
301                .map_err(|err| {
302                    ApiError::new_with_html(400, &format!("File upload error: {}", err))
303                })?;
304            return Ok((response, owned_headers));
305        }
306        _ => parse_body(&headers, reader, &mut buffer)?,
307    };
308
309    let response = router.route(path, method, body.as_deref(), peer_addr, &headers)?;
310
311    Ok((response, owned_headers))
312}
313
314fn parse_body<'a>(
315    headers: &HashMap<&str, &str>,
316    reader: BufReader<&mut Box<dyn ReadWrite>>,
317    buffer: &'a mut Vec<u8>,
318) -> Result<Option<Cow<'a, str>>, Box<dyn std::error::Error>> {
319    match headers.get("Content-Length") {
320        Some(content_length) => {
321            let content_length = content_length.parse::<usize>()?;
322            let mut body_reader = reader.take(content_length.try_into()?);
323            body_reader.read_to_end(buffer)?;
324            let body = String::from_utf8_lossy(&buffer[..]);
325            Ok(Some(body))
326        }
327        None => Ok(None),
328    }
329}
330
331fn handle_multipart_file_upload(
332    content_type: &str,
333    headers: &HashMap<&str, &str>,
334    reader: &mut BufReader<&mut Box<dyn ReadWrite>>,
335    path: &str,
336) -> Result<HttpResponse, Box<dyn std::error::Error>> {
337    let idx = content_type
338        .find("boundary=")
339        .ok_or("Missing multipart boundary")?;
340    let boundary = format!("--{}", &content_type[(idx + "boundary=".len())..]);
341    let mut multipart_headers = HashMap::new();
342    let mut header_size = 0;
343
344    //read headers
345    loop {
346        let mut line = String::new();
347        header_size += reader.read_line(&mut line)?;
348        if line.trim() == boundary {
349            continue;
350        }
351        if line == "\r\n" {
352            break;
353        }
354
355        let parts: Vec<&str> = line.trim().split(':').map(|s| s.trim()).collect();
356        if parts.len() < 2 {
357            return Err("Error parsing multipart request".into());
358        }
359        multipart_headers.insert(parts[0].to_owned(), parts[1].to_owned());
360    }
361
362    //get file name from content disposition and form target path
363    let content_disposition = multipart_headers
364        .get("Content-Disposition")
365        .ok_or("Missing content disposition")?;
366    let filename = content_disposition
367        .split("filename=\"")
368        .nth(1)
369        .and_then(|s| s.split('\"').next())
370        .ok_or("Error parsing file name")?;
371    let mut target_path = PathBuf::from("./").canonicalize()?.join(path);
372    target_path.push(filename);
373
374    let current_dir = std::env::current_dir()?;
375    if !target_path.starts_with(current_dir) {
376        return Err("Only paths relative to the current directory are allowed".into());
377    }
378
379    //calculate file size based on whole content length so that reading the stream can be stopped
380    let mut file = File::create(target_path)?;
381    let content_length = headers
382        .get("Content-Length")
383        .ok_or("Missing content length")?
384        .parse::<usize>()?;
385    let file_bytes = content_length - boundary.len() - header_size - 6;
386
387    //take only the file length from the main buf reader
388    let mut limited_reader = reader.take(file_bytes.try_into()?);
389
390    //copy streams
391    io::copy(&mut limited_reader, &mut file)?;
392
393    let response = HttpResponse::new(
394        Some(Body::Text(format!(
395            "File {} uploaded successfully.",
396            filename
397        ))),
398        Some(String::from("text/plain")),
399        200,
400    );
401    Ok(response)
402}