Update auth error handling

v3xlabs · Jul 29, 2024 · 33c2d2b · 33c2d2b
1 parent 6c5cd33
commit 33c2d2b
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 17 deletions.
diff --git a/engine/src/auth/middleware.rs b/engine/src/auth/middleware.rs
@@ -1,7 +1,6 @@
 use std::sync::Arc;
 
 use poem::{web::Data, Error, FromRequest, Request, RequestBody, Result};
-use reqwest::StatusCode;
 use uuid::Uuid;
 
 use crate::state::AppState;

diff --git a/engine/src/routes/auth.rs b/engine/src/routes/auth.rs
@@ -8,8 +8,9 @@ use poem::{
     handler,
     http::HeaderMap,
     web::{Data, Json, Query, RealIp, Redirect},
-    IntoResponse,
+    Error, IntoResponse,
 };
+use reqwest::StatusCode;
 use serde::Deserialize;
 use std::{collections::HashSet, sync::Arc};
 use url::Url;
@@ -109,27 +110,44 @@ pub async fn callback(
 
 #[handler]
 pub async fn me(state: Data<&Arc<AppState>>, token: AuthToken) -> impl IntoResponse {
-    let user = UserData::get_by_id(token.session.user_id, &state.database)
-        .await
-        .unwrap();
-
-    Json(user)
+    match token {
+        AuthToken::Active(active_user) => {
+            let user = UserData::get_by_id(active_user.session.user_id, &state.database)
+                .await
+                .unwrap();
+
+            Json(user).into_response()
+        }
+        _ => Error::from_string("Not Authenticated", StatusCode::UNAUTHORIZED).into_response(),
+    }
 }
 
 #[handler]
 pub async fn get_sessions(state: Data<&Arc<AppState>>, token: AuthToken) -> impl IntoResponse {
-    let sessions = SessionState::get_by_user_id(token.session.user_id, &state.database)
-        .await
-        .unwrap();
-
-    Json(sessions)
+    match token {
+        AuthToken::Active(active_user) => {
+            let sessions =
+                SessionState::get_by_user_id(active_user.session.user_id, &state.database)
+                    .await
+                    .unwrap();
+
+            Json(sessions).into_response()
+        }
+        _ => Error::from_string("Not Authenticated", StatusCode::UNAUTHORIZED).into_response(),
+    }
 }
 
 #[handler]
 pub async fn delete_sessions(state: Data<&Arc<AppState>>, token: AuthToken) -> impl IntoResponse {
-    let sessions = SessionState::invalidate_by_user_id(token.session.user_id, &state.database)
-        .await
-        .unwrap();
-
-    Json(sessions)
+    match token {
+        AuthToken::Active(active_user) => {
+            let sessions =
+                SessionState::invalidate_by_user_id(active_user.session.user_id, &state.database)
+                    .await
+                    .unwrap();
+
+            Json(sessions).into_response()
+        }
+        _ => Error::from_string("Not Authenticated", StatusCode::UNAUTHORIZED).into_response(),
+    }
 }