persistence

Cargo.toml

@@ -1,3 +1,6 @@
+[workspace]
+members = [".", "migration"]
+
 [package]
 name = "chat"
 version = "0.1.0"
@@ -14,3 +17,9 @@ serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.149"
 time = "0.3.47"
 tokio = { version = "1.49.0", features = ["full"] }
+sea-orm = { version = "1.1.0", features = ["runtime-tokio-rustls", "sqlx-postgres", "macros", "with-chrono", "with-uuid"] }
+sqlx = { version = "0.8.2", features = ["runtime-tokio-rustls", "postgres", "chrono", "uuid", "macros"] }
+chrono = { version = "0.4", features = ["serde"] }
+uuid = { version = "1.10", features = ["v4", "serde"] }
+migration = { path = "migration" }
+# Removed tower-sessions

index.html

@@ -113,6 +113,11 @@
             box-shadow: 0 0 5px var(--online-green);
         }
 
+        .contact.offline .status-dot {
+            background: #666;
+            box-shadow: none;
+        }
+
         .contact .username {
             flex: 1;
             font-weight: bold;
@@ -338,7 +343,8 @@
         const messageInput = document.getElementById("message");
 
         // State
-        let onlineUsers = [];
+        let allUsers = [];
+        let onlineUsers = new Set();
         let selectedUser = null;
         let messageHistory = {}; // { username: [messages...] }
         let unreadCounts = {}; // { username: count }
@@ -351,7 +357,6 @@
 
             if (code) {
                 // Exchange code for token
-                // NOTE: This assumes the Authentik client is "Public" and doesn't require a secret
                 try {
                     const response = await fetch("https://idm.flegr.me/application/o/token/", {
                         method: "POST",
@@ -366,6 +371,9 @@
                     const data = await response.json();
                     if (data.access_token) {
                         localStorage.setItem("access_token", data.access_token);
+                        if (data.refresh_token) {
+                            localStorage.setItem("refresh_token", data.refresh_token);
+                        }
                         window.history.replaceState({}, document.title, window.location.pathname);
                     } else {
                         console.error("Token exchange failed:", data);
@@ -375,13 +383,64 @@
                 }
             }
 
-            const storedToken = localStorage.getItem("access_token");
+            let accessToken = localStorage.getItem("access_token");
-            if (!storedToken) {
+            const refreshToken = localStorage.getItem("refresh_token");
-                const url = `${AUTH_URL}?client_id=${CLIENT_ID}&response_type=code&redirect_uri=${encodeURIComponent(REDIRECT_URI)}&scope=openid+profile+email`;
+
+            // Check if token is expired or about to expire (within 60 seconds)
+            if (accessToken && isTokenExpired(accessToken)) {
+                console.log("Access token expired, attempting refresh...");
+                if (refreshToken) {
+                    accessToken = await performTokenRefresh(refreshToken);
+                } else {
+                    accessToken = null;
+                }
+            }
+
+            if (!accessToken) {
+                // Redirect to Authentik with offline_access scope to get a refresh_token
+                const url = `${AUTH_URL}?client_id=${CLIENT_ID}&response_type=code&redirect_uri=${encodeURIComponent(REDIRECT_URI)}&scope=openid+profile+email+offline_access`;
                 window.location.href = url;
                 return null;
             }
-            return storedToken;
+            return accessToken;
+        }
+
+        function isTokenExpired(token) {
+            try {
+                const payload = JSON.parse(atob(token.split('.')[1]));
+                const now = Math.floor(Date.now() / 1000);
+                return payload.exp < (now + 60);
+            } catch (e) {
+                return true;
+            }
+        }
+
+        async function performTokenRefresh(refreshToken) {
+            try {
+                const response = await fetch("https://idm.flegr.me/application/o/token/", {
+                    method: "POST",
+                    headers: { "Content-Type": "application/x-www-form-urlencoded" },
+                    body: new URLSearchParams({
+                        grant_type: "refresh_token",
+                        client_id: CLIENT_ID,
+                        refresh_token: refreshToken,
+                    })
+                });
+                const data = await response.json();
+                if (data.access_token) {
+                    localStorage.setItem("access_token", data.access_token);
+                    if (data.refresh_token) {
+                        localStorage.setItem("refresh_token", data.refresh_token);
+                    }
+                    console.log("Token refreshed successfully");
+                    return data.access_token;
+                }
+                console.error("Token refresh failed:", data);
+                return null;
+            } catch (e) {
+                console.error("Error during token refresh:", e);
+                return null;
+            }
         }
 
         (async () => {
@@ -413,20 +472,32 @@
 
             function handleMessage(msg) {
                 switch (msg.type) {
+                    case "identity":
+                        currentUsername = msg.username;
+                        console.log("Identity confirmed:", currentUsername);
+                        // Re-filter contacts now that we know our username
+                        allUsers = allUsers.filter(u => u !== currentUsername);
+                        renderContacts();
+                        break;
+
                     case "user_list":
-                        onlineUsers = msg.users.filter(u => u !== currentUsername);
+                        allUsers = msg.users;
+                        onlineUsers = new Set(msg.online);
                         renderContacts();
                         break;
 
                     case "user_joined":
-                        if (msg.username !== currentUsername && !onlineUsers.includes(msg.username)) {
+                        if (msg.username !== currentUsername) {
-                            onlineUsers.push(msg.username);
+                            if (!allUsers.includes(msg.username)) {
+                                allUsers.push(msg.username);
+                            }
+                            onlineUsers.add(msg.username);
                             renderContacts();
                         }
                         break;
 
                     case "user_left":
-                        onlineUsers = onlineUsers.filter(u => u !== msg.username);
+                        onlineUsers.delete(msg.username);
                         renderContacts();
                         if (selectedUser === msg.username) {
                             addSystemMessage(`${msg.username} has disconnected.`);
@@ -438,13 +509,6 @@
                         break;
 
                     case "system":
-                        // Extract username from welcome message
-                        if (msg.content.startsWith("Welcome, ")) {
-                            currentUsername = msg.content.replace("Welcome, ", "").replace("!", "");
-                            // Re-filter contacts now that we know our username
-                            onlineUsers = onlineUsers.filter(u => u !== currentUsername);
-                            renderContacts();
-                        }
                         addSystemMessage(msg.content);
                         break;
                 }
@@ -478,9 +542,12 @@
 
             function renderContacts() {
                 contactsList.innerHTML = "";
-                onlineUsers.forEach(username => {
+                allUsers.forEach(username => {
                     const div = document.createElement("div");
-                    div.className = "contact" + (selectedUser === username ? " selected" : "");
+                    const isOnline = onlineUsers.has(username);
+                    div.className = "contact" +
+                        (selectedUser === username ? " selected" : "") +
+                        (isOnline ? "" : " offline");
                     div.onclick = () => selectUser(username);
 
                     const statusDot = document.createElement("span");
@@ -505,7 +572,7 @@
                 });
             }
 
-            function selectUser(username) {
+            async function selectUser(username) {
                 selectedUser = username;
                 unreadCounts[username] = 0; // Clear unread
                 chatHeader.textContent = `Chat with ${username}`;
@@ -513,6 +580,27 @@
                 chatDiv.classList.add("active");
                 controls.classList.add("active");
                 renderContacts();
+
+                // Fetch history from DB
+                try {
+                    const token = localStorage.getItem("access_token");
+                    const response = await fetch(`/api/history?to=${encodeURIComponent(username)}`, {
+                        headers: {
+                            "Authorization": `Bearer ${token}`
+                        }
+                    });
+                    if (response.ok) {
+                        const history = await response.json();
+                        messageHistory[username] = history.map(msg => ({
+                            from: msg.from_user,
+                            content: msg.content,
+                            isSent: msg.from_user === currentUsername
+                        }));
+                    }
+                } catch (e) {
+                    console.error("Failed to fetch history:", e);
+                }
+
                 renderChat();
                 messageInput.focus();
             }

migration/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "migration"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+name = "migration"
+path = "src/lib.rs"
+
+[dependencies]
+async-trait = "0.1"
+sea-orm-migration = { version = "1.1.0", features = ["runtime-tokio-rustls", "sqlx-postgres"] }
+tokio = { version = "1", features = ["full"] }

migration/src/lib.rs

@@ -0,0 +1,12 @@
+pub use sea_orm_migration::prelude::*;
+
+mod m20220101_000001_create_table;
+
+pub struct Migrator;
+
+#[async_trait::async_trait]
+impl MigratorTrait for Migrator {
+    fn migrations() -> Vec<Box<dyn MigrationTrait>> {
+        vec![Box::new(m20220101_000001_create_table::Migration)]
+    }
+}

migration/src/m20220101_000001_create_table.rs

@@ -0,0 +1,98 @@
+use sea_orm_migration::prelude::*;
+
+#[derive(Iden)]
+enum Users {
+    Table,
+    Id,
+    Username,
+    CreatedAt,
+}
+
+#[derive(Iden)]
+enum Messages {
+    Table,
+    Id,
+    FromUser,
+    ToUser,
+    Content,
+    CreatedAt,
+}
+
+#[derive(DeriveMigrationName)]
+pub struct Migration;
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        manager
+            .create_table(
+                Table::create()
+                    .table(Users::Table)
+                    .if_not_exists()
+                    .col(
+                        ColumnDef::new(Users::Id)
+                            .integer()
+                            .not_null()
+                            .auto_increment()
+                            .primary_key(),
+                    )
+                    .col(
+                        ColumnDef::new(Users::Username)
+                            .string()
+                            .not_null()
+                            .unique_key(),
+                    )
+                    .col(
+                        ColumnDef::new(Users::CreatedAt)
+                            .date_time()
+                            .not_null()
+                            .extra("DEFAULT CURRENT_TIMESTAMP".to_string()),
+                    )
+                    .to_owned(),
+            )
+            .await?;
+
+        manager
+            .create_table(
+                Table::create()
+                    .table(Messages::Table)
+                    .if_not_exists()
+                    .col(ColumnDef::new(Messages::Id).uuid().not_null().primary_key())
+                    .col(ColumnDef::new(Messages::FromUser).string().not_null())
+                    .col(ColumnDef::new(Messages::ToUser).string().not_null())
+                    .col(ColumnDef::new(Messages::Content).text().not_null())
+                    .col(
+                        ColumnDef::new(Messages::CreatedAt)
+                            .date_time()
+                            .not_null()
+                            .extra("DEFAULT CURRENT_TIMESTAMP".to_string()),
+                    )
+                    .foreign_key(
+                        ForeignKey::create()
+                            .name("fk-messages-from_user")
+                            .from(Messages::Table, Messages::FromUser)
+                            .to(Users::Table, Users::Username),
+                    )
+                    .foreign_key(
+                        ForeignKey::create()
+                            .name("fk-messages-to_user")
+                            .from(Messages::Table, Messages::ToUser)
+                            .to(Users::Table, Users::Username),
+                    )
+                    .to_owned(),
+            )
+            .await?;
+
+        Ok(())
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        manager
+            .drop_table(Table::drop().table(Messages::Table).to_owned())
+            .await?;
+        manager
+            .drop_table(Table::drop().table(Users::Table).to_owned())
+            .await?;
+        Ok(())
+    }
+}

migration/src/main.rs

@@ -0,0 +1,6 @@
+use sea_orm_migration::prelude::*;
+
+#[tokio::main]
+async fn main() {
+    cli::run_cli(migration::Migrator).await;
+}

src/entities/messages.rs

@@ -0,0 +1,37 @@
+use sea_orm::entity::prelude::*;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
+#[sea_orm(table_name = "messages")]
+pub struct Model {
+    #[sea_orm(primary_key, auto_increment = false)]
+    pub id: Uuid,
+    pub from_user: String,
+    pub to_user: String,
+    pub content: String,
+    pub created_at: DateTime,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::users::Entity",
+        from = "Column::FromUser",
+        to = "super::users::Column::Username"
+    )]
+    FromUser,
+    #[sea_orm(
+        belongs_to = "super::users::Entity",
+        from = "Column::ToUser",
+        to = "super::users::Column::Username"
+    )]
+    ToUser,
+}
+
+impl Related<super::users::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::FromUser.def() // Or ToUser, depending on what you want as the 'default' relationship
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}

src/entities/mod.rs

@@ -0,0 +1,7 @@
+pub mod messages;
+pub mod users;
+
+pub mod prelude {
+    pub use super::messages::Entity as Messages;
+    pub use super::users::Entity as Users;
+}

src/entities/users.rs

@@ -0,0 +1,26 @@
+use sea_orm::entity::prelude::*;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
+#[sea_orm(table_name = "users")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: i32,
+    #[sea_orm(unique)]
+    pub username: String,
+    pub created_at: DateTime,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(has_many = "super::messages::Entity")]
+    Messages,
+}
+
+impl Related<super::messages::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::Messages.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}

src/main.rs

@@ -4,17 +4,26 @@ use axum::{
         Query, State,
         ws::{Message, WebSocket, WebSocketUpgrade},
     },
-    http::StatusCode,
+    http::{HeaderMap, StatusCode},
     response::{Html, IntoResponse},
     routing::get,
 };
 use dotenvy::dotenv;
 use futures::{sink::SinkExt, stream::StreamExt};
+use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
 use serde::{Deserialize, Serialize};
 use std::{collections::HashMap, sync::Arc};
 use tokio::sync::{RwLock, mpsc};
 
+use migration::{Migrator, MigratorTrait};
+use sea_orm::{
+    ActiveModelTrait, ActiveValue, ColumnTrait, Condition, Database, DatabaseConnection,
+    EntityTrait, QueryFilter, QueryOrder,
+};
+
+mod entities;
 mod user;
+use entities::{messages, prelude::*, users};
 use user::User;
 
 // WebSocket message types for client-server communication
@@ -22,13 +31,18 @@ use user::User;
 #[serde(tag = "type")]
 enum WsMessage {
     #[serde(rename = "user_list")]
-    UserList { users: Vec<String> },
+    UserList {
+        users: Vec<String>,
+        online: Vec<String>,
+    },
     #[serde(rename = "private_message")]
     PrivateMessage {
         from: String,
         to: String,
         content: String,
     },
+    #[serde(rename = "identity")]
+    Identity { username: String },
     #[serde(rename = "user_joined")]
     UserJoined { username: String },
     #[serde(rename = "system")]
@@ -88,19 +102,38 @@ impl ChatState {
 // Application state
 struct AppState {
     chat: ChatState,
+    jwks: RwLock<HashMap<String, (String, String)>>,
+    db: DatabaseConnection,
 }
 
 #[tokio::main]
 async fn main() {
     dotenv().ok();
 
-    let app_state = Arc::new(AppState {
+    let database_url = dotenvy::var("DATABASE_URL").expect("DATABASE_URL must be set");
-        chat: ChatState::new(),
+    let db = Database::connect(database_url)
-    });
+        .await
+        .expect("Failed to connect to database");
+
+    // Run migrations
+    Migrator::up(&db, None)
+        .await
+        .expect("Failed to run migrations");
+
+    let chat = ChatState::new();
+    let jwks = RwLock::new(HashMap::new());
+
+    let app_state = Arc::new(AppState { chat, jwks, db });
+
+    // Initial JWKS fetch
+    if let Err(e) = fetch_jwks(&app_state).await {
+        eprintln!("Warning: Failed to fetch initial JWKS: {}", e);
+    }
 
     // Build application with routes
     let app = Router::new()
         .route("/", get(index))
+        .route("/api/history", get(get_history))
         .route("/ws", get(websocket_handler))
         .with_state(app_state);
 
@@ -113,6 +146,34 @@ async fn main() {
     axum::serve(listener, app).await.unwrap();
 }
 
+async fn fetch_jwks(state: &AppState) -> Result<(), Box<dyn std::error::Error>> {
+    let resp: JwkSet = reqwest::get("https://idm.flegr.me/application/o/chat/jwks/")
+        .await?
+        .json()
+        .await?;
+
+    let mut keys = state.jwks.write().await;
+    for jwk in resp.keys {
+        if let (Some(kid), Some(n), Some(e)) = (jwk.kid, jwk.n, jwk.e) {
+            keys.insert(kid, (n, e));
+        }
+    }
+    println!("Fetched {} public keys from Authentik", keys.len());
+    Ok(())
+}
+
+#[derive(Deserialize)]
+struct JwkSet {
+    keys: Vec<Jwk>,
+}
+
+#[derive(Deserialize)]
+struct Jwk {
+    kid: Option<String>,
+    n: Option<String>,
+    e: Option<String>,
+}
+
 // Handlers
 
 async fn index() -> impl IntoResponse {
@@ -130,30 +191,138 @@ struct Claims {
     // Add other fields as needed
 }
 
+async fn verify_token(token: &str, state: &AppState) -> Result<Claims, StatusCode> {
+    let header = match decode_header(token) {
+        Ok(h) => h,
+        Err(_) => return Err(StatusCode::BAD_REQUEST),
+    };
+
+    let kid = match header.kid {
+        Some(k) => k,
+        None => return Err(StatusCode::UNAUTHORIZED),
+    };
+
+    let components = {
+        let jwks = state.jwks.read().await;
+        if let Some(c) = jwks.get(&kid).cloned() {
+            c
+        } else {
+            // Key not found, might need to refresh JWKS
+            drop(jwks);
+            if let Err(e) = fetch_jwks(state).await {
+                eprintln!("Failed to refresh JWKS: {}", e);
+                return Err(StatusCode::UNAUTHORIZED);
+            }
+            let jwks = state.jwks.read().await;
+            match jwks.get(&kid).cloned() {
+                Some(c) => c,
+                None => return Err(StatusCode::UNAUTHORIZED),
+            }
+        }
+    };
+
+    let key = match DecodingKey::from_rsa_components(&components.0, &components.1) {
+        Ok(k) => k,
+        Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
+    };
+
+    let mut validation = Validation::new(Algorithm::RS256);
+    validation.set_audience(&["xLL8a23JBnBvEarjQUY6Jgd4zLnnqIi2y3oB6laM"]);
+    validation.set_issuer(&["https://idm.flegr.me/application/o/chat/"]);
+
+    match decode::<Claims>(token, &key, &validation) {
+        Ok(c) => Ok(c.claims),
+        Err(e) => {
+            eprintln!("Token validation failed: {}", e);
+            Err(StatusCode::UNAUTHORIZED)
+        }
+    }
+}
+
 async fn websocket_handler(
     ws: WebSocketUpgrade,
     Query(params): Query<WsParams>,
     State(state): State<Arc<AppState>>,
 ) -> impl IntoResponse {
-    // For now, we'll verify the token by calling the userinfo endpoint.
+    let claims = match verify_token(&params.token, &state).await {
-    // In a production app, you should verify the JWT signature locally using JWKS.
+        Ok(c) => c,
-    let client = reqwest::Client::new();
+        Err(status) => return status.into_response(),
-    let user_info_resp = client
+    };
-        .get("https://idm.flegr.me/application/o/userinfo/")
+
-        .header("User-Agent", "axum-chat-app")
+    let user = User {
-        .bearer_auth(&params.token)
+        login: claims.preferred_username,
-        .send()
+        avatar_url: String::new(),
+    };
+
+    // Ensure user exists in DB
+    let user_active_model = users::ActiveModel {
+        username: ActiveValue::Set(user.login.clone()),
+        ..Default::default()
+    };
+    let _ = Users::insert(user_active_model)
+        .on_conflict(
+            sea_orm::sea_query::OnConflict::column(users::Column::Username)
+                .do_nothing()
+                .to_owned(),
+        )
+        .exec(&state.db)
         .await;
 
-    match user_info_resp {
+    let state_clone = state.clone();
-        Ok(resp) if resp.status().is_success() => {
+    ws.on_upgrade(move |socket| handle_websocket(socket, state_clone, user))
-            let user: User = match resp.json().await {
+}
-                Ok(u) => u,
+
-                Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
+#[derive(Deserialize)]
+struct HistoryParams {
+    to: String,
+}
+
+async fn get_history(
+    headers: HeaderMap,
+    Query(params): Query<HistoryParams>,
+    State(state): State<Arc<AppState>>,
+) -> impl IntoResponse {
+    let auth_header = match headers.get(axum::http::header::AUTHORIZATION) {
+        Some(h) => h.to_str().unwrap_or(""),
+        None => return StatusCode::UNAUTHORIZED.into_response(),
     };
-            ws.on_upgrade(move |socket| handle_websocket(socket, state, user))
+
+    if !auth_header.starts_with("Bearer ") {
+        return StatusCode::UNAUTHORIZED.into_response();
+    }
+
+    let token = &auth_header[7..];
+    let claims = match verify_token(token, &state).await {
+        Ok(c) => c,
+        Err(status) => return status.into_response(),
+    };
+
+    let current_user = claims.preferred_username;
+
+    let messages = Messages::find()
+        .filter(
+            Condition::any()
+                .add(
+                    Condition::all()
+                        .add(messages::Column::FromUser.eq(current_user.clone()))
+                        .add(messages::Column::ToUser.eq(params.to.clone())),
+                )
+                .add(
+                    Condition::all()
+                        .add(messages::Column::FromUser.eq(params.to))
+                        .add(messages::Column::ToUser.eq(current_user)),
+                ),
+        )
+        .order_by_asc(messages::Column::CreatedAt)
+        .all(&state.db)
+        .await;
+
+    match messages {
+        Ok(msgs) => axum::Json(msgs).into_response(),
+        Err(e) => {
+            eprintln!("Failed to fetch history: {}", e);
+            StatusCode::INTERNAL_SERVER_ERROR.into_response()
         }
-        _ => StatusCode::UNAUTHORIZED.into_response(),
     }
 }
 
@@ -169,18 +338,41 @@ async fn handle_websocket(stream: WebSocket, state: Arc<AppState>, user: User) {
     state.chat.add_user(username.clone(), tx).await;
     println!("User connected: {}", username);
 
-    // Send current user list to the newly connected user
+    // Send identity
-    let users = state.chat.get_online_users().await;
+    let identity_msg = serde_json::to_string(&WsMessage::Identity {
-    println!("Sending user list to {}: {:?}", username, users);
+        username: username.clone(),
-    let user_list_msg = serde_json::to_string(&WsMessage::UserList { users }).unwrap();
-    let _ = ws_sender.send(Message::Text(user_list_msg.into())).await;
-
-    // Send welcome message
-    let welcome = serde_json::to_string(&WsMessage::System {
-        content: format!("Welcome, {}!", username),
     })
     .unwrap();
-    let _ = ws_sender.send(Message::Text(welcome.into())).await;
+    let _ = ws_sender.send(Message::Text(identity_msg.into())).await;
+
+    // Send global user list and online status to the newly connected user
+    let (db_users, online_users) = match Users::find().all(&state.db).await {
+        Ok(all_users) => {
+            let db_list = all_users
+                .into_iter()
+                .map(|u| u.username)
+                .filter(|u| u != &username)
+                .collect::<Vec<String>>();
+            let online_list = state.chat.get_online_users().await;
+            (db_list, online_list)
+        }
+        Err(e) => {
+            eprintln!("Failed to fetch users from DB: {}", e);
+            (Vec::new(), Vec::new())
+        }
+    };
+    println!(
+        "Sending global user list to {}: DB Count={}, Online Count={}",
+        username,
+        db_users.len(),
+        online_users.len()
+    );
+    let user_list_msg = serde_json::to_string(&WsMessage::UserList {
+        users: db_users,
+        online: online_users,
+    })
+    .unwrap();
+    let _ = ws_sender.send(Message::Text(user_list_msg.into())).await;
 
     // Task to forward messages from channel to WebSocket
     let send_task = tokio::spawn(async move {
@@ -205,6 +397,16 @@ async fn handle_websocket(stream: WebSocket, state: Arc<AppState>, user: User) {
         while let Some(Ok(Message::Text(text))) = ws_receiver.next().await {
             // Parse incoming message
             if let Ok(msg) = serde_json::from_str::<ClientMessage>(&text) {
+                // Persist to DB
+                let new_message = messages::ActiveModel {
+                    id: ActiveValue::Set(uuid::Uuid::new_v4()),
+                    from_user: ActiveValue::Set(username_clone.clone()),
+                    to_user: ActiveValue::Set(msg.to.clone()),
+                    content: ActiveValue::Set(msg.content.clone()),
+                    ..Default::default()
+                };
+                let _ = new_message.insert(&state_clone.db).await;
+
                 // Create private message
                 let private_msg = serde_json::to_string(&WsMessage::PrivateMessage {
                     from: username_clone.clone(),