From 5cc46855889ba6f3b7fe0a627bc19e2bd85ce4ea Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Mon, 23 Feb 2026 07:57:10 +0100
Subject: [PATCH] Session management

---
 Cargo.lock     |  4 +--
 Cargo.toml     |  2 +-
 src/main.rs    | 79 +++++++++++++++++++++++++++++++++++++++++++++++---
 src/methods.rs | 20 +++++++++++++
 4 files changed, 98 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fe36849..611ad9c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1757,9 +1757,9 @@ dependencies = [
 
 [[package]]
 name = "session-rs"
-version = "0.1.1"
+version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "51e9fdc5adcdd0ddc8c5077df84a2158d2446780efdbee8c379573e9eed17094"
+checksum = "7166810f8fb16dd8df2796a217296c9c8a4a7d59c50df415cf758ebd2932003f"
 dependencies = [
  "base64",
  "rand 0.10.0",
diff --git a/Cargo.toml b/Cargo.toml
index f4947d9..3f789d2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,7 +7,7 @@ edition = "2024"
 reqwest = { version = "0.13.2", features = ["json"] }
 serde = { version = "1.0.228", features = ["serde_derive"] }
 serde_json = "1.0.149"
-session-rs = "0.1.1"
+session-rs = "0.1.2"
 sqlx = { version = "0.8.6", features = [
     "sqlite",
     "runtime-tokio-native-tls",
diff --git a/src/main.rs b/src/main.rs
index e41f364..fcbb233 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,31 +2,64 @@ mod cosmetics;
 mod methods;
 mod user;
 
-use std::{pin::Pin, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    pin::Pin,
+    sync::Arc,
+};
 
 use serde::{Deserialize, Serialize};
-use session_rs::server::SessionServer;
+use session_rs::{server::SessionServer, session::Session};
 use sqlx::SqlitePool;
 use tokio::sync::Mutex;
 
 use crate::user::User;
 
+type SessionMap = Arc<Mutex<HashMap<String, HashSet<Session>>>>;
+
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> session_rs::Result<()> {
     let pool = Arc::new(user::init_db().await);
     let server = SessionServer::bind("127.0.0.1:8080").await?;
 
+    let sessions: SessionMap = Arc::new(Mutex::new(HashMap::new()));
+
     server
         .session_loop({
             let pool = Arc::clone(&pool);
             move |session, _| {
                 let pool = Arc::clone(&pool);
+                let sessions = Arc::clone(&sessions);
 
                 Box::pin(async move {
                     println!("Connected");
                     let uuid: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
 
-                    session.on_close(async || Ok(println!("Closed"))).await;
+                    session
+                        .on_close({
+                            let session = session.clone();
+                            let uuid = uuid.clone();
+                            let sessions = Arc::clone(&sessions);
+
+                            move || {
+                                let uuid = uuid.clone();
+                                let sessions = Arc::clone(&sessions);
+                                let session = session.clone();
+
+                                Box::pin(async move {
+                                    let uuid_lock = uuid.lock().await;
+                                    if !uuid_lock.is_empty() {
+                                        if let Some(sessions) =
+                                            sessions.lock().await.get_mut(uuid_lock.as_str())
+                                        {
+                                            sessions.remove(&session);
+                                        }
+                                    }
+                                    Ok(())
+                                })
+                            }
+                        })
+                        .await;
 
                     session.start_ping(
                         tokio::time::Duration::from_secs(30),
@@ -37,8 +70,18 @@ async fn main() -> session_rs::Result<()> {
                         .on_request::<methods::Auth, _>({
                             let pool = Arc::clone(&pool);
                             let uuid = Arc::clone(&uuid);
+                            let session = session.clone();
+                            let sessions = sessions.clone();
 
-                            move |_, token| authenticate(uuid.clone(), token, pool.clone())
+                            move |_, token| {
+                                authenticate(
+                                    sessions.clone(),
+                                    session.clone(),
+                                    uuid.clone(),
+                                    token,
+                                    pool.clone(),
+                                )
+                            }
                         })
                         .await;
 
@@ -118,6 +161,25 @@ async fn main() -> session_rs::Result<()> {
                         })
                         .await;
 
+                    session
+                        .on_request::<methods::SetHat, _>({
+                            let uuid = Arc::clone(&uuid);
+                            let pool = Arc::clone(&pool);
+
+                            move |_, item_id| {
+                                let pool = Arc::clone(&pool);
+                                let item_id = item_id.clone();
+
+                                cosmetics::equip(
+                                    cosmetics::CosmeticKind::Hat,
+                                    uuid.clone(),
+                                    item_id,
+                                    pool,
+                                )
+                            }
+                        })
+                        .await;
+
                     Ok::<(), session_rs::Error>(())
                 }) as Pin<Box<dyn Future<Output = _> + Send>>
             }
@@ -132,6 +194,8 @@ struct MinecraftAuthResponse {
 }
 
 async fn authenticate(
+    sessions: SessionMap,
+    session: Session,
     uuid: Arc<Mutex<String>>,
     session_token: String,
     pool: Arc<SqlitePool>,
@@ -163,5 +227,12 @@ async fn authenticate(
 
     *uuid.lock().await = auth.id.clone();
 
+    sessions
+        .lock()
+        .await
+        .entry(auth.id.clone())
+        .or_default()
+        .insert(session);
+
     user::get_put(&auth.id, &pool).await
 }
diff --git a/src/methods.rs b/src/methods.rs
index dcafb2e..51708fc 100644
--- a/src/methods.rs
+++ b/src/methods.rs
@@ -52,3 +52,23 @@ impl Method for BuyHat {
     type Response = String;
     type Error = String;
 }
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Emote;
+
+impl Method for Emote {
+    const NAME: &'static str = "emote";
+    type Request = String;
+    type Response = String;
+    type Error = String;
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct EmoteEvent;
+
+impl Method for EmoteEvent {
+    const NAME: &'static str = "emote_event";
+    type Request = String;
+    type Response = String;
+    type Error = String;
+}