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 }
@@ -466,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)) {
-                            onlineUsers.push(msg.username);
+                        if (msg.username !== currentUsername) {
+                            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.`);
@@ -491,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;
                 }
@@ -531,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");
@@ -558,7 +572,7 @@
                 });
             }
 
-            function selectUser(username) {
+            async function selectUser(username) {
                 selectedUser = username;
                 unreadCounts[username] = 0; // Clear unread
                 chatHeader.textContent = `Chat with ${username}`;
@@ -566,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,7 +4,7 @@ use axum::{
         Query, State,
         ws::{Message, WebSocket, WebSocketUpgrade},
     },
-    http::StatusCode,
+    http::{HeaderMap, StatusCode},
     response::{Html, IntoResponse},
     routing::get,
 };
@@ -15,7 +15,15 @@ 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
@@ -23,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")]
@@ -90,16 +103,27 @@ impl ChatState {
 struct AppState {
     chat: ChatState,
     jwks: RwLock<HashMap<String, (String, String)>>,
+    db: DatabaseConnection,
 }
 
 #[tokio::main]
 async fn main() {
     dotenv().ok();
 
+    let database_url = dotenvy::var("DATABASE_URL").expect("DATABASE_URL must be set");
+    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 });
+    let app_state = Arc::new(AppState { chat, jwks, db });
 
     // Initial JWKS fetch
     if let Err(e) = fetch_jwks(&app_state).await {
@@ -109,6 +133,7 @@ async fn main() {
     // 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);
 
@@ -166,19 +191,15 @@ struct Claims {
     // Add other fields as needed
 }
 
-async fn websocket_handler(
-    ws: WebSocketUpgrade,
-    Query(params): Query<WsParams>,
-    State(state): State<Arc<AppState>>,
-) -> impl IntoResponse {
-    let header = match decode_header(&params.token) {
+async fn verify_token(token: &str, state: &AppState) -> Result<Claims, StatusCode> {
+    let header = match decode_header(token) {
         Ok(h) => h,
-        Err(_) => return StatusCode::BAD_REQUEST.into_response(),
+        Err(_) => return Err(StatusCode::BAD_REQUEST),
     };
 
     let kid = match header.kid {
         Some(k) => k,
-        None => return StatusCode::UNAUTHORIZED.into_response(),
+        None => return Err(StatusCode::UNAUTHORIZED),
     };
 
     let components = {
@@ -188,44 +209,123 @@ async fn websocket_handler(
         } else {
             // Key not found, might need to refresh JWKS
             drop(jwks);
-            if let Err(e) = fetch_jwks(&state).await {
+            if let Err(e) = fetch_jwks(state).await {
                 eprintln!("Failed to refresh JWKS: {}", e);
-                return StatusCode::UNAUTHORIZED.into_response();
+                return Err(StatusCode::UNAUTHORIZED);
             }
             let jwks = state.jwks.read().await;
             match jwks.get(&kid).cloned() {
                 Some(c) => c,
-                None => return StatusCode::UNAUTHORIZED.into_response(),
+                None => return Err(StatusCode::UNAUTHORIZED),
             }
         }
     };
 
     let key = match DecodingKey::from_rsa_components(&components.0, &components.1) {
         Ok(k) => k,
-        Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
+        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/"]);
 
-    let token_data = match decode::<Claims>(&params.token, &key, &validation) {
-        Ok(c) => c,
+    match decode::<Claims>(token, &key, &validation) {
+        Ok(c) => Ok(c.claims),
         Err(e) => {
             eprintln!("Token validation failed: {}", e);
-            return StatusCode::UNAUTHORIZED.into_response();
+            Err(StatusCode::UNAUTHORIZED)
         }
+    }
+}
+
+async fn websocket_handler(
+    ws: WebSocketUpgrade,
+    Query(params): Query<WsParams>,
+    State(state): State<Arc<AppState>>,
+) -> impl IntoResponse {
+    let claims = match verify_token(&params.token, &state).await {
+        Ok(c) => c,
+        Err(status) => return status.into_response(),
     };
 
     let user = User {
-        login: token_data.claims.preferred_username,
+        login: claims.preferred_username,
         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;
+
     let state_clone = state.clone();
     ws.on_upgrade(move |socket| handle_websocket(socket, state_clone, user))
 }
 
+#[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(),
+    };
+
+    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()
+        }
+    }
+}
+
 // WebSocket connection handler
 async fn handle_websocket(stream: WebSocket, state: Arc<AppState>, user: User) {
     let (mut ws_sender, mut ws_receiver) = stream.split();
@@ -238,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
-    let users = state.chat.get_online_users().await;
-    println!("Sending user list to {}: {:?}", username, users);
-    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),
+    // Send identity
+    let identity_msg = serde_json::to_string(&WsMessage::Identity {
+        username: username.clone(),
     })
     .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 {
@@ -274,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(),