1 files changed, 52 insertions, 36 deletions

diff --git a/src/main.rs b/src/main.rs
index f227be9..f26369a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,7 +1,7 @@
 #![feature(proc_macro_hygiene)]
 
-use actix_web::http::Method;
-use actix_web::{fs, middleware, server, App};
+use actix_web::http::{Method, StatusCode};
+use actix_web::{fs, middleware, server, App, HttpRequest, HttpResponse};
 use clap::crate_version;
 use simplelog::{Config, LevelFilter, TermLogger};
 use std::io::{self, Write};
@@ -19,7 +19,7 @@ mod listing;
 mod renderer;
 mod themes;
 
-use crate::errors::{ContextualError};
+use crate::errors::ContextualError;
 
 #[derive(Clone)]
 /// Configuration of the Miniserve application
@@ -83,12 +83,7 @@ fn run() -> Result<(), ContextualError> {
         && miniserve_config
             .path
             .symlink_metadata()
-            .map_err(|e|
-                ContextualError::IOError(
-                    "Failed to retrieve symlink's metadata".to_string(),
-                    e,
-                )
-            )?
+            .map_err(|e| ContextualError::IOError("Failed to retrieve symlink's metadata".to_string(), e))?
             .file_type()
             .is_symlink()
     {
@@ -117,10 +112,7 @@ fn run() -> Result<(), ContextualError> {
         .collect::<Vec<String>>();
 
     let canon_path = miniserve_config.path.canonicalize().map_err(|e| {
-        ContextualError::IOError(
-            "Failed to resolve path to be served".to_string(),
-            e,
-        )
+        ContextualError::IOError("Failed to resolve path to be served".to_string(), e)
     })?;
     let path_string = canon_path.to_string_lossy();
 
@@ -135,20 +127,14 @@ fn run() -> Result<(), ContextualError> {
             "      Invoke with -h|--help to see options or invoke as `miniserve .` to hide this advice."
         );
         print!("Starting server in ");
-        io::stdout().flush().map_err(|e| {
-            ContextualError::IOError(
-                "Failed to write data".to_string(),
-                e,
-            )
-        })?;
+        io::stdout()
+            .flush()
+            .map_err(|e| ContextualError::IOError("Failed to write data".to_string(), e))?;
         for c in "3… 2… 1… \n".chars() {
             print!("{}", c);
-            io::stdout().flush().map_err(|e| {
-                ContextualError::IOError(
-                    "Failed to write data".to_string(),
-                    e,
-                )
-            })?;
+            io::stdout()
+                .flush()
+                .map_err(|e| ContextualError::IOError("Failed to write data".to_string(), e))?;
             thread::sleep(Duration::from_millis(500));
         }
     }
@@ -210,12 +196,7 @@ fn run() -> Result<(), ContextualError> {
             .configure(configure_app)
     })
     .bind(socket_addresses.as_slice())
-    .map_err(|e| {
-        ContextualError::IOError(
-            "Failed to bind server".to_string(),
-            e,
-        )
-    })?
+    .map_err(|e| ContextualError::IOError("Failed to bind server".to_string(), e))?
     .shutdown_timeout(0)
     .start();
 
@@ -239,7 +220,7 @@ fn configure_app(app: App<MiniserveConfig>) -> App<MiniserveConfig> {
         let path = &app.state().path;
         let no_symlinks = app.state().no_symlinks;
         let random_route = app.state().random_route.clone();
-        let default_color_scheme = app.state().default_color_scheme.clone();
+        let default_color_scheme = app.state().default_color_scheme;
         let file_upload = app.state().file_upload;
         upload_route = if let Some(random_route) = app.state().random_route.clone() {
             format!("/{}/upload", random_route)
@@ -261,10 +242,11 @@ fn configure_app(app: App<MiniserveConfig>) -> App<MiniserveConfig> {
                             no_symlinks,
                             file_upload,
                             random_route.clone(),
-                            default_color_scheme.clone(),
+                            default_color_scheme,
                             u_r.clone(),
                         )
-                    }),
+                    })
+                    .default_handler(error_404),
             )
         }
     };
@@ -274,18 +256,52 @@ fn configure_app(app: App<MiniserveConfig>) -> App<MiniserveConfig> {
 
     if let Some(s) = s {
         if app.state().file_upload {
+            let default_color_scheme = app.state().default_color_scheme;
             // Allow file upload
-            app.resource(&upload_route, |r| {
-                r.method(Method::POST).f(file_upload::upload_file)
+            app.resource(&upload_route, move |r| {
+                r.method(Method::POST)
+                    .f(move |file| file_upload::upload_file(file, default_color_scheme))
             })
             // Handle directories
             .handler(&full_route, s)
+            .default_resource(|r| r.method(Method::GET).f(error_404))
         } else {
             // Handle directories
             app.handler(&full_route, s)
+                .default_resource(|r| r.method(Method::GET).f(error_404))
         }
     } else {
         // Handle single files
         app.resource(&full_route, |r| r.f(listing::file_handler))
+            .default_resource(|r| r.method(Method::GET).f(error_404))
     }
 }
+
+fn error_404(req: &HttpRequest<crate::MiniserveConfig>) -> Result<HttpResponse, io::Error> {
+    let err_404 = ContextualError::RouteNotFoundError(req.path().to_string());
+    let default_color_scheme = req.state().default_color_scheme;
+    let return_address = match &req.state().random_route {
+        Some(random_route) => format!("/{}", random_route),
+        None => "/".to_string(),
+    };
+
+    let query_params = listing::extract_query_parameters(req);
+    let color_scheme = query_params.theme.unwrap_or(default_color_scheme);
+
+    errors::log_error_chain(err_404.to_string());
+
+    Ok(actix_web::HttpResponse::NotFound().body(
+        renderer::render_error(
+            &err_404.to_string(),
+            StatusCode::NOT_FOUND,
+            &return_address,
+            query_params.sort,
+            query_params.order,
+            color_scheme,
+            default_color_scheme,
+            false,
+            true,
+        )
+        .into_string(),
+    ))
+}