Fall back to index.gmi, more modular code

Signed-off-by: Alexey Yerin <yerinalexey98fd@gmail.com>

src/gemini.rs

@@ -14,20 +14,13 @@ impl GeminiRequest {
         Ok(gemini_request)
     }
 
-    fn unsafe_file_path(&self) -> Option<&str> {
+    /// Get file path
+    pub fn file_path(&self) -> &str {
         self.path
             .path()
             .chars()
             .next()
-            .map(|c| &self.path.path()[c.len_utf8()..])
-    }
-
-    pub fn file_path(&self) -> Option<&str> {
-        match self.unsafe_file_path() {
-            Some(path) if path.contains("..") || path.starts_with("/") => None,
-            Some(path) => Some(path),
-            None => None,
-        }
+            .map_or("", |c| &self.path.path()[c.len_utf8()..])
     }
 }
 

src/main.rs

@@ -5,10 +5,12 @@ mod gemini;
 
 use native_tls::{Identity, TlsAcceptor, TlsStream};
 use std::fs::File;
+use std::io;
 use std::io::Read;
 use std::io::Write;
 use std::net::TcpListener;
 use std::net::TcpStream;
+use std::path;
 use std::sync::Arc;
 use std::thread;
 
@@ -35,16 +37,55 @@ fn main() {
             Ok(stream) => {
                 let acceptor = acceptor.clone();
                 thread::spawn(move || match acceptor.accept(stream) {
-                    Ok(stream) => handle_client(stream),
+                    Ok(stream) => {
+                        handle_client(stream).unwrap_or_else(|e| println!("Error: {}", e))
+                    }
                     Err(e) => println!("Can't handle stream: {}", e),
                 });
             }
-            Err(_e) => println!("Error: {}", _e),
+            Err(err) => println!("Error: {}", err),
         }
     }
 }
 
-fn handle_client(mut stream: TlsStream<TcpStream>) {
+/// Helper function to read a file into Vec
+fn read_file(file_path: &str) -> Result<Vec<u8>, io::Error> {
+    let mut file = File::open(file_path)?;
+    let mut buf = Vec::new();
+
+    file.read_to_end(&mut buf)?;
+
+    Ok(buf)
+}
+
+/// Send file as a response
+fn send_file(path: &str, response: &mut gemini::GeminiResonse) {
+    match read_file(path) {
+        Ok(buf) => {
+            response.body = Some(buf);
+        }
+        Err(err) => {
+            // Cannot read file or it doesn't exist
+
+            println!("Error ({}): {}", path, err);
+
+            response.status = [b'5', b'1'];
+            response.meta = format!("Resource not found: {}", path).into();
+        }
+    }
+}
+
+fn not_found(path: &str, response: &mut gemini::GeminiResonse) {
+    response.status = [b'5', b'1'];
+    response.meta = format!("Resource not found: {}", path).into();
+}
+
+fn redirect(path: &str, response: &mut gemini::GeminiResonse) {
+    response.status = [b'3', b'1'];
+    response.meta = path.into();
+}
+
+fn handle_client(mut stream: TlsStream<TcpStream>) -> Result<(), String> {
     let mut buffer = [0; 1024];
     if let Err(e) = stream.read(&mut buffer) {
         println!("Could not read from stream: {}", e)
@@ -60,34 +101,42 @@ fn handle_client(mut stream: TlsStream<TcpStream>) {
     let request = gemini::GeminiRequest::from_string(&raw_request).unwrap();
     let mut response = gemini::GeminiResonse::new();
 
-    match request.file_path() {
-        Some(file_path) => {
-            println!("Reading {}", file_path);
+    let url_path = request.file_path();
+    let file_path = path::Path::new(url_path);
 
-            match File::open(file_path) {
-                Ok(mut file) => {
-                    let mut body_buf = Vec::new();
-                    if let Err(e) = file.read_to_end(&mut body_buf) {
-                        println!("Could not read file {}", e);
-                    }
+    if file_path.has_root() {
+        // File starts with `/` (*nix) or `\\` (Windows), decline it
+        not_found(url_path, &mut response);
+    } else {
+        let path = path::Path::new(".").join(file_path).as_path().to_owned();
 
-                    response.body = Some(body_buf);
-                }
-                Err(e) => {
-                    let file_path = request.file_path().unwrap_or("");
-                    println!("Error ({}): {}", file_path, e);
-                    response.status = [b'5', b'1'];
-                    response.meta = format!("Resource not found: {}", file_path).into();
-                }
+        // Check if file/dir exists
+        if path.exists() {
+            // If it's a directory, try to find index.gmi
+            if path.is_dir() {
+                redirect(
+                    path.join("index.gmi")
+                        .iter()
+                        .skip(1)
+                        .collect::<path::PathBuf>()
+                        .to_str()
+                        .ok_or("Invalid Unicode".to_owned())?,
+                    &mut response,
+                );
+            } else {
+                send_file(
+                    path.to_str().ok_or("Invalid Unicode".to_owned())?,
+                    &mut response,
+                );
             }
-        }
-        None => {
-            println!("No file found in request");
-            response.status = [b'5', b'1']
+        } else {
+            not_found(url_path, &mut response);
         }
     }
 
     if let Err(e) = stream.write(&response.build()) {
         println!("Could not write to stream: {}", e);
     }
+
+    Ok(())
 }