persistence

.forgejo/workflows/pipeline.yaml

@@ -11,7 +11,7 @@ jobs:
         with:
           node-version: 24
       - uses: actions/checkout@v6
-      - run: cargo build -r
+      - run: ~/.cargo/bin/cargo build -r
       - run: |
           curl -X POST \
             -H "Authorization: token ${{ github.token }}" \
@@ -20,6 +20,7 @@ jobs:
       - run: |
           export XDG_RUNTIME_DIR=/run/user/$(id -u)
           export DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/$(id -u)/bus
+          systemctl --user stop chat
           cp target/release/chat ~/chat/chat
           chmod +x ~/chat/chat
-          systemctl --user restart chat
+          systemctl --user start chat

Cargo.toml

@@ -1,3 +1,6 @@
+[workspace]
+members = [".", "migration"]
+
 [package]
 name = "chat"
 version = "0.1.0"
@@ -8,16 +11,15 @@ axum = { version = "0.8.8", features = ["ws"] }
 dotenvy = "0.15.7"
 futures = "0.3.31"
 oauth2 = "4.4.2"
+jsonwebtoken = "9.3.0"
 reqwest = { version = "0.13.2", features = ["json"] }
 serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.149"
 time = "0.3.47"
 tokio = { version = "1.49.0", features = ["full"] }
-tower-sessions = "0.15.0"
-
-[profile.release]
-strip = true  # Automatically strip symbols from the binary.
-opt-level = "z"  # Optimize for size.
-lto = true # Enable Link Time Optimization
-codegen-units = 1 #  Maximize LTO
-panic = "abort" # Abort on panic
+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

@@ -11,6 +11,11 @@
             --parchment: #f5e6d3;
             --steam: #e0e0e0;
             --gear-color: rgba(0, 0, 0, 0.2);
+            --online-green: #4caf50;
+        }
+
+        * {
+            box-sizing: border-box;
         }
 
         body {
@@ -34,11 +39,102 @@
             padding-bottom: 10px;
             letter-spacing: 2px;
             text-transform: uppercase;
+            margin-bottom: 20px;
         }
 
-        #chat-container {
+        #main-container {
             width: 100%;
-            max-width: 800px;
+            max-width: 1000px;
+            display: flex;
+            gap: 15px;
+            height: 600px;
+        }
+
+        /* Contacts Sidebar */
+        #contacts-panel {
+            width: 250px;
+            background: rgba(43, 29, 14, 0.9);
+            border: 4px solid var(--copper);
+            border-radius: 10px;
+            box-shadow:
+                0 0 15px var(--brass),
+                inset 0 0 20px #000;
+            display: flex;
+            flex-direction: column;
+            position: relative;
+        }
+
+        #contacts-header {
+            padding: 15px;
+            border-bottom: 2px solid var(--copper);
+            font-weight: bold;
+            color: var(--brass);
+            text-transform: uppercase;
+            letter-spacing: 1px;
+        }
+
+        #contacts-list {
+            flex: 1;
+            overflow-y: auto;
+            padding: 10px;
+            scrollbar-width: thin;
+            scrollbar-color: var(--brass) var(--leather);
+        }
+
+        .contact {
+            padding: 12px 15px;
+            margin-bottom: 8px;
+            background: rgba(0, 0, 0, 0.3);
+            border: 1px solid var(--copper);
+            border-radius: 5px;
+            cursor: pointer;
+            display: flex;
+            align-items: center;
+            gap: 10px;
+            transition: all 0.2s;
+        }
+
+        .contact:hover {
+            background: rgba(212, 175, 55, 0.2);
+            border-color: var(--brass);
+        }
+
+        .contact.selected {
+            background: rgba(212, 175, 55, 0.3);
+            border-color: var(--brass);
+            box-shadow: 0 0 10px rgba(212, 175, 55, 0.3);
+        }
+
+        .contact .status-dot {
+            width: 10px;
+            height: 10px;
+            background: var(--online-green);
+            border-radius: 50%;
+            box-shadow: 0 0 5px var(--online-green);
+        }
+
+        .contact.offline .status-dot {
+            background: #666;
+            box-shadow: none;
+        }
+
+        .contact .username {
+            flex: 1;
+            font-weight: bold;
+        }
+
+        .contact .unread {
+            background: var(--copper);
+            color: var(--parchment);
+            padding: 2px 8px;
+            border-radius: 10px;
+            font-size: 0.8em;
+            font-weight: bold;
+        }
+
+        /* Chat Panel */
+        #chat-container {
+            flex: 1;
             background: rgba(43, 29, 14, 0.9);
             border: 4px solid var(--copper);
             border-radius: 10px;
@@ -74,14 +170,37 @@
             right: 10px;
         }
 
+        #chat-header {
+            padding: 10px;
+            border-bottom: 2px solid var(--copper);
+            font-weight: bold;
+            color: var(--brass);
+            text-transform: uppercase;
+            letter-spacing: 1px;
+        }
+
+        #no-chat-selected {
+            flex: 1;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            color: var(--copper);
+            font-style: italic;
+        }
+
         #chat {
-            height: 500px;
-            overflow-y: scroll;
+            flex: 1;
+            overflow-y: auto;
             border: 2px inset var(--copper);
             background: rgba(0, 0, 0, 0.3);
             padding: 15px;
             scrollbar-width: thin;
             scrollbar-color: var(--brass) var(--leather);
+            display: none;
+        }
+
+        #chat.active {
+            display: block;
         }
 
         #chat::-webkit-scrollbar {
@@ -107,16 +226,28 @@
             box-shadow: 2px 2px 2px rgba(0, 0, 0, 0.3);
             font-weight: bold;
             position: relative;
+            max-width: 80%;
         }
 
-        .message::after {
-            content: '';
-            position: absolute;
-            bottom: -5px;
-            left: 10px;
-            border-width: 5px 5px 0;
-            border-style: solid;
-            border-color: var(--copper) transparent;
+        .message.sent {
+            margin-left: auto;
+            background: linear-gradient(to bottom, var(--brass), var(--copper));
+            color: #2b1d0e;
+        }
+
+        .message.received {
+            margin-right: auto;
+        }
+
+        .message .sender {
+            font-size: 0.8em;
+            color: var(--copper);
+            margin-bottom: 4px;
+        }
+
+        .message.sent .sender {
+            color: #2b1d0e;
+            opacity: 0.7;
         }
 
         .system-msg {
@@ -128,14 +259,11 @@
             background: none;
             border: none;
             box-shadow: none;
-        }
-
-        .system-msg::after {
-            display: none;
+            max-width: 100%;
         }
 
         #controls {
-            display: flex;
+            display: none;
             gap: 10px;
             padding: 10px;
             background: #1a1a1a;
@@ -143,6 +271,10 @@
             border-radius: 5px;
         }
 
+        #controls.active {
+            display: flex;
+        }
+
         input[type="text"] {
             flex-grow: 1;
             padding: 10px;
@@ -179,33 +311,155 @@
 
 <body>
     <h1>&#9881; AetherChat &#9881;</h1>
+    <div id="main-container">
+        <!-- Contacts Sidebar -->
+        <div id="contacts-panel">
+            <div id="contacts-header">&#9881; Contacts</div>
+            <div id="contacts-list"></div>
+        </div>
+
+        <!-- Chat Area -->
         <div id="chat-container">
+            <div id="chat-header">Select a contact to start chatting</div>
+            <div id="no-chat-selected">Click on a contact to begin transmission...</div>
             <div id="chat"></div>
             <div id="controls">
                 <input type="text" id="message" placeholder="Transmit payload..." autofocus>
                 <button onclick="sendMessage()">Engage</button>
             </div>
         </div>
+    </div>
 
     <script>
-        const chatCheck = document.getElementById("chat");
+        const CLIENT_ID = "xLL8a23JBnBvEarjQUY6Jgd4zLnnqIi2y3oB6laM";
+        const AUTH_URL = "https://idm.flegr.me/application/o/authorize/";
+        const REDIRECT_URI = window.location.origin + window.location.pathname;
+
+        const contactsList = document.getElementById("contacts-list");
+        const chatDiv = document.getElementById("chat");
+        const chatHeader = document.getElementById("chat-header");
+        const noChatSelected = document.getElementById("no-chat-selected");
+        const controls = document.getElementById("controls");
         const messageInput = document.getElementById("message");
 
-        // Connect to WebSocket
-        const proto = window.location.protocol === 'https:' ? 'wss' : 'ws';
-        const ws = new WebSocket(`${proto}://${window.location.host}/ws`);
+        // State
+        let allUsers = [];
+        let onlineUsers = new Set();
+        let selectedUser = null;
+        let messageHistory = {}; // { username: [messages...] }
+        let unreadCounts = {}; // { username: count }
+        let currentUsername = null; // Will be set from welcome message
 
-        ws.onmessage = (event) => {
-            const div = document.createElement("div");
-            div.textContent = event.data;
-            div.className = "message";
+        // Handle OAuth Callback
+        async function handleAuth() {
+            const params = new URLSearchParams(window.location.search);
+            const code = params.get("code");
 
-            if (event.data.startsWith("System:")) {
-                div.className += " system-msg";
+            if (code) {
+                // Exchange code for token
+                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: "authorization_code",
+                            client_id: CLIENT_ID,
+                            code: code,
+                            redirect_uri: REDIRECT_URI,
+                        })
+                    });
+                    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);
+                    }
+                } catch (e) {
+                    console.error("Error exchanging code:", e);
+                }
             }
 
-            chatCheck.appendChild(div);
-            chatCheck.scrollTop = chatCheck.scrollHeight;
+            let accessToken = localStorage.getItem("access_token");
+            const refreshToken = localStorage.getItem("refresh_token");
+
+            // 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 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 () => {
+            const token = await handleAuth();
+            if (!token) return;
+
+            // Connect to WebSocket with token
+            const proto = window.location.protocol === 'https:' ? 'wss' : 'ws';
+            const ws = new WebSocket(`${proto}://${window.location.host}/ws?token=${token}`);
+
+            ws.onmessage = (event) => {
+                try {
+                    const msg = JSON.parse(event.data);
+                    handleMessage(msg);
+                } catch (e) {
+                    console.error("Failed to parse message:", e);
+                    // For debugging: show what we received
+                    addSystemMessage("Error parsing: " + event.data);
+                }
             };
 
             ws.onopen = () => {
@@ -216,18 +470,177 @@
                 addSystemMessage("Connection severed. Check steam pressure.");
             };
 
+            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":
+                        allUsers = msg.users;
+                        onlineUsers = new Set(msg.online);
+                        renderContacts();
+                        break;
+
+                    case "user_joined":
+                        if (msg.username !== currentUsername) {
+                            if (!allUsers.includes(msg.username)) {
+                                allUsers.push(msg.username);
+                            }
+                            onlineUsers.add(msg.username);
+                            renderContacts();
+                        }
+                        break;
+
+                    case "user_left":
+                        onlineUsers.delete(msg.username);
+                        renderContacts();
+                        if (selectedUser === msg.username) {
+                            addSystemMessage(`${msg.username} has disconnected.`);
+                        }
+                        break;
+
+                    case "private_message":
+                        handlePrivateMessage(msg);
+                        break;
+
+                    case "system":
+                        addSystemMessage(msg.content);
+                        break;
+                }
+            }
+
+            function handlePrivateMessage(msg) {
+                const otherUser = msg.from === currentUsername ? msg.to : msg.from;
+                const isSent = msg.from === currentUsername;
+
+                // Initialize history if needed
+                if (!messageHistory[otherUser]) {
+                    messageHistory[otherUser] = [];
+                }
+
+                // Store message
+                messageHistory[otherUser].push({
+                    from: msg.from,
+                    content: msg.content,
+                    isSent: isSent
+                });
+
+                // If this chat is open, render it
+                if (selectedUser === otherUser) {
+                    renderChat();
+                } else if (!isSent) {
+                    // Increment unread count for received messages
+                    unreadCounts[otherUser] = (unreadCounts[otherUser] || 0) + 1;
+                    renderContacts();
+                }
+            }
+
+            function renderContacts() {
+                contactsList.innerHTML = "";
+                allUsers.forEach(username => {
+                    const div = document.createElement("div");
+                    const isOnline = onlineUsers.has(username);
+                    div.className = "contact" +
+                        (selectedUser === username ? " selected" : "") +
+                        (isOnline ? "" : " offline");
+                    div.onclick = () => selectUser(username);
+
+                    const statusDot = document.createElement("span");
+                    statusDot.className = "status-dot";
+
+                    const nameSpan = document.createElement("span");
+                    nameSpan.className = "username";
+                    nameSpan.textContent = username;
+
+                    div.appendChild(statusDot);
+                    div.appendChild(nameSpan);
+
+                    // Show unread badge
+                    if (unreadCounts[username] && unreadCounts[username] > 0) {
+                        const badge = document.createElement("span");
+                        badge.className = "unread";
+                        badge.textContent = unreadCounts[username];
+                        div.appendChild(badge);
+                    }
+
+                    contactsList.appendChild(div);
+                });
+            }
+
+            async function selectUser(username) {
+                selectedUser = username;
+                unreadCounts[username] = 0; // Clear unread
+                chatHeader.textContent = `Chat with ${username}`;
+                noChatSelected.style.display = "none";
+                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();
+            }
+
+            function renderChat() {
+                chatDiv.innerHTML = "";
+                const messages = messageHistory[selectedUser] || [];
+                messages.forEach(msg => {
+                    const div = document.createElement("div");
+                    div.className = "message " + (msg.isSent ? "sent" : "received");
+
+                    const sender = document.createElement("div");
+                    sender.className = "sender";
+                    sender.textContent = msg.from;
+
+                    const content = document.createElement("div");
+                    content.textContent = msg.content;
+
+                    div.appendChild(sender);
+                    div.appendChild(content);
+                    chatDiv.appendChild(div);
+                });
+                chatDiv.scrollTop = chatDiv.scrollHeight;
+            }
+
             function addSystemMessage(text) {
                 const div = document.createElement("div");
-            div.textContent = `System: ${text}`;
+                div.textContent = text;
                 div.className = "message system-msg";
-            chatCheck.appendChild(div);
-            chatCheck.scrollTop = chatCheck.scrollHeight;
+                chatDiv.appendChild(div);
+                chatDiv.scrollTop = chatDiv.scrollHeight;
             }
 
             function sendMessage() {
-            const msg = messageInput.value;
-            if (msg) {
-                ws.send(msg);
+                const content = messageInput.value.trim();
+                if (content && selectedUser) {
+                    ws.send(JSON.stringify({
+                        to: selectedUser,
+                        content: content
+                    }));
                     messageInput.value = "";
                 }
             }
@@ -235,6 +648,7 @@
             messageInput.addEventListener("keypress", (e) => {
                 if (e.key === "Enter") sendMessage();
             });
+        })();
     </script>
 </body>
 

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

@@ -1,211 +1,445 @@
 use axum::{
-    extract::{
-        ws::{Message, WebSocket, WebSocketUpgrade},
-        Query, State,
-    },
-    http::StatusCode,
-    response::{Html, IntoResponse, Redirect},
-    routing::get,
     Router,
+    extract::{
+        Query, State,
+        ws::{Message, WebSocket, WebSocketUpgrade},
+    },
+    http::{HeaderMap, StatusCode},
+    response::{Html, IntoResponse},
+    routing::get,
 };
 use dotenvy::dotenv;
 use futures::{sink::SinkExt, stream::StreamExt};
-use oauth2::{
-    basic::BasicClient,AuthUrl, ClientId, ClientSecret, RedirectUrl, TokenResponse,
-    TokenUrl,
-};
-use std::sync::Arc;
-use tokio::sync::broadcast;
-use tower_sessions::{cookie::SameSite, Expiry, MemoryStore, Session, SessionManagerLayer};
+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;
 
-// Define application state
-struct AppState {
-    tx: broadcast::Sender<String>,
-    oauth_client: BasicClient,
+// WebSocket message types for client-server communication
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type")]
+enum WsMessage {
+    #[serde(rename = "user_list")]
+    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")]
+    System { content: String },
 }
 
-const AFTER_LOGIN_URL: &str = "/";
+// Per-user connection channel
+type UserTx = mpsc::UnboundedSender<String>;
+
+// Shared state for tracking online users and their connections
+struct ChatState {
+    // Map of username -> sender channel
+    connections: RwLock<HashMap<String, UserTx>>,
+}
+
+impl ChatState {
+    fn new() -> Self {
+        Self {
+            connections: RwLock::new(HashMap::new()),
+        }
+    }
+
+    async fn add_user(&self, username: String, tx: UserTx) {
+        let mut conns = self.connections.write().await;
+        conns.insert(username, tx);
+    }
+
+    async fn remove_user(&self, username: &str) {
+        let mut conns = self.connections.write().await;
+        conns.remove(username);
+    }
+
+    async fn get_online_users(&self) -> Vec<String> {
+        let conns = self.connections.read().await;
+        conns.keys().cloned().collect()
+    }
+
+    async fn send_to_user(&self, username: &str, message: &str) -> bool {
+        let conns = self.connections.read().await;
+        if let Some(tx) = conns.get(username) {
+            tx.send(message.to_string()).is_ok()
+        } else {
+            false
+        }
+    }
+
+    async fn broadcast_except(&self, message: &str, exclude: &str) {
+        let conns = self.connections.read().await;
+        for (username, tx) in conns.iter() {
+            if username != exclude {
+                let _ = tx.send(message.to_string());
+            }
+        }
+    }
+}
+
+// Application state
+struct AppState {
+    chat: ChatState,
+    jwks: RwLock<HashMap<String, (String, String)>>,
+    db: DatabaseConnection,
+}
 
 #[tokio::main]
 async fn main() {
     dotenv().ok();
 
-    // Create a broadcast channel
-    let (tx, _rx) = broadcast::channel(100);
+    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");
 
-    // OAuth configuration
-    let client_id = dotenvy::var("CLIENT_ID").unwrap();
-    let client_secret = dotenvy::var("CLIENT_SECRET").unwrap();
-    let port = dotenvy::var("PORT").unwrap_or_else(|_| "3001".to_string());
-    let host = dotenvy::var("HOST").unwrap_or_else(|_| "http://localhost:".to_string() + &port);
-    // NOTE: In production, do not hardcode localhost
-    let auth_url = AuthUrl::new("https://idm.flegr.me/application/o/authorize/".to_string())
-        .expect("Invalid authorization endpoint URL");
-    let token_url = TokenUrl::new("https://idm.flegr.me/application/o/token/".to_string())
-        .expect("Invalid token endpoint URL");
+    // Run migrations
+    Migrator::up(&db, None)
+        .await
+        .expect("Failed to run migrations");
 
-    let client = BasicClient::new(
-        ClientId::new(client_id),
-        Some(ClientSecret::new(client_secret)),
-        auth_url,
-        Some(token_url),
-    )
-    .set_redirect_uri(
-        RedirectUrl::new(host + "/auth/callback")
-            .expect("Invalid redirect URL"),
-    );
+    let chat = ChatState::new();
+    let jwks = RwLock::new(HashMap::new());
 
-    let app_state = Arc::new(AppState {
-        tx,
-        oauth_client: client,
-    });
+    let app_state = Arc::new(AppState { chat, jwks, db });
 
-    // Session configuration
-    let session_store = MemoryStore::default();
-    let session_layer = SessionManagerLayer::new(session_store)
-        .with_secure(false) // For localhost; set to true in production with HTTPS
-        .with_same_site(SameSite::Lax) // Ensure we can receive cookies from OAuth redirect
-        .with_expiry(Expiry::OnInactivity(time::Duration::minutes(30)));       
+    // 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("/login", get(login))
-        .route("/logout", get(logout))
-        .route("/auth/callback", get(auth_callback))
+        .route("/api/history", get(get_history))
         .route("/ws", get(websocket_handler))
-        .layer(session_layer)
         .with_state(app_state);
 
     // Run the app
     let port = dotenvy::var("PORT").unwrap_or_else(|_| "3001".to_string());
-    let listener = tokio::net::TcpListener::bind("0.0.0.0:".to_string() + &port).await.unwrap();
+    let listener = tokio::net::TcpListener::bind("0.0.0.0:".to_string() + &port)
+        .await
+        .unwrap();
     println!("Chat server listening on http://0.0.0.0:{}", port);
     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(session: Session) -> impl IntoResponse {
-    if let Some(_user) = session.get::<User>("user").await.unwrap() {
-        return Html(include_str!("../index.html")).into_response();
+async fn index() -> impl IntoResponse {
+    Html(include_str!("../index.html"))
+}
+
+#[derive(Deserialize)]
+struct WsParams {
+    token: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct Claims {
+    preferred_username: String,
+    // 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);
             }
-    Redirect::to("/login").into_response()
-}
+            let jwks = state.jwks.read().await;
+            match jwks.get(&kid).cloned() {
+                Some(c) => c,
+                None => return Err(StatusCode::UNAUTHORIZED),
+            }
+        }
+    };
 
-async fn login(State(state): State<Arc<AppState>>) -> impl IntoResponse {
-    let (auth_url, _csrf_token) = state
-        .oauth_client
-        .authorize_url(oauth2::CsrfToken::new_random)
-        .url();
+    let key = match DecodingKey::from_rsa_components(&components.0, &components.1) {
+        Ok(k) => k,
+        Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
+    };
 
-    Redirect::to(auth_url.as_str())
-}
+    let mut validation = Validation::new(Algorithm::RS256);
+    validation.set_audience(&["xLL8a23JBnBvEarjQUY6Jgd4zLnnqIi2y3oB6laM"]);
+    validation.set_issuer(&["https://idm.flegr.me/application/o/chat/"]);
 
-async fn logout(session: Session) -> impl IntoResponse {
-    session.delete().await.unwrap();
-    Redirect::to("/login")
-}
-
-#[derive(serde::Deserialize)]
-struct AuthRequest {
-    code: String,
-    state: String,
-}
-
-async fn auth_callback(
-    Query(query): Query<AuthRequest>,
-    State(state): State<Arc<AppState>>,
-    session: Session,
-) -> impl IntoResponse {
-    let token = match state
-        .oauth_client
-        .exchange_code(oauth2::AuthorizationCode::new(query.code))
-        .request_async(oauth2::reqwest::async_http_client)
-        .await
-    {
-        Ok(t) => t,
+    match decode::<Claims>(token, &key, &validation) {
+        Ok(c) => Ok(c.claims),
         Err(e) => {
-            return (
-                StatusCode::BAD_REQUEST,
-                format!("Failed to exchange authorization code: {}", e),
-            )
-            .into_response();
+            eprintln!("Token validation failed: {}", e);
+            Err(StatusCode::UNAUTHORIZED)
+        }
     }
-    };
-
-    let client = reqwest::Client::new();
-    let user_info_resp = client
-        .get("https://idm.flegr.me/application/o/userinfo/")
-        .header("User-Agent", "axum-chat-app")
-        .bearer_auth(token.access_token().secret())
-        .send()
-        .await;
-    
-    let user_data: User = match user_info_resp {
-        Ok(resp) => match resp.json().await {
-            Ok(u) => u,
-            Err(e) => return (
-                StatusCode::INTERNAL_SERVER_ERROR, 
-                format!("Failed to parse user info: {}", e)
-            ).into_response()
-        },
-        Err(e) => return (
-            StatusCode::INTERNAL_SERVER_ERROR, 
-            format!("Failed to fetch user info: {}", e)
-        ).into_response()
-    };
-
-    session.insert("user", user_data).await.unwrap();
-
-    Redirect::to(AFTER_LOGIN_URL).into_response()
 }
 
 async fn websocket_handler(
     ws: WebSocketUpgrade,
+    Query(params): Query<WsParams>,
     State(state): State<Arc<AppState>>,
-    session: Session,
 ) -> impl IntoResponse {
-    if let Some(user) = session.get::<User>("user").await.unwrap() {
-        return ws.on_upgrade(move |socket| websocket(socket, state, user));
-    }
-    StatusCode::UNAUTHORIZED.into_response()
+    let claims = match verify_token(&params.token, &state).await {
+        Ok(c) => c,
+        Err(status) => return status.into_response(),
+    };
+
+    let user = User {
+        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))
 }
 
-// WebSocket handler
-async fn websocket(stream: WebSocket, state: Arc<AppState>, user: User) {
-    let (mut sender, mut receiver) = stream.split();
-    let mut rx = state.tx.subscribe();
+#[derive(Deserialize)]
+struct HistoryParams {
+    to: String,
+}
 
-    // Send welcome message
-    let msg = format!("System: Welcome, {}!", user.login);
-    let _ = sender.send(Message::Text(msg.into())).await;
+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(),
+    };
 
-    // Send task
-    let mut send_task = tokio::spawn(async move {
-        while let Ok(msg) = rx.recv().await {
-            if sender.send(Message::Text(msg.into())).await.is_err() {
+    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();
+    let username = user.login.clone();
+
+    // Create channel for this user
+    let (tx, mut rx) = mpsc::unbounded_channel::<String>();
+
+    // Add user to connections
+    state.chat.add_user(username.clone(), tx).await;
+    println!("User connected: {}", username);
+
+    // Send identity
+    let identity_msg = serde_json::to_string(&WsMessage::Identity {
+        username: username.clone(),
+    })
+    .unwrap();
+    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 {
+        while let Some(msg) = rx.recv().await {
+            if ws_sender.send(Message::Text(msg.into())).await.is_err() {
                 break;
             }
         }
     });
 
-    // Receive task
-    let tx = state.tx.clone();
-    let username = user.login.clone();
-    let mut recv_task = tokio::spawn(async move {
-        while let Some(Ok(Message::Text(text))) = receiver.next().await {
-            // Prepend username
-            let msg = format!("{}: {}", username, text);
-            let _ = tx.send(msg);
+    // Broadcast user joined to everyone else (after send_task is spawned so others receive it)
+    let join_msg = serde_json::to_string(&WsMessage::UserJoined {
+        username: username.clone(),
+    })
+    .unwrap();
+    state.chat.broadcast_except(&join_msg, &username).await;
+
+    // Task to receive messages from WebSocket and route them
+    let state_clone = state.clone();
+    let username_clone = username.clone();
+    let recv_task = tokio::spawn(async move {
+        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(),
+                    to: msg.to.clone(),
+                    content: msg.content.clone(),
+                })
+                .unwrap();
+
+                // Send to recipient
+                state_clone.chat.send_to_user(&msg.to, &private_msg).await;
+
+                // Also send back to sender (for their own chat view)
+                state_clone
+                    .chat
+                    .send_to_user(&username_clone, &private_msg)
+                    .await;
+            }
         }
     });
 
+    // Wait for either task to complete
     tokio::select! {
-        _ = (&mut send_task) => recv_task.abort(),
-        _ = (&mut recv_task) => send_task.abort(),
-    };
+        _ = send_task => {},
+        _ = recv_task => {},
+    }
+
+    // User disconnected - clean up
+    state.chat.remove_user(&username).await;
+}
+
+// Message structure from client
+#[derive(Debug, Deserialize)]
+struct ClientMessage {
+    to: String,
+    content: String,
 }