From 1931b07b0d8ec9e0a451ed7261a7643650f24abe Mon Sep 17 00:00:00 2001
From: Jackz <me@jackz.me>
Date: Wed, 16 Apr 2025 13:33:52 -0500
Subject: [PATCH] Add 404 and 401 check

---
 server/src/guards.rs                   | 37 ++++++++++++++++++++++++++
 server/src/main.rs                     | 27 ++++++++++++++++---
 server/src/routes/ui/user.rs           | 10 ++++---
 server/templates/auth/login.html.hbs   |  5 ++--
 server/templates/partials/nav.html.hbs |  1 +
 5 files changed, 70 insertions(+), 10 deletions(-)
 create mode 100644 server/src/guards.rs

diff --git a/server/src/guards.rs b/server/src/guards.rs
new file mode 100644
index 0000000..d3875cd
--- /dev/null
+++ b/server/src/guards.rs
@@ -0,0 +1,37 @@
+use rocket::http::Status;
+use rocket::Request;
+use rocket::request::{FromRequest, Outcome};
+use rocket_session_store::{Session, SessionResult};
+use crate::models::user::UserModel;
+use crate::{LoginSessionData, SessionData};
+
+pub struct AuthUser {
+    pub user: LoginSessionData
+}
+
+#[derive(Debug)]
+pub enum UserError {
+}
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for AuthUser {
+    type Error = UserError;
+
+    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
+        let sess = match request.guard::<Session<SessionData>>().await {
+            Outcome::Success(sess ) => sess,
+            _ => return Outcome::Forward(Status::Unauthorized)
+        };
+        let sess = match sess.get().await {
+            Ok(Some(sess)) => {
+                sess
+            }
+            _ => return Outcome::Forward(Status::Unauthorized),
+        };
+        if let Some(login) = &sess.login {
+            Outcome::Success(Self { user: login.clone() })
+        } else {
+            Outcome::Forward(Status::Unauthorized)
+        }
+    }
+}
\ No newline at end of file
diff --git a/server/src/main.rs b/server/src/main.rs
index 23b1b51..48b4b17 100644
--- a/server/src/main.rs
+++ b/server/src/main.rs
@@ -2,14 +2,15 @@ use std::net::IpAddr;
 use std::sync::Arc;
 use std::time::Duration;
 use log::{debug, error, info, trace, warn};
-use rocket::{catch, launch, routes, Request, State};
+use rocket::{catch, catchers, launch, routes, uri, Request, Route, State};
 use rocket::data::ByteUnit;
 use rocket::fs::{relative, FileServer};
 use rocket::futures::AsyncWriteExt;
 use rocket::http::private::cookie::CookieBuilder;
+use rocket::response::Redirect;
 use rocket::serde::Serialize;
 use rocket_dyn_templates::handlebars::{handlebars_helper, Context, Handlebars, Helper, HelperResult, Output, RenderContext};
-use rocket_dyn_templates::Template;
+use rocket_dyn_templates::{context, Template};
 use rocket_session_store::memory::MemoryStore;
 use rocket_session_store::SessionStore;
 use sqlx::{migrate, Pool, Postgres};
@@ -36,6 +37,7 @@ mod managers;
 mod objs;
 mod helpers;
 mod consts;
+mod guards;
 
 pub type DB = Pool<Postgres>;
 
@@ -128,10 +130,29 @@ async fn rocket() -> _ {
             ui::user::index, ui::user::redirect_list_library_files, ui::user::list_library_files, ui::user::get_library_file,
             ui::help::test_get
         ])
+        .register("/api", catchers![
+            not_found_api,
+        ])
+        .register("/", catchers![
+            not_found, not_authorized
+        ])
+}
+
+#[catch(401)]
+pub fn not_authorized(req: &Request) -> Redirect {
+    // uri!(ui::auth::login) doesn't work, it redirects to /login instead
+    Redirect::to(format!("/auth/login?path={}", req.uri()))
 }
 
 #[catch(404)]
-fn not_found(req: &Request) -> ResponseError {
+fn not_found(req: &Request) -> Template {
+    Template::render("errors/404", context! {
+        path: req.uri()
+    })
+}
+
+#[catch(404)]
+fn not_found_api(req: &Request) -> ResponseError {
     ResponseError::NotFound(
         JsonErrorResponse {
             code: "ROUTE_NOT_FOUND".to_string(),
diff --git a/server/src/routes/ui/user.rs b/server/src/routes/ui/user.rs
index 8278f3e..eb65883 100644
--- a/server/src/routes/ui/user.rs
+++ b/server/src/routes/ui/user.rs
@@ -2,7 +2,7 @@ use std::io::Cursor;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use log::debug;
-use rocket::{get, uri, Response, Route, State};
+use rocket::{catch, get, uri, Response, Route, State};
 use rocket::fs::NamedFile;
 use rocket::http::{ContentType, Header};
 use rocket::http::hyper::body::Buf;
@@ -13,9 +13,12 @@ use rocket_dyn_templates::{context, Template};
 use serde::Serialize;
 use serde_json::Value;
 use tokio::sync::Mutex;
+use crate::guards::{AuthUser};
 use crate::managers::libraries::LibraryManager;
+use crate::routes::ui::auth;
 use crate::util::{JsonErrorResponse, ResponseError};
 
+
 #[get("/")]
 pub async fn index(route: &Route) -> Template {
     Template::render("index", context! { user: true, route: route.uri.path(), test: "value" })
@@ -32,7 +35,7 @@ pub async fn redirect_list_library_files(libraries: &State<Arc<Mutex<LibraryMana
 
 
 #[get("/library/<library_id>/<_>/<path..>")]
-pub async fn list_library_files(route: &Route, libraries: &State<Arc<Mutex<LibraryManager>>>, library_id: &str, path: PathBuf)
+pub async fn list_library_files(user: AuthUser, route: &Route, libraries: &State<Arc<Mutex<LibraryManager>>>, library_id: &str, path: PathBuf)
     -> Result<Template, ResponseError>
 {
     let libs = libraries.lock().await;
@@ -63,9 +66,8 @@ pub async fn list_library_files(route: &Route, libraries: &State<Arc<Mutex<Libra
         .collect();
     debug!("parent={:?}", parent);
     debug!("segments={:?}", segments);
-
     Ok(Template::render("libraries", context! {
-        user: true,
+        user: user.user,
         route: route.uri.path(),
         library: library.model(),
         files: files,
diff --git a/server/templates/auth/login.html.hbs b/server/templates/auth/login.html.hbs
index e593b65..5d22ee9 100644
--- a/server/templates/auth/login.html.hbs
+++ b/server/templates/auth/login.html.hbs
@@ -4,8 +4,7 @@
         <h1 class="title is-1 has-text-centered">storage-app</h1>
         <div class="box is-radiusless">
             <h4 class="title is-4 has-text-centered">Login</h4>
-            {{debug this}}
-            {{#if (len form.form_errors) includeZero=true  }}
+            {{#unless (eq (len form.form_errors) 0)  }}
             <div class="notification is-danger is-light">
                 <b>Login failed with errors:</b>
                 <ul>
@@ -14,7 +13,7 @@
                     {{/each}}
                 </ul>
             </div>
-            {{/if}}
+            {{/unless}}
             <form method="post" action="/auth/login">
                 <input type="hidden" name="_csrf" value="{{ csrf_token }}">
                 <div class="field">
diff --git a/server/templates/partials/nav.html.hbs b/server/templates/partials/nav.html.hbs
index 039b7a2..d23d8bf 100644
--- a/server/templates/partials/nav.html.hbs
+++ b/server/templates/partials/nav.html.hbs
@@ -37,6 +37,7 @@
         <div class="navbar-item has-dropdown is-hoverable">
             <a class="navbar-link">
                 <img src="/static/img/default_user.png" alt="User Image" />
+                {{ debug user }}
             </a>
         
             <div class="navbar-dropdown">