Format, Status, User Stati & Endpoints

src/sql/connection_status.rs

@@ -0,0 +1,38 @@
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ConnectionType {
+    Online,
+    UserOffline,
+    IotaOffline,
+    Away,
+    DoNotDisturb,
+}
+impl ConnectionType {
+    pub fn to_str(&self) -> &str {
+        match self {
+            ConnectionType::Online => "online",
+            ConnectionType::UserOffline => "user_offline",
+            ConnectionType::IotaOffline => "iota_offline",
+            ConnectionType::Away => "away",
+            ConnectionType::DoNotDisturb => "do_not_disturb",
+        }
+    }
+    pub fn to_string(&self) -> String {
+        match self {
+            ConnectionType::Online => "online".to_string(),
+            ConnectionType::UserOffline => "user_offline".to_string(),
+            ConnectionType::IotaOffline => "iota_offline".to_string(),
+            ConnectionType::Away => "away".to_string(),
+            ConnectionType::DoNotDisturb => "do_not_disturb".to_string(),
+        }
+    }
+    pub fn from_str(s: &str) -> Option<ConnectionType> {
+        match s.to_lowercase().as_str() {
+            "online" => Some(ConnectionType::Online),
+            "user_offline" => Some(ConnectionType::UserOffline),
+            "iota_offline" => Some(ConnectionType::IotaOffline),
+            "away" => Some(ConnectionType::Away),
+            "do_not_disturb" => Some(ConnectionType::DoNotDisturb),
+            _ => None,
+        }
+    }
+}

src/sql/mod.rs

@@ -1,3 +1,4 @@
+pub mod connection_status;
 pub mod iota_omikron_tracker;
 pub mod sql;
 pub mod user_online_tracker;

src/sql/sql.rs

@@ -202,11 +202,11 @@ pub async fn get_by_username(
                 String::from_utf8_lossy(&token).to_string(),
             ))
         }
-        None => Err(sqlx::Error::RowNotFound),
+        _ => Err(sqlx::Error::RowNotFound),
     }
 }
 
-pub async fn get_by_id(
+pub async fn get_by_user_id(
     id: i64,
 ) -> Result<
     (
@@ -265,10 +265,73 @@ pub async fn get_by_id(
                 String::from_utf8_lossy(&token).to_string(),
             ))
         }
-        None => Err(sqlx::Error::RowNotFound),
+        _ => Err(sqlx::Error::RowNotFound),
     }
 }
 
+pub async fn get_users_by_iota_id(
+    iota_id_param: i64,
+) -> Result<
+    Vec<(
+        i64,
+        i64,
+        String,
+        String,
+        String,
+        String,
+        String,
+        i32,
+        i64,
+        String,
+        String,
+        String,
+    )>,
+    sqlx::Error,
+> {
+    let db_lock = SQL_DB.read().await;
+    let pool = db_lock.as_ref().expect("Database pool is not initialized");
+
+    let rows = sqlx::query(
+        "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE iota_id = ?",
+    )
+    .bind(iota_id_param)
+    .fetch_all(pool)
+    .await?;
+
+    let mut users = Vec::new();
+    for row in rows {
+        let id: i64 = row.get("id");
+        let iota_id: i64 = row.get("iota_id");
+        let username: String = row.get("username");
+        let display: Vec<u8> = row.get("display");
+        let status: Vec<u8> = row.get("status");
+        let about: Vec<u8> = row.get("about");
+        let avatar: Vec<u8> = row.get("avatar");
+        let sub_level: i32 = row.get("sub_level");
+        let sub_end: i64 = row.get("sub_end");
+        let public_key: String = row.get("public_key");
+        let private_key_hash: String = row.get("private_key_hash");
+        let token: Vec<u8> = row.get("token");
+
+        users.push((
+            id,
+            iota_id,
+            username,
+            String::from_utf8_lossy(&display).to_string(),
+            String::from_utf8_lossy(&status).to_string(),
+            String::from_utf8_lossy(&about).to_string(),
+            String::from_utf8_lossy(&avatar).to_string(),
+            sub_level,
+            sub_end,
+            public_key,
+            private_key_hash,
+            String::from_utf8_lossy(&token).to_string(),
+        ));
+    }
+
+    Ok(users)
+}
+
 pub async fn change_username(id: i64, new_username: String) -> Result<(), sqlx::Error> {
     let db_lock = SQL_DB.read().await;
     let pool = db_lock.as_ref().expect("Database pool is not initialized");
@@ -514,7 +577,7 @@ pub async fn get_random_omikron() -> Result<(i64, String, String), sqlx::Error>
             String::from_utf8_lossy(&public_key).to_string(),
             String::from_utf8_lossy(&ip_address).to_string(),
         )),
-        None => Err(sqlx::Error::RowNotFound),
+        _ => Err(sqlx::Error::RowNotFound),
     }
 }
 
@@ -534,7 +597,7 @@ pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error>
             String::from_utf8_lossy(&public_key).to_string(),
             String::from_utf8_lossy(&ip_address).to_string(),
         )),
-        None => Err(sqlx::Error::RowNotFound),
+        _ => Err(sqlx::Error::RowNotFound),
     }
 }
 

src/sql/user_online_tracker.rs

@@ -1,27 +1,100 @@
+use crate::sql;
+use crate::sql::connection_status::ConnectionType;
 use once_cell::sync::Lazy;
 use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::sync::RwLock;
 
-static USER_OMIKRON_MAP: Lazy<Arc<RwLock<HashMap<i64, i64>>>> =
+#[derive(Debug, Clone)]
+pub struct UserStatus {
+    pub connection_type: ConnectionType,
+    pub omikron_id: i64,
+}
+
+// IotaID -> Vec<OmikronID>
+static IOTA_OMIKRON_CONNECTIONS: Lazy<Arc<RwLock<HashMap<i64, Vec<i64>>>>> =
     Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
 
-pub async fn track_user_omikron(user: i64, omikron: i64) {
-    let mut c = USER_OMIKRON_MAP.write().await;
-    c.insert(user, omikron);
+// UserID -> UserStatus
+static USER_STATUS_MAP: Lazy<Arc<RwLock<HashMap<i64, UserStatus>>>> =
+    Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
+
+pub async fn track_iota_connection(iota_id: i64, omikron_id: i64) {
+    let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await;
+    let connections = iota_map.entry(iota_id).or_default();
+    if !connections.contains(&omikron_id) {
+        connections.push(omikron_id);
+    }
 }
 
-pub async fn get_omikron_for_user(user: i64) -> Option<i64> {
-    let c = USER_OMIKRON_MAP.read().await;
-    c.get(&user).cloned()
+pub async fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
+    let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await;
+    if let Some(connections) = iota_map.get_mut(&iota_id) {
+        connections.retain(|&id| id != omikron_id);
+        if connections.is_empty() {
+            iota_map.remove(&iota_id);
+            return true; // Iota is now offline
+        }
+    }
+    false
 }
 
-pub async fn untrack_user(user: i64) {
-    let mut c = USER_OMIKRON_MAP.write().await;
-    c.remove(&user);
+pub async fn get_iota_omikron_connections(iota_id: i64) -> Option<Vec<i64>> {
+    let iota_map = IOTA_OMIKRON_CONNECTIONS.read().await;
+    iota_map.get(&iota_id).cloned()
 }
 
-pub async fn untrack_by_omikron(omikron: i64) {
-    let mut c = USER_OMIKRON_MAP.write().await;
-    c.retain(|_, v| *v != omikron);
+pub async fn track_user_status(user_id: i64, status: ConnectionType, omikron_id: i64) {
+    let mut user_map = USER_STATUS_MAP.write().await;
+    user_map.insert(
+        user_id,
+        UserStatus {
+            connection_type: status,
+            omikron_id,
+        },
+    );
+}
+
+pub async fn get_user_status(user_id: i64) -> Option<UserStatus> {
+    let user_map = USER_STATUS_MAP.read().await;
+    user_map.get(&user_id).cloned()
+}
+
+pub async fn untrack_user(user_id: i64) {
+    let mut user_map = USER_STATUS_MAP.write().await;
+    user_map.remove(&user_id);
+}
+
+pub async fn untrack_many_users(user_ids: &[i64]) {
+    let mut user_map = USER_STATUS_MAP.write().await;
+    for user_id in user_ids {
+        user_map.remove(user_id);
+    }
+}
+
+pub async fn untrack_omikron(omikron_id: i64) {
+    let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await;
+    let mut user_map = USER_STATUS_MAP.write().await;
+
+    let mut offline_iotas = Vec::new();
+
+    iota_map.retain(|iota_id, connections| {
+        connections.retain(|id| *id != omikron_id);
+        if connections.is_empty() {
+            offline_iotas.push(*iota_id);
+            false
+        } else {
+            true
+        }
+    });
+
+    user_map.retain(|_, status| status.omikron_id != omikron_id);
+
+    for iota_id in offline_iotas {
+        if let Ok(users) = sql::sql::get_users_by_iota_id(iota_id).await {
+            for user in users {
+                user_map.remove(&user.0);
+            }
+        }
+    }
 }