user update

README.md

@@ -9,24 +9,7 @@ Based is a micro framework providing web dev primitives.
 - Templates (Shell)
 
 ## User Auth
-To use the user auth feature, make sure a migration has added the following to your PostgresDB:
+To use the user auth feature, make sure a migration has added [these tables](src/auth/auth.sql) to your PostgresDB:
-
-```sql
-CREATE TYPE user_role AS ENUM ('regular', 'admin');
-
-CREATE TABLE IF NOT EXISTS users (
-    username VARCHAR(255) NOT NULL PRIMARY KEY,
-    "password" text NOT NULL,
-    user_role user_role NOT NULL DEFAULT 'regular'
-);
-
-CREATE TABLE IF NOT EXISTS user_session (
-    id UUID NOT NULL PRIMARY KEY DEFAULT gen_random_uuid(),
-    token text NOT NULL,
-    "user" varchar(255) NOT NULL,
-    FOREIGN KEY("user") REFERENCES users(username)
-);
-```
 
 ## HTMX
 Based has a route for serving HTMX at `based::htmx::htmx_script_route` which you can include in your rocket `routes!`. The HTMX script will be available at `/assets/htmx.min.js`.

src/auth/auth.sql

@@ -1,4 +1,5 @@
 CREATE TYPE user_role AS ENUM ('regular', 'admin');
+CREATE TYPE session_kind AS ENUM ('api', 'user');
 
 CREATE TABLE IF NOT EXISTS users (
     username VARCHAR(255) NOT NULL PRIMARY KEY,
@@ -11,5 +12,13 @@ CREATE TABLE IF NOT EXISTS user_session (
     token text NOT NULL,
     "user" varchar(255) NOT NULL,
     "created" timestamptz NOT NULL DEFAULT NOW(),
+    "csrf" UUID NOT NULL DEFAULT gen_random_uuid(),
+    "name" varchar(255)
+    kind session_kind NOT NULL DEFAULT 'user',
     FOREIGN KEY("user") REFERENCES users(username)
 );
+
+CREATE TABLE IF NOT EXISTS user_profile_pic (
+    username VARCHAR(255) NOT NULL PRIMARY KEY,
+    "image" bytea NOT NULL
+);

src/auth/csrf.rs

@@ -0,0 +1,36 @@
+use super::User;
+use crate::get_pg;
+use std::str::FromStr;
+
+pub trait CSRF {
+    fn get_csrf(&self) -> impl std::future::Future<Output = uuid::Uuid>;
+    fn verify_csrf(&self, csrf: &str) -> impl std::future::Future<Output = bool>;
+}
+
+impl CSRF for User {
+    /// Get CSRF Token for the current session
+    async fn get_csrf(&self) -> uuid::Uuid {
+        let res: (uuid::Uuid,) = sqlx::query_as("SELECT csrf FROM user_session WHERE token = $1")
+            .bind(&self.session)
+            .fetch_one(get_pg!())
+            .await
+            .unwrap();
+
+        res.0
+    }
+
+    /// Verify CSRF and generate a new one
+    async fn verify_csrf(&self, csrf: &str) -> bool {
+        if self.get_csrf().await == uuid::Uuid::from_str(csrf).unwrap() {
+            sqlx::query("UPDATE user_session SET csrf = gen_random_uuid() WHERE token = $1")
+                .bind(&self.session)
+                .execute(get_pg!())
+                .await
+                .unwrap();
+
+            return true;
+        }
+
+        false
+    }
+}

src/auth/mod.rs

@@ -1,10 +1,12 @@
-use chrono::Utc;
 use data_encoding::HEXUPPER;
 use rand::RngCore;
-use serde::{Deserialize, Serialize};
-use sqlx::FromRow;
 
+pub mod csrf;
+pub mod profile_pic;
+mod session;
 mod user;
+pub use session::Session;
+pub use session::Sessions;
 pub use user::AdminUser;
 pub use user::MaybeUser;
 pub use user::User;
@@ -18,18 +20,6 @@ fn gen_token(token_length: usize) -> String {
     HEXUPPER.encode(&token_bytes)
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
-pub struct Session {
-    /// The unique ID of the session token
-    pub id: uuid::Uuid,
-    /// The generated session token
-    pub token: String,
-    /// The username associated with the session token
-    pub user: String,
-    /// Session creation time
-    pub created: chrono::DateTime<Utc>
-}
-
 /// A macro to check if a user has admin privileges.
 ///
 /// This macro checks whether the provided user has admin privileges by calling the `is_admin` method on it.

src/auth/profile_pic.rs

@@ -0,0 +1,31 @@
+use crate::get_pg;
+
+use super::User;
+pub trait ProfilePic {
+    fn profile_pic(&self) -> impl std::future::Future<Output = Option<Vec<u8>>>;
+    fn set_profile_pic(&self, image: Vec<u8>) -> impl std::future::Future<Output = ()>;
+}
+
+impl ProfilePic for User {
+    /// Get a user's profile picture from the database
+    async fn profile_pic(&self) -> Option<Vec<u8>> {
+        let res: Option<(Vec<u8>,)> =
+            sqlx::query_as("SELECT image FROM user_profile_pic WHERE username = $1;")
+                .bind(&self.username)
+                .fetch_optional(get_pg!())
+                .await
+                .unwrap();
+
+        res.map(|x| x.0)
+    }
+
+    /// Set a user's profile picture in the database
+    async fn set_profile_pic(&self, image: Vec<u8>) {
+        sqlx::query("INSERT INTO user_profile_pic (username, image) VALUES ($1, $2) ON CONFLICT (id) DO UPDATE SET image = EXCLUDED.image",
+        )
+        .bind(&self.username)
+        .bind(&image)
+        .execute(get_pg!())
+        .await.unwrap();
+    }
+}

src/auth/session.rs

@@ -0,0 +1,116 @@
+use chrono::Utc;
+use serde::{Deserialize, Serialize};
+use sqlx::FromRow;
+
+use crate::get_pg;
+
+use super::{User, UserRole, gen_token};
+
+#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
+pub struct Session {
+    /// The unique ID of the session token
+    pub id: uuid::Uuid,
+    /// The generated session token
+    pub token: String,
+    /// The username associated with the session token
+    pub user: String,
+    /// Session creation time
+    pub created: chrono::DateTime<Utc>,
+    /// Internal CSRF value
+    csrf: uuid::Uuid,
+    /// Named session value
+    pub name: Option<String>,
+    /// Kind of session
+    pub kind: SessionKind,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, sqlx::Type)]
+#[sqlx(type_name = "session_kind", rename_all = "lowercase")]
+pub enum SessionKind {
+    API,
+    USER,
+}
+
+pub trait Sessions {
+    fn from_session(session: &str) -> impl std::future::Future<Output = Option<User>>;
+    fn login(
+        username: &str,
+        password: &str,
+    ) -> impl std::future::Future<Output = Option<(Session, UserRole)>>;
+    fn api_key(&self, name: &str) -> impl std::future::Future<Output = Session>;
+    fn session(&self) -> impl std::future::Future<Output = Session>;
+    fn list_sessions(&self) -> impl std::future::Future<Output = Vec<Session>>;
+    fn end_session(&self) -> impl std::future::Future<Output = ()>;
+}
+
+impl Sessions for User {
+    /// Generate a new API Key session
+    async fn api_key(&self, name: &str) -> Session {
+        sqlx::query_as(
+            "INSERT INTO user_session (token, \"user\", kind, name) VALUES ($1, $2, $3) RETURNING *",
+        )
+        .bind(gen_token(64))
+        .bind(&self.username)
+        .bind(SessionKind::API)
+        .bind(name)
+        .fetch_one(get_pg!())
+        .await
+        .unwrap()
+    }
+
+    /// End a user session
+    async fn end_session(&self) -> () {
+        sqlx::query("DELETE FROM user_session WHERE token = $1")
+            .bind(&self.session)
+            .execute(get_pg!())
+            .await
+            .unwrap();
+    }
+
+    /// Get all sessions for a user
+    async fn list_sessions(&self) -> Vec<Session> {
+        sqlx::query_as("SELECT * FROM user_session WHERE user = $1")
+            .bind(&self.username)
+            .fetch_all(get_pg!())
+            .await
+            .unwrap()
+    }
+
+    // Get a user from session ID
+    async fn from_session(session: &str) -> Option<User> {
+        let user: Option<Self> = sqlx::query_as("SELECT * FROM users WHERE username = (SELECT \"user\" FROM user_session WHERE token = $1)").bind(session).fetch_optional(get_pg!()).await.unwrap();
+
+        if let Some(mut user) = user {
+            user.session = session.to_string();
+            return Some(user);
+        }
+
+        None
+    }
+
+    /// Login a user with the given username and password
+    async fn login(username: &str, password: &str) -> Option<(Session, UserRole)> {
+        let u = Self::find(username).await?;
+
+        if !u.verify_pw(password) {
+            return None;
+        }
+
+        Some((u.session().await, u.user_role))
+    }
+
+    /// Generate a new session token for the user
+    ///
+    /// Returns a Session instance containing the generated token and associated user
+    async fn session(&self) -> Session {
+        sqlx::query_as(
+            "INSERT INTO user_session (token, \"user\", kind) VALUES ($1, $2, $3) RETURNING *",
+        )
+        .bind(gen_token(64))
+        .bind(&self.username)
+        .bind(SessionKind::USER)
+        .fetch_one(get_pg!())
+        .await
+        .unwrap()
+    }
+}

src/auth/user.rs

@@ -3,9 +3,11 @@ use serde::{Deserialize, Serialize};
 use serde_json::json;
 use sqlx::FromRow;
 
-use super::{Session, gen_token};
+use super::Sessions;
 use crate::{get_pg, request::api::ToAPI};
 
+// TODO : 2FA
+
 /// User
 ///
 /// # Example:
@@ -27,7 +29,7 @@ pub struct User {
     /// The role of the user
     pub user_role: UserRole,
     #[sqlx(default)]
-    pub session: String
+    pub session: String,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, sqlx::Type)]
@@ -40,18 +42,6 @@ pub enum UserRole {
 }
 
 impl User {
-    // Get a user from session ID
-    pub async fn from_session(session: &str) -> Option<Self> {
-        let user: Option<Self> = sqlx::query_as("SELECT * FROM users WHERE username = (SELECT \"user\" FROM user_session WHERE token = $1)").bind(session).fetch_optional(get_pg!()).await.unwrap();
-        
-        if let Some(mut user) = user {
-            user.session = session.to_string();
-            return Some(user);
-        }
-
-        None
-    }
-
     /// Find a user by their username
     pub async fn find(username: &str) -> Option<Self> {
         sqlx::query_as("SELECT * FROM users WHERE username = $1")
@@ -74,6 +64,7 @@ impl User {
             username: username.to_string(),
             password: bcrypt::hash(password, bcrypt::DEFAULT_COST).unwrap(),
             user_role: role,
+            session: String::new(),
         };
 
         sqlx::query("INSERT INTO users (username, \"password\", user_role) VALUES ($1, $2, $3)")
@@ -87,17 +78,6 @@ impl User {
         Some(u)
     }
 
-    /// Login a user with the given username and password
-    pub async fn login(username: &str, password: &str) -> Option<(Session, UserRole)> {
-        let u = Self::find(username).await?;
-
-        if !u.verify_pw(password) {
-            return None;
-        }
-
-        Some((u.session().await, u.user_role))
-    }
-
     /// Change the password of a User
     ///
     /// Returns a Result indicating whether the password change was successful or not
@@ -124,20 +104,6 @@ impl User {
             .unwrap()
     }
 
-    /// Generate a new session token for the user
-    ///
-    /// Returns a Session instance containing the generated token and associated user
-    pub async fn session(&self) -> Session {
-        sqlx::query_as(
-            "INSERT INTO user_session (token, \"user\") VALUES ($1, $2) RETURNING id, token, \"user\"",
-        )
-        .bind(gen_token(64))
-        .bind(&self.username)
-        .fetch_one(get_pg!())
-        .await
-        .unwrap()
-    }
-
     /// Check if the user is an admin
     pub const fn is_admin(&self) -> bool {
         matches!(self.user_role, UserRole::Admin)
@@ -160,6 +126,7 @@ impl ToAPI for User {
     }
 }
 
+/// extracts a user from a request with `session` cookie
 async fn extract_user<'r>(request: &'r Request<'_>) -> Option<User> {
     if let Some(session_id) = request.cookies().get("session") {
         if let Some(user) = User::from_session(session_id.value()).await {

src/lib.rs

@@ -2,17 +2,13 @@ use tokio::sync::OnceCell;
 
 pub mod auth;
 pub mod format;
+pub mod htmx;
 pub mod page;
 pub mod request;
 pub mod result;
-pub mod htmx;
 
-// TODO : Cache Headers
-// TODO : Refactor Responders
-// TODO : Streaming Responses + Ranges
 // TODO : API Pagination?
 // TODO : CORS?
-// TODO : CSRF?
 
 // Postgres