1 files changed, 30 insertions, 2 deletions

diff --git a/src/types/extract.rs b/src/types/extract.rs
index 65d9f1a..f05215f 100644
--- a/src/types/extract.rs
+++ b/src/types/extract.rs
@@ -14,9 +14,11 @@ use axum::{
 use bytes::Bytes;
 use image::{io::Reader, DynamicImage, ImageFormat};
 use serde::de::DeserializeOwned;
+use tokio::sync::Mutex;
 use tower_cookies::Cookies;
 
 use crate::{
+    database,
     public::admin,
     public::console,
     types::{
@@ -97,11 +99,17 @@ where
             return Err(ResponseCode::Forbidden.text("No auth token provided"))
         };
 
-        let Ok(session) = Session::from_token(token.value()) else {
+        let Some(db) = parts.extensions.get::<DatabaseExtention>() else {
+            return Err(ResponseCode::Forbidden.text("Could not connect to database"))
+        };
+
+        let db = db.0.lock().await;
+
+        let Ok(session) = Session::from_token(&db, token.value()) else {
             return Err(ResponseCode::Unauthorized.text("Auth token invalid"))
         };
 
-        let Ok(user) = User::from_user_id(session.user_id, true)  else {
+        let Ok(user) = User::from_user_id(&db, session.user_id, true)  else {
             tracing::error!("Valid token but no valid user");
             return Err(ResponseCode::InternalServerError.text("Valid token but no valid user"))
         };
@@ -260,6 +268,26 @@ where
     }
 }
 
+pub struct DatabaseExtention(pub Mutex<database::Database>);
+pub struct Database(pub database::Database);
+
+#[async_trait]
+impl<S> FromRequestParts<S> for Database
+where
+    S: Send + Sync,
+{
+    type Rejection = Response;
+
+    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self> {
+        let db = parts.extensions.remove::<DatabaseExtention>();
+        let Some(db) = db else {
+            return Err(ResponseCode::InternalServerError.text("Database is not loaded"))
+        };
+
+        Ok(Self(db.0.into_inner()))
+    }
+}
+
 async fn read_body<S, B>(mut req: Request<B>, state: &S) -> Result<Vec<u8>>
 where
     B: HttpBody + Sync + Send + 'static,