Fix return link when random_route is set

1 files changed, 8 insertions, 1 deletions

diff --git a/src/main.rs b/src/main.rs
index 4b2118c..f0c23cc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -266,26 +266,33 @@ fn configure_app(app: App<MiniserveConfig>) -> App<MiniserveConfig> {
             })
             // Handle directories
             .handler(&full_route, s)
+            .default_resource(|r| r.method(Method::GET).f(p404))
         } else {
             // Handle directories
             app.handler(&full_route, s)
+                .default_resource(|r| r.method(Method::GET).f(p404))
         }
     } else {
         // Handle single files
         app.resource(&full_route, |r| r.f(listing::file_handler))
+            .default_resource(|r| r.method(Method::GET).f(p404))
     }
 }
 
 fn p404(req: &HttpRequest<crate::MiniserveConfig>) -> Result<HttpResponse, io::Error> {
     let err_404 = ContextualError::RouteNotFoundError(req.uri().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 => req.path().to_string(),
+    };
 
     errors::log_error_chain(err_404.to_string());
 
     Ok(actix_web::HttpResponse::NotFound().body(
         renderer::render_error(
             &err_404.to_string(),
-            "/",
+            &return_address,
             None,
             None,
             default_color_scheme,